def test_default_from_params_params(self):
beam_search = BeamSearch.from_params(
Params({
"beam_size": 2,
"end_index": 7
}))
assert beam_search.beam_size == 2
assert beam_search._end_index == 7
def test_params_no_sampling(self):
beam_search = BeamSearch.from_params(
Params({
"beam_size": 2,
"end_index": 7
}))
assert beam_search.beam_size == 2
assert beam_search._end_index == 7
assert beam_search.sampler is None
def test_params_p_sampling(self):
beam_search = BeamSearch.from_params(
Params({
"type": "top_p_sampling",
"beam_size": 2,
"end_index": 7,
"p": 0.4,
}))
assert beam_search.beam_size == 2
assert beam_search._end_index == 7
assert beam_search.sampler is not None
def test_params_sampling(self):
beam_search = BeamSearch.from_params(
Params({
"sampler": {
"type": "top-k",
"k": 4,
},
"beam_size": 2,
"end_index": 7,
}))
assert beam_search.beam_size == 2
assert beam_search._end_index == 7
assert beam_search.sampler is not None
