(model, orig_item, preproc_item, beam_size, max_steps, visualize_flag=False)
| 15 | |
| 16 | |
| 17 | def beam_search(model, orig_item, preproc_item, beam_size, max_steps, visualize_flag=False): |
| 18 | inference_state, next_choices = model.begin_inference(orig_item, preproc_item) |
| 19 | beam = [Hypothesis(inference_state, next_choices)] |
| 20 | finished = [] |
| 21 | |
| 22 | for step in range(max_steps): |
| 23 | if visualize_flag: |
| 24 | print('step:') |
| 25 | print(step) |
| 26 | # Check if all beams are finished |
| 27 | if len(finished) == beam_size: |
| 28 | break |
| 29 | |
| 30 | candidates = [] |
| 31 | |
| 32 | # For each hypothesis, get possible expansions |
| 33 | # Score each expansion |
| 34 | for hyp in beam: |
| 35 | candidates += [(hyp, choice, choice_score.item(), |
| 36 | hyp.score + choice_score.item()) |
| 37 | for choice, choice_score in hyp.next_choices] |
| 38 | |
| 39 | # Keep the top K expansions |
| 40 | candidates.sort(key=operator.itemgetter(3), reverse=True) |
| 41 | candidates = candidates[:beam_size - len(finished)] |
| 42 | |
| 43 | # Create the new hypotheses from the expansions |
| 44 | beam = [] |
| 45 | for hyp, choice, choice_score, cum_score in candidates: |
| 46 | inference_state = hyp.inference_state.clone() |
| 47 | next_choices = inference_state.step(choice) |
| 48 | if next_choices is None: |
| 49 | finished.append(Hypothesis( |
| 50 | inference_state, |
| 51 | None, |
| 52 | cum_score, |
| 53 | hyp.choice_history + [choice], |
| 54 | hyp.score_history + [choice_score])) |
| 55 | else: |
| 56 | beam.append( |
| 57 | Hypothesis(inference_state, next_choices, cum_score, |
| 58 | hyp.choice_history + [choice], |
| 59 | hyp.score_history + [choice_score])) |
| 60 | |
| 61 | finished.sort(key=operator.attrgetter('score'), reverse=True) |
| 62 | return finished |
nothing calls this directly
no test coverage detected