def _infer_one(self,
model,
data_item,
preproc_item,
beam_size,
output_history=False,
use_heuristic=True):
if isinstance(model.decoder, HeadCornerDecoder):
beams = spider_beam_search.head_corner_beam_search_with_heuristics(
model,
data_item,
preproc_item,
beam_size=beam_size,
max_steps=1000,
from_cond=False)
else:
if use_heuristic:
# TODO: from_cond should be true from non-bert model
beams = spider_beam_search.beam_search_with_heuristics(
model,
data_item,
preproc_item,
beam_size=beam_size,
max_steps=1000,
from_cond=False)
else:
beams = beam_search.beam_search(model,
data_item,
preproc_item,
beam_size=beam_size,
max_steps=1000)
decoded = []
for beam in beams:
model_output, inferred_code = beam.inference_state.finalize()
decoded.append({
'orig_question':
data_item.orig["question"],
'model_output':
model_output,
'inferred_code':
inferred_code,
'score':
beam.score,
**({
'choice_history': beam.choice_history,
'score_history': beam.score_history,
} if output_history else {})
})
return decoded
def _infer_one(self,
model,
data_item,
preproc_item,
beam_size,
output_history=False,
use_heuristic=True):
if use_heuristic:
# TODO: from_cond should be true from non-bert model
beams = spider_beam_search.beam_search_with_heuristics(
model,
data_item,
preproc_item,
beam_size=beam_size,
max_steps=1000,
from_cond=False)
elif use_heuristic == 1: #oracle sketch
beams = spider_beam_search.beam_search_with_oracle_sketch(
model,
data_item,
preproc_item,
beam_size=beam_size,
max_steps=1000)
else:
beams = beam_search.beam_search(model,
data_item,
preproc_item,
beam_size=beam_size,
max_steps=1000)
decoded = []
for beam in beams:
model_output, inferred_code = beam.inference_state.finalize(
gold=data_item.code, oracle=[])
_, oracle_select_inferred_code = beam.inference_state.finalize(
gold=data_item.code, oracle=['select'])
_, oracle_select_from_inferred_code = beam.inference_state.finalize(
gold=data_item.code, oracle=['select', 'from'])
decoded.append({
'orig_question':
data_item.orig["question"],
'orig_code':
data_item.code,
'preproc_item':
preproc_item[0],
'model_output':
model_output,
'inferred_code':
inferred_code,
'oracle_select_inferred_code':
oracle_select_inferred_code,
'oracle_select_from_inferred_code':
oracle_select_from_inferred_code,
'score':
beam.score,
**({
'choice_history': beam.choice_history,
'score_history': beam.score_history,
'attention_history': beam.attention_history,
'column_history': beam.column_history
} if output_history else {})
})
return decoded