File size: 2,079 Bytes
d758c99 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 |
import copy
import operator
import attr
@attr.s
class Hypothesis:
inference_state = attr.ib()
next_choices = attr.ib()
score = attr.ib(default=0)
choice_history = attr.ib(factory=list)
score_history = attr.ib(factory=list)
def beam_search(model, orig_item, preproc_item, beam_size, max_steps, visualize_flag=False):
inference_state, next_choices = model.begin_inference(orig_item, preproc_item)
beam = [Hypothesis(inference_state, next_choices)]
finished = []
for step in range(max_steps):
if visualize_flag:
print('step:')
print(step)
# Check if all beams are finished
if len(finished) == beam_size:
break
candidates = []
# For each hypothesis, get possible expansions
# Score each expansion
for hyp in beam:
candidates += [(hyp, choice, choice_score.item(),
hyp.score + choice_score.item())
for choice, choice_score in hyp.next_choices]
# Keep the top K expansions
candidates.sort(key=operator.itemgetter(3), reverse=True)
candidates = candidates[:beam_size - len(finished)]
# Create the new hypotheses from the expansions
beam = []
for hyp, choice, choice_score, cum_score in candidates:
inference_state = hyp.inference_state.clone()
next_choices = inference_state.step(choice)
if next_choices is None:
finished.append(Hypothesis(
inference_state,
None,
cum_score,
hyp.choice_history + [choice],
hyp.score_history + [choice_score]))
else:
beam.append(
Hypothesis(inference_state, next_choices, cum_score,
hyp.choice_history + [choice],
hyp.score_history + [choice_score]))
finished.sort(key=operator.attrgetter('score'), reverse=True)
return finished
|