(model, orig_item, preproc_item, beam_size, max_steps)
| 33 | |
| 34 | |
| 35 | def beam_search(model, orig_item, preproc_item, beam_size, max_steps): |
| 36 | inference_state, next_choices = model.begin_inference(orig_item, preproc_item) |
| 37 | beam = [Hypothesis(inference_state, next_choices)] |
| 38 | finished = [] |
| 39 | |
| 40 | for step in range(max_steps): |
| 41 | # Check if all beams are finished |
| 42 | if len(finished) == beam_size: |
| 43 | break |
| 44 | |
| 45 | candidates = [] |
| 46 | |
| 47 | # For each hypothesis, get possible expansions |
| 48 | # Score each expansion |
| 49 | for hyp in beam: |
| 50 | candidates += [(hyp, choice, choice_score.item(), |
| 51 | hyp.score + choice_score.item()) |
| 52 | for choice, choice_score in hyp.next_choices] |
| 53 | |
| 54 | # Keep the top K expansions |
| 55 | candidates.sort(key=operator.itemgetter(3), reverse=True) |
| 56 | candidates = candidates[:beam_size - len(finished)] |
| 57 | |
| 58 | # Create the new hypotheses from the expansions |
| 59 | beam = [] |
| 60 | for hyp, choice, choice_score, cum_score in candidates: |
| 61 | inference_state = hyp.inference_state.clone() |
| 62 | next_choices = inference_state.step(choice) |
| 63 | if next_choices is None: |
| 64 | finished.append(Hypothesis( |
| 65 | inference_state, |
| 66 | None, |
| 67 | cum_score, |
| 68 | hyp.choice_history + [choice], |
| 69 | hyp.score_history + [choice_score])) |
| 70 | else: |
| 71 | beam.append( |
| 72 | Hypothesis(inference_state, next_choices, cum_score, |
| 73 | hyp.choice_history + [choice], |
| 74 | hyp.score_history + [choice_score])) |
| 75 | |
| 76 | finished.sort(key=operator.attrgetter('score'), reverse=True) |
| 77 | return finished |
no test coverage detected