def _get_initial_rnn_state(self, sentence: Dict[str, torch.LongTensor]):
embedded_input = self._sentence_embedder(sentence)
# (batch_size, sentence_length)
sentence_mask = util.get_text_field_mask(sentence).float()
batch_size = embedded_input.size(0)
# (batch_size, sentence_length, encoder_output_dim)
encoder_outputs = self._encoder(embedded_input, sentence_mask)
final_encoder_output = util.get_final_encoder_states(
encoder_outputs, sentence_mask, self._encoder.is_bidirectional())
memory_cell = util.new_variable_with_size(
encoder_outputs, (batch_size, self._encoder.get_output_dim()), 0)
attended_sentence = self._decoder_step.attend_on_sentence(
final_encoder_output, encoder_outputs, sentence_mask)
encoder_outputs_list = [encoder_outputs[i] for i in range(batch_size)]
sentence_mask_list = [sentence_mask[i] for i in range(batch_size)]
initial_rnn_state = []
for i in range(batch_size):
initial_rnn_state.append(
RnnState(final_encoder_output[i], memory_cell[i],
self._first_action_embedding, attended_sentence[i],
encoder_outputs_list, sentence_mask_list))
return initial_rnn_state
def forward(
self, # type: ignore
sentence: Dict[str, torch.LongTensor],
worlds: List[List[NlvrWorld]],
actions: List[List[ProductionRuleArray]],
agenda: torch.LongTensor,
labels: torch.LongTensor = None,
epoch_num: List[int] = None) -> Dict[str, torch.Tensor]:
# pylint: disable=arguments-differ
"""
Decoder logic for producing type constrained target sequences that maximize coverage of
their respective agendas, and minimize a denotation based loss.
"""
# We look at the epoch number and adjust the checklist cost weight if needed here.
instance_epoch_num = epoch_num[0] if epoch_num is not None else None
if self._dynamic_cost_rate is not None:
if self.training and instance_epoch_num is None:
raise RuntimeError(
"If you want a dynamic cost weight, use the "
"EpochTrackingBucketIterator!")
if instance_epoch_num != self._last_epoch_in_forward:
if instance_epoch_num >= self._dynamic_cost_wait_epochs:
decrement = self._checklist_cost_weight * self._dynamic_cost_rate
self._checklist_cost_weight -= decrement
logger.info("Checklist cost weight is now %f",
self._checklist_cost_weight)
self._last_epoch_in_forward = instance_epoch_num
batch_size = len(worlds)
action_embeddings, action_indices = self._embed_actions(actions)
initial_rnn_state = self._get_initial_rnn_state(sentence)
initial_score_list = [
util.new_variable_with_data(
list(sentence.values())[0], torch.Tensor([0.0]))
for i in range(batch_size)
]
# TODO (pradeep): Assuming all worlds give the same set of valid actions.
initial_grammar_state = [
self._create_grammar_state(worlds[i][0], actions[i])
for i in range(batch_size)
]
label_strings = self._get_label_strings(
labels) if labels is not None else None
# Each instance's agenda is of size (agenda_size, 1)
agenda_list = [agenda[i] for i in range(batch_size)]
initial_checklist_states = []
for instance_actions, instance_agenda in zip(actions, agenda_list):
checklist_info = self._get_checklist_info(instance_agenda,
instance_actions)
checklist_target, terminal_actions, checklist_mask = checklist_info
initial_checklist = util.new_variable_with_size(
checklist_target, checklist_target.size(), 0)
initial_checklist_states.append(
ChecklistState(terminal_actions=terminal_actions,
checklist_target=checklist_target,
checklist_mask=checklist_mask,
checklist=initial_checklist))
initial_state = NlvrDecoderState(
batch_indices=list(range(batch_size)),
action_history=[[] for _ in range(batch_size)],
score=initial_score_list,
rnn_state=initial_rnn_state,
grammar_state=initial_grammar_state,
action_embeddings=action_embeddings,
action_indices=action_indices,
possible_actions=actions,
worlds=worlds,
label_strings=label_strings,
checklist_state=initial_checklist_states)
agenda_data = [agenda_[:, 0].cpu().data for agenda_ in agenda_list]
outputs = self._decoder_trainer.decode(initial_state,
self._decoder_step,
self._get_state_cost)
best_action_sequences = outputs['best_action_sequences']
batch_action_strings = self._get_action_strings(
actions, best_action_sequences)
batch_denotations = self._get_denotations(batch_action_strings, worlds)
if labels is not None:
# We're either training or validating.
self._update_metrics(action_strings=batch_action_strings,
worlds=worlds,
label_strings=label_strings,
possible_actions=actions,
agenda_data=agenda_data)
else:
# We're testing.
outputs["best_action_strings"] = batch_action_strings
outputs["denotations"] = batch_denotations
return outputs