def _predictDistributions(
self, in_datas: List[TacticContext]) -> torch.FloatTensor:
assert self._tokenizer
assert self._embedding
assert self.training_args
goals_batch = [
normalizeSentenceLength(self._tokenizer.toTokenList(goal),
self.training_args.max_length)
for _, _, _, goal in in_datas
]
hyps = [
get_closest_hyp(hyps, goal, self.training_args.max_length)
for _, _, hyps, goal in in_datas
]
hyp_types = [serapi_instance.get_hyp_type(hyp) for hyp in hyps]
hyps_batch = [
normalizeSentenceLength(self._tokenizer.toTokenList(hyp_type),
self.training_args.max_length)
for hyp_type in hyp_types
]
word_features_batch = [
self._get_word_features(in_data) for in_data in in_datas
]
vec_features_batch = [
self._get_vec_features(in_data) for in_data in in_datas
]
stem_distribution = self._model(LongTensor(goals_batch),
LongTensor(hyps_batch),
FloatTensor(vec_features_batch),
LongTensor(word_features_batch))
return stem_distribution
def _predictStemDistributions(self, in_datas : List[TacticContext]) \
-> torch.FloatTensor:
word_features_batch = LongTensor(
[self._get_word_features(in_data) for in_data in in_datas])
vec_features_batch = FloatTensor(
[self._get_vec_features(in_data) for in_data in in_datas])
encoded_word_features = self._model.word_features_encoder(
word_features_batch)
stem_distribution = \
self._softmax(self._model.features_classifier(torch.cat((
encoded_word_features, vec_features_batch), dim=1)))
return stem_distribution
def predict_stems(self, k: int,
word_features: List[List[int]],
vec_features: List[List[float]]
) -> Tuple[torch.FloatTensor, torch.LongTensor]:
assert self._model
assert len(word_features) == len(vec_features)
batch_size = len(word_features)
stem_distribution = self._model.stem_classifier(
LongTensor(word_features), FloatTensor(vec_features))
stem_probs, stem_idxs = stem_distribution.topk(k)
assert stem_probs.size() == torch.Size([batch_size, k])
assert stem_idxs.size() == torch.Size([batch_size, k])
return stem_probs, stem_idxs
def hyp_name_scores(self,
stem_idxs: torch.LongTensor,
tokenized_goal: List[int],
tokenized_premises: List[List[int]],
premise_features: List[List[float]]
) -> torch.FloatTensor:
assert self._model
assert len(stem_idxs.size()) == 1
stem_width = stem_idxs.size()[0]
num_hyps = len(tokenized_premises)
encoded_goals = self._model.goal_encoder(LongTensor([tokenized_goal]))
hyp_arg_values = self.runHypModel(stem_idxs.unsqueeze(0),
encoded_goals,
LongTensor([tokenized_premises]),
FloatTensor([premise_features]))
assert hyp_arg_values.size() == torch.Size([1, stem_width, num_hyps])
return hyp_arg_values
def _predictDistribution(self, in_data : TacticContext) -> \
Tuple[torch.FloatTensor, str]:
if len(in_data.hypotheses) > 0:
relevant_hyp, relevance = \
max([(hyp,
term_relevance(in_data.goal,
serapi_instance.get_hyp_type(hyp)))
for hyp in in_data.hypotheses], key=lambda x: x[1])
else:
relevant_hyp = ":"
relevance = 0
encoded_hyp = self._encode_term(serapi_instance.get_hyp_type(relevant_hyp))
encoded_relevance = [relevance]
encoded_goal = self._encode_term(in_data.goal)
stem_distribution = self._run_model(encoded_hyp, encoded_relevance, encoded_goal)
return FloatTensor(stem_distribution), \
serapi_instance.get_first_var_in_hyp(relevant_hyp)
def _getBatchPredictionLoss(self, data_batch: Sequence[torch.Tensor],
model: CopyArgModel) -> torch.FloatTensor:
goals_batch, word_features_batch, vec_features_batch, \
stems_batch, arg_idxs_batch = \
cast(Tuple[torch.LongTensor, torch.FloatTensor,
torch.LongTensor, torch.LongTensor,
torch.LongTensor],
data_batch)
batch_size = goals_batch.size()[0]
encoded_word_features = model.word_features_encoder(
maybe_cuda(Variable(word_features_batch)))
catted_data = torch.cat(
(encoded_word_features, maybe_cuda(Variable(vec_features_batch))),
dim=1)
stemDistributions = model.features_classifier(catted_data)
stem_var = maybe_cuda(Variable(stems_batch)).view(batch_size)
argTokenIdxDistributions = model.find_arg_rnn(goals_batch, stems_batch)
argToken_var = maybe_cuda(Variable(arg_idxs_batch)).view(batch_size)
loss = FloatTensor([0.])
loss += self._criterion(stemDistributions, stem_var)
loss += self._criterion(argTokenIdxDistributions, argToken_var)
return loss
def _getBatchPredictionLoss(
self, arg_values: Namespace, batch: Sequence[torch.Tensor],
model: FeaturesPolyArgModel) -> torch.FloatTensor:
tokenized_hyp_types_batch, hyp_features_batch, num_hyps_batch, \
tokenized_goals_batch, goal_masks_batch, \
word_features_batch, vec_features_batch, \
stem_idxs_batch, arg_total_idxs_batch = \
cast(Tuple[torch.LongTensor, torch.FloatTensor, torch.LongTensor,
torch.LongTensor, torch.ByteTensor,
torch.LongTensor, torch.FloatTensor,
torch.LongTensor, torch.LongTensor],
data_batch)
batch_size = tokenized_goals_batch.size()[0]
goal_size = tokenized_goals_batch.size()[1]
stemDistributions = model.stem_classifier(word_features_batch,
vec_features_batch)
num_stem_poss = stemDistributions.size()[1]
stem_width = min(arg_values.max_beam_width, num_stem_poss)
stem_var = maybe_cuda(Variable(stem_idxs_batch))
predictedProbs, predictedStemIdxs = stemDistributions.topk(stem_width)
mergedStemIdxs = []
for stem_idx, predictedStemIdxList in zip(stem_idxs_batch,
predictedStemIdxs):
if stem_idx.item() in predictedStemIdxList:
mergedStemIdxs.append(predictedStemIdxList)
else:
mergedStemIdxs.append(
torch.cat((maybe_cuda(stem_idx.view(1)),
predictedStemIdxList[:stem_width - 1])))
mergedStemIdxsT = torch.stack(mergedStemIdxs)
correctPredictionIdxs = torch.LongTensor([
list(idxList).index(stem_idx)
for idxList, stem_idx in zip(mergedStemIdxs, stem_var)
])
if arg_values.hyp_rnn:
tokenized_hyps_var = maybe_cuda(
Variable(tokenized_hyp_types_batch))
else:
tokenized_hyps_var = maybe_cuda(
Variable(torch.zeros_like(tokenized_hyp_types_batch)))
if arg_values.hyp_features:
hyp_features_var = maybe_cuda(Variable(hyp_features_batch))
else:
hyp_features_var = maybe_cuda(
Variable(torch.zeros_like(hyp_features_batch)))
goal_arg_values = model.goal_args_model(
mergedStemIdxsT.view(batch_size * stem_width),
tokenized_goals_batch.view(batch_size, 1, goal_size).expand(-1, stem_width, -1)
.contiguous().view(batch_size * stem_width, goal_size))\
.view(batch_size, stem_width, goal_size + 1)
goal_arg_values = torch.where(
maybe_cuda(
goal_masks_batch.view(batch_size, 1,
arg_values.max_length + 1)).expand(
-1, stem_width, -1), goal_arg_values,
maybe_cuda(torch.full_like(goal_arg_values, -float("Inf"))))
encoded_goals = model.goal_encoder(tokenized_goals_batch)
hyp_lists_length = tokenized_hyp_types_batch.size()[1]
hyp_length = tokenized_hyp_types_batch.size()[2]
hyp_features_size = hyp_features_batch.size()[2]
encoded_goal_size = encoded_goals.size()[1]
encoded_goals_expanded = \
encoded_goals.view(batch_size, 1, 1, encoded_goal_size)\
.expand(-1, stem_width, hyp_lists_length, -1).contiguous()\
.view(batch_size * stem_width * hyp_lists_length, encoded_goal_size)
if not arg_values.goal_rnn:
encoded_goals_expanded = torch.zeros_like(encoded_goals_expanded)
stems_expanded = \
mergedStemIdxsT.view(batch_size, stem_width, 1)\
.expand(-1, -1, hyp_lists_length).contiguous()\
.view(batch_size * stem_width * hyp_lists_length)
hyp_arg_values_concatted = \
model.hyp_model(stems_expanded,
encoded_goals_expanded,
tokenized_hyps_var
.view(batch_size, 1, hyp_lists_length, hyp_length)
.expand(-1, stem_width, -1, -1).contiguous()
.view(batch_size * stem_width * hyp_lists_length,
hyp_length),
hyp_features_var
.view(batch_size, 1, hyp_lists_length, hyp_features_size)
.expand(-1, stem_width, -1, -1).contiguous()
.view(batch_size * stem_width * hyp_lists_length,
hyp_features_size))
assert hyp_arg_values_concatted.size() == torch.Size(
[batch_size * stem_width * hyp_lists_length,
1]), hyp_arg_values_concatted.size()
hyp_arg_values = hyp_arg_values_concatted.view(batch_size, stem_width,
hyp_lists_length)
total_arg_values = torch.cat((goal_arg_values, hyp_arg_values), dim=2)
num_probs = hyp_lists_length + goal_size + 1
total_arg_distribution = \
self._softmax(total_arg_values.view(batch_size, stem_width * num_probs))
total_arg_var = maybe_cuda(Variable(arg_total_idxs_batch +
(correctPredictionIdxs * num_probs)))\
.view(batch_size)
loss = FloatTensor([0.])
loss += self._criterion(stemDistributions, stem_var)
loss += self._criterion(total_arg_distribution, total_arg_var)
return loss
def predictKTactics(self, context: TacticContext,
k: int) -> List[Prediction]:
assert self.training_args
assert self._model
num_stem_poss = get_num_tokens(self._metadata)
stem_width = min(self.training_args.max_beam_width, num_stem_poss,
k**2)
tokenized_premises, hyp_features, \
nhyps_batch, tokenized_goal, \
goal_mask, \
word_features, vec_features = \
sample_fpa(extract_dataloader_args(self.training_args),
self._metadata,
context.relevant_lemmas,
context.prev_tactics,
context.hypotheses,
context.goal)
num_hyps = nhyps_batch[0]
stem_distribution = self._model.stem_classifier(
LongTensor(word_features), FloatTensor(vec_features))
stem_certainties, stem_idxs = stem_distribution.topk(stem_width)
goals_batch = LongTensor(tokenized_goal)
goal_arg_values = self._model.goal_args_model(
stem_idxs.view(1 * stem_width),
goals_batch.view(1, 1, self.training_args.max_length)
.expand(-1, stem_width, -1).contiguous()
.view(1 * stem_width,
self.training_args.max_length))\
.view(1, stem_width,
self.training_args.max_length + 1)
goal_arg_values = torch.where(
torch.ByteTensor(goal_mask).view(1, 1,
self.training_args.max_length +
1).expand(-1, stem_width, -1),
goal_arg_values, torch.full_like(goal_arg_values, -float("Inf")))
assert goal_arg_values.size() == torch.Size([1, stem_width,
self.training_args.max_length + 1]),\
"goal_arg_values.size(): {}; stem_width: {}".format(goal_arg_values.size(),
stem_width)
num_probs = 1 + num_hyps + self.training_args.max_length
goal_symbols = get_fpa_words(context.goal)
num_valid_probs = (1 + num_hyps + len(goal_symbols)) * stem_width
if num_hyps > 0:
encoded_goals = self._model.goal_encoder(goals_batch)\
.view(1, 1, self.training_args.hidden_size)
hyps_batch = LongTensor(tokenized_premises)
assert hyps_batch.size() == torch.Size([1, num_hyps,
self.training_args.max_length]), \
(hyps_batch.size(),
num_hyps,
self.training_args.max_length)
hypfeatures_batch = FloatTensor(hyp_features)
assert hypfeatures_batch.size() == \
torch.Size([1, num_hyps, hypFeaturesSize()]), \
(hypfeatures_batch.size(), num_hyps, hypFeaturesSize())
hyp_arg_values = self.runHypModel(stem_idxs, encoded_goals,
hyps_batch, hypfeatures_batch)
assert hyp_arg_values.size() == \
torch.Size([1, stem_width, num_hyps])
total_values = torch.cat((goal_arg_values, hyp_arg_values), dim=2)
else:
total_values = goal_arg_values
all_prob_batches = self._softmax((total_values +
stem_certainties.view(1, stem_width, 1)
.expand(-1, -1, num_probs))
.contiguous()
.view(1, stem_width * num_probs))\
.view(stem_width * num_probs)
final_probs, final_idxs = all_prob_batches.topk(k)
assert not torch.isnan(final_probs).any()
assert final_probs.size() == torch.Size([k])
row_length = self.training_args.max_length + num_hyps + 1
stem_keys = final_idxs // row_length
assert stem_keys.size() == torch.Size([k])
assert stem_idxs.size() == torch.Size([1,
stem_width]), stem_idxs.size()
prediction_stem_idxs = stem_idxs.view(stem_width).index_select(
0, stem_keys)
assert prediction_stem_idxs.size() == torch.Size([k]), \
prediction_stem_idxs.size()
arg_idxs = final_idxs % row_length
assert arg_idxs.size() == torch.Size([k])
if self.training_args.lemma_args:
all_hyps = context.hypotheses + context.relevant_lemmas
else:
all_hyps = context.hypotheses
return [
Prediction(
decode_fpa_result(extract_dataloader_args(self.training_args),
self._metadata, all_hyps, context.goal,
stem_idx.item(), arg_idx.item()),
math.exp(prob)) for stem_idx, arg_idx, prob in islice(
zip(prediction_stem_idxs, arg_idxs, final_probs),
min(k, num_valid_probs))
]
def predictKTactics_batch(self, context_batch: List[TacticContext],
k: int) \
-> List[List[Prediction]]:
assert self.training_args
assert self._model
num_stem_poss = get_num_tokens(self._metadata)
stem_width = min(self.training_args.max_beam_width, num_stem_poss,
k**2)
batch_size = len(context_batch)
tprems_batch, pfeat_batch, \
nhyps_batch, tgoals_batch, \
goal_masks_batch, \
wfeats_batch, vfeats_batch = \
sample_fpa_batch(extract_dataloader_args(self.training_args),
self._metadata,
[context_py2r(context)
for context in context_batch])
for tprem, pfeat, nhyp, tgoal, masks, wfeat, vfeat, context \
in zip(tprems_batch, pfeat_batch,
nhyps_batch, tgoals_batch,
goal_masks_batch, wfeats_batch,
vfeats_batch, context_batch):
s_tprem, s_pfeat, s_nhyp, s_tgoal, s_masks, s_wfeat, s_vfeat = \
sample_fpa(extract_dataloader_args(self.training_args),
self._metadata,
context.relevant_lemmas,
context.prev_tactics,
context.hypotheses,
context.goal)
assert len(s_tprem) == 1
assert len(tprem) == len(s_tprem[0])
for p1, p2 in zip(tprem, s_tprem[0]):
assert p1 == p2, (p1, p2)
assert len(s_pfeat) == 1
assert len(pfeat) == len(s_pfeat[0])
for f1, f2 in zip(pfeat, s_pfeat[0]):
assert f1 == f2, (f1, f2)
assert s_nhyp[0] == nhyp, (s_nhyp[0], nhyp)
assert s_tgoal[0] == tgoal, (s_tgoal[0], tgoal)
assert s_masks[0] == masks, (s_masks[0], masks)
assert s_wfeat[0] == wfeat, (s_wfeat[0], wfeat)
assert s_vfeat[0] == vfeat, (s_vfeat[0], vfeat)
stem_distribution = self._model.stem_classifier(
LongTensor(wfeats_batch), FloatTensor(vfeats_batch))
stem_certainties_batch, stem_idxs_batch = stem_distribution.topk(
stem_width)
goals_batch = LongTensor(tgoals_batch)
goal_arg_values = self._model.goal_args_model(
stem_idxs_batch.view(batch_size * stem_width),
goals_batch.view(batch_size, 1, self.training_args.max_length)
.expand(-1, stem_width, -1).contiguous()
.view(batch_size * stem_width,
self.training_args.max_length))\
.view(batch_size, stem_width,
self.training_args.max_length + 1)
goal_arg_values = torch.where(
torch.ByteTensor(goal_masks_batch).view(
batch_size, 1,
self.training_args.max_length + 1).expand(-1, stem_width, -1),
goal_arg_values, torch.full_like(goal_arg_values, -float("Inf")))
encoded_goals_batch = self._model.goal_encoder(goals_batch)
stems_expanded_batch = torch.cat([
stem_idxs.view(1, stem_width).expand(
num_hyps, stem_width).contiguous().view(num_hyps * stem_width)
for stem_idxs, num_hyps, in zip(stem_idxs_batch, nhyps_batch)
])
egoals_expanded_batch = torch.cat([
encoded_goal.view(1, self.training_args.hidden_size).expand(
num_hyps * stem_width, -1).contiguous()
for encoded_goal, num_hyps in zip(encoded_goals_batch, nhyps_batch)
])
tprems_expanded_batch = torch.cat([
LongTensor(tpremises).view(
1, -1, self.training_args.max_length).expand(
stem_width, -1,
-1).contiguous().view(-1, self.training_args.max_length)
for tpremises in tprems_batch
])
pfeat_expanded_batch = torch.cat([
FloatTensor(premise_features).view(1, num_hyps, 2).expand(
stem_width, -1, -1).contiguous().view(num_hyps * stem_width, 2)
for premise_features, num_hyps in zip(pfeat_batch, nhyps_batch)
])
prem_arg_values = self._model.hyp_model(stems_expanded_batch,
egoals_expanded_batch,
tprems_expanded_batch,
pfeat_expanded_batch)
prem_arg_values_split = prem_arg_values.split(
[num_hyps * stem_width for num_hyps in nhyps_batch])
total_values_list = [
torch.cat((goal_values, prem_values.view(stem_width, -1)),
dim=1) for goal_values, prem_values in zip(
goal_arg_values, prem_arg_values_split)
]
all_probs_list = [
self._softmax((total_values + stem_certainties.view(
stem_width, 1).expand_as(total_values)).contiguous().view(
1, -1)).view(-1) for total_values, stem_certainties in zip(
total_values_list, stem_certainties_batch)
]
final_probs_list, final_idxs_list = zip(
*[probs.topk(k) for probs in all_probs_list])
stem_keys_list = [
final_idxs // (self.training_args.max_length + num_hyps + 1)
for final_idxs, num_hyps in zip(final_idxs_list, nhyps_batch)
]
stem_idxs_list = [
stem_idxs.view(stem_width).index_select(0, stem_keys)
for stem_idxs, stem_keys in zip(stem_idxs_batch, stem_keys_list)
]
arg_idxs_list = [
final_idxs % (self.training_args.max_length + num_hyps + 1)
for final_idxs, num_hyps in zip(final_idxs_list, nhyps_batch)
]
predictions = [[
Prediction(
decode_fpa_result(extract_dataloader_args(self.training_args),
self._metadata,
context.hypotheses + context.relevant_lemmas,
context.goal,
stem_idx.item(), arg_idx.item()),
math.exp(prob)) for stem_idx, arg_idx, prob in islice(
zip(stem_idxs, arg_idxs, final_probs),
min(k, 1 + num_hyps + len(get_fpa_words(context.goal))))
] for stem_idxs, arg_idxs, final_probs, context, num_hyps in zip(
stem_idxs_list, arg_idxs_list, final_probs_list, context_batch,
nhyps_batch)]
for context, pred_list in zip(context_batch, predictions):
for batch_pred, single_pred in zip(
pred_list, self.predictKTactics(context, k)):
assert batch_pred.prediction == single_pred.prediction, \
(batch_pred, single_pred)
return predictions
def predict(self, inputs: List[List[float]]) -> List[List[float]]:
with self._lock:
return self._model(FloatTensor(inputs)).data
def decodeKTactics(decoder : DecoderRNN, encoder_hidden : torch.FloatTensor,
beam_width : int, max_length : int):
v = decoder.output_size
pos_index = Variable(LongTensor([0]) * beam_width).view(-1, 1)
hidden = _inflate(encoder_hidden, beam_width)
sequence_scores = FloatTensor(beam_width, 1)
sequence_scores.fill_(-float('Inf'))
sequence_scores.index_fill_(0, LongTensor([0]), 0.0)
sequence_scores = Variable(sequence_scores)
input_var = Variable(LongTensor([[SOS_token] * beam_width]))
stored_predecessors = list()
stored_emitted_symbols = list()
for j in range(max_length):
decoder_output, hidden = decoder(input_var, hidden)
sequence_scores = _inflate(sequence_scores, v)
sequence_scores += decoder_output
scores, candidates = sequence_scores.view(1, -1).topk(beam_width)
input_var = (candidates % v).view(1, beam_width)
sequence_scores = scores.view(beam_width, 1)
predecessors = (candidates / v +
pos_index.expand_as(candidates)).view(beam_width, 1)
hidden = hidden.index_select(1, cast(torch.LongTensor, predecessors.squeeze()))
eos_indices = input_var.data.eq(EOS_token)
if eos_indices.nonzero().dim() > 0:
sequence_scores.data.masked_fill_(torch.transpose(eos_indices, 0, 1),
-float('inf'))
stored_predecessors.append(predecessors)
stored_emitted_symbols.append(torch.transpose(input_var, 0, 1))
# Trace back from the final three highest scores
_, next_idxs = sequence_scores.view(beam_width).sort(descending=True)
seqs = [] # type: List[List[SomeLongTensor]]
eos_found = 0
for i in range(max_length - 1, -1, -1):
# The next column of symbols from the end
next_symbols = stored_emitted_symbols[i].view(beam_width) \
.index_select(0, next_idxs).data
# The predecessors of that column
next_idxs = stored_predecessors[i].view(beam_width).index_select(0, next_idxs)
# Handle sequences that ended early
eos_indices = stored_emitted_symbols[i].data.squeeze(1).eq(EOS_token).nonzero()
if eos_indices.dim() > 0:
for j in range(eos_indices.size(0)-1, -1, -1):
idx = eos_indices[j]
res_k_idx = beam_width - (eos_found % beam_width) - 1
eos_found += 1
res_idx = res_k_idx
next_idxs[res_idx] = stored_predecessors[i][idx[0]]
next_symbols[res_idx] = stored_emitted_symbols[i][idx[0]].data[0]
# Commit the result
seqs.insert(0, next_symbols)
# Transpose
int_seqs = [[data[i][0] for data in seqs] for i in range(beam_width)]
# Cut off EOS tokens
int_seqs = [list(takewhile(lambda x: x != EOS_token, seq)) for seq in int_seqs]
return int_seqs
def _run_model(self, hyp : List[float], rel : List[float], goal : List[float]) -> \
torch.FloatTensor:
# return FloatTensor(self._model.predict_log_proba([hyp])[0])
# return FloatTensor(self._model.predict_log_proba([goal])[0])
return FloatTensor(self._model.predict([list(hyp) + rel + list(goal)])[0])