| 107 | return decoded |
| 108 | |
| 109 | def _inner_infer(self, model, beam_size, output_history, sliced_orig_data, sliced_preproc_data, output, use_heuristic=False): |
| 110 | for i, (orig_item, preproc_item) in enumerate( |
| 111 | tqdm.tqdm(zip(sliced_orig_data, sliced_preproc_data), |
| 112 | total=len(sliced_orig_data))): |
| 113 | if use_heuristic: |
| 114 | #TODO: from_cond should be true from non-bert model |
| 115 | beams = spider_beam_search.beam_search_with_heuristics( |
| 116 | model, orig_item, preproc_item, beam_size=beam_size, max_steps=1000, from_cond=False) |
| 117 | else: |
| 118 | beams = beam_search.beam_search( |
| 119 | model, orig_item, preproc_item, beam_size=beam_size, max_steps=1000) |
| 120 | |
| 121 | decoded = [] |
| 122 | for beam in beams: |
| 123 | model_output, inferred_code = beam.inference_state.finalize() |
| 124 | |
| 125 | decoded.append({ |
| 126 | 'orig_question': orig_item.orig["question"], |
| 127 | 'model_output': model_output, |
| 128 | 'inferred_code': inferred_code, |
| 129 | 'score': beam.score, |
| 130 | **({ |
| 131 | 'choice_history': beam.choice_history, |
| 132 | 'score_history': beam.score_history, |
| 133 | } if output_history else {})}) |
| 134 | |
| 135 | output.write( |
| 136 | json.dumps({ |
| 137 | 'index': i, |
| 138 | 'beams': decoded, |
| 139 | }) + '\n') |
| 140 | output.flush() |
| 141 | |
| 142 | |
| 143 | def _debug(self, model, sliced_data, output): |