def _get_state_cost(self, state: NlvrDecoderState) -> torch.Tensor:
"""
Return the costs a finished state. Since it is a finished state, the group size will be 1,
and hence we'll return just one cost.
"""
if not state.is_finished():
raise RuntimeError("_get_state_cost() is not defined for unfinished states!")
# Our checklist cost is a sum of squared error from where we want to be, making sure we
# take into account the mask.
checklist_balance = state.checklist_state[0].get_balance()
checklist_cost = torch.sum((checklist_balance) ** 2)
# This is the number of items on the agenda that we want to see in the decoded sequence.
# We use this as the denotation cost if the path is incorrect.
# Note: If we are penalizing the model for producing non-agenda actions, this is not the
# upper limit on the checklist cost. That would be the number of terminal actions.
denotation_cost = torch.sum(state.checklist_state[0].checklist_target.float())
checklist_cost = self._checklist_cost_weight * checklist_cost
# TODO (pradeep): The denotation based cost below is strict. May be define a cost based on
# how many worlds the logical form is correct in?
# label_strings being None happens when we are testing. We do not care about the cost then.
# TODO (pradeep): Make this cleaner.
if state.label_strings is None or all(self._check_state_denotations(state)):
cost = checklist_cost
else:
cost = checklist_cost + (1 - self._checklist_cost_weight) * denotation_cost
return cost
def _get_state_cost(self, state: NlvrDecoderState) -> torch.Tensor:
"""
Return the costs a finished state. Since it is a finished state, the group size will be 1,
and hence we'll return just one cost.
"""
if not state.is_finished():
raise RuntimeError("_get_state_cost() is not defined for unfinished states!")
# Our checklist cost is a sum of squared error from where we want to be, making sure we
# take into account the mask.
checklist_balance = state.checklist_state[0].get_balance()
checklist_cost = torch.sum((checklist_balance) ** 2)
# This is the number of items on the agenda that we want to see in the decoded sequence.
# We use this as the denotation cost if the path is incorrect.
# Note: If we are penalizing the model for producing non-agenda actions, this is not the
# upper limit on the checklist cost. That would be the number of terminal actions.
denotation_cost = torch.sum(state.checklist_state[0].checklist_target.float())
checklist_cost = self._checklist_cost_weight * checklist_cost
# TODO (pradeep): The denotation based cost below is strict. May be define a cost based on
# how many worlds the logical form is correct in?
# label_strings being None happens when we are testing. We do not care about the cost then.
# TODO (pradeep): Make this cleaner.
if state.label_strings is None or all(self._check_state_denotations(state)):
cost = checklist_cost
else:
cost = checklist_cost + (1 - self._checklist_cost_weight) * denotation_cost
return cost
def forward(
self, # type: ignore
sentence: Dict[str, torch.LongTensor],
worlds: List[List[NlvrWorld]],
actions: List[List[ProductionRuleArray]],
target_action_sequences: torch.LongTensor = None,
labels: torch.LongTensor = None) -> Dict[str, torch.Tensor]:
# pylint: disable=arguments-differ
"""
Decoder logic for producing type constrained target sequences, trained to maximize marginal
likelihod over a set of approximate logical forms.
"""
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)
]
label_strings = self._get_label_strings(
labels) if labels is not None else None
# 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)
]
worlds_list = [worlds[i] for i in range(batch_size)]
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_list,
label_strings=label_strings)
if target_action_sequences is not None:
# Remove the trailing dimension (from ListField[ListField[IndexField]]).
target_action_sequences = target_action_sequences.squeeze(-1)
target_mask = target_action_sequences != self._action_padding_index
else:
target_mask = None
outputs: Dict[str, torch.Tensor] = {}
if target_action_sequences is not None:
outputs = self._decoder_trainer.decode(
initial_state, self._decoder_step,
(target_action_sequences, target_mask))
best_final_states = self._decoder_beam_search.search(
self._max_decoding_steps,
initial_state,
self._decoder_step,
keep_final_unfinished_states=False)
best_action_sequences: Dict[int, List[List[int]]] = {}
for i in range(batch_size):
# Decoding may not have terminated with any completed logical forms, if `num_steps`
# isn't long enough (or if the model is not trained enough and gets into an
# infinite action loop).
if i in best_final_states:
best_action_indices = [
best_final_states[i][0].action_history[0]
]
best_action_sequences[i] = best_action_indices
batch_action_strings = self._get_action_strings(
actions, best_action_sequences)
batch_denotations = self._get_denotations(batch_action_strings, worlds)
if target_action_sequences is not None:
self._update_metrics(action_strings=batch_action_strings,
worlds=worlds,
label_strings=label_strings)
else:
outputs["best_action_strings"] = batch_action_strings
outputs["denotations"] = batch_denotations
return outputs
def take_step(
self, # type: ignore
state: NlvrDecoderState,
max_actions: int = None,
allowed_actions: List[Set[int]] = None) -> List[NlvrDecoderState]:
"""
Given a ``NlvrDecoderState``, returns a list of next states that are sorted by their scores.
This method is very similar to ``WikiTablesDecoderStep._take_step``. The differences are
that depending on the type of supervision being used, we may not have a notion of
"allowed actions" here, and we do not perform entity linking here.
"""
# Outline here: first we'll construct the input to the decoder, which is a concatenation of
# an embedding of the decoder input (the last action taken) and an attention over the
# sentence. Then we'll update our decoder's hidden state given this input, and recompute
# an attention over the sentence given our new hidden state. We'll use a concatenation of
# the new hidden state, the new attention, and optionally the checklist balance to predict an
# output, then yield new states. We will compute and use a checklist balance when
# ``allowed_actions`` is None, with the assumption that the ``DecoderTrainer`` that is
# calling this method is trying to train a parser without logical form supervision.
# TODO (pradeep): Make the distinction between the two kinds of trainers in the way they
# call this method more explicit.
# Each new state corresponds to one valid action that can be taken from the current state,
# and they are ordered by model scores.
attended_sentence = torch.stack(
[rnn_state.attended_input for rnn_state in state.rnn_state])
hidden_state = torch.stack(
[rnn_state.hidden_state for rnn_state in state.rnn_state])
memory_cell = torch.stack(
[rnn_state.memory_cell for rnn_state in state.rnn_state])
previous_action_embedding = torch.stack([
rnn_state.previous_action_embedding
for rnn_state in state.rnn_state
])
# (group_size, decoder_input_dim)
decoder_input = self._input_projection_layer(
torch.cat([attended_sentence, previous_action_embedding], -1))
decoder_input = torch.nn.functional.tanh(decoder_input)
hidden_state, memory_cell = self._decoder_cell(
decoder_input, (hidden_state, memory_cell))
hidden_state = self._dropout(hidden_state)
# (group_size, encoder_output_dim)
encoder_outputs = torch.stack([
state.rnn_state[0].encoder_outputs[i] for i in state.batch_indices
])
encoder_output_mask = torch.stack([
state.rnn_state[0].encoder_output_mask[i]
for i in state.batch_indices
])
attended_sentence = self.attend_on_sentence(hidden_state,
encoder_outputs,
encoder_output_mask)
# We get global indices of actions to embed here. The following logic is similar to
# ``WikiTablesDecoderStep._get_actions_to_consider``, except that we do not have any actions
# to link.
valid_actions = state.get_valid_actions()
global_valid_actions: List[List[Tuple[int, int]]] = []
for batch_index, valid_action_list in zip(state.batch_indices,
valid_actions):
global_valid_actions.append([])
for action_index in valid_action_list:
# state.action_indices is a dictionary that maps (batch_index, batch_action_index)
# to global_action_index
global_action_index = state.action_indices[(batch_index,
action_index)]
global_valid_actions[-1].append(
(global_action_index, action_index))
global_actions_to_embed: List[List[int]] = []
local_actions: List[List[int]] = []
for global_action_list in global_valid_actions:
global_action_list.sort()
global_actions_to_embed.append([])
local_actions.append([])
for global_action_index, action_index in global_action_list:
global_actions_to_embed[-1].append(global_action_index)
local_actions[-1].append(action_index)
max_num_actions = max(
[len(action_list) for action_list in global_actions_to_embed])
# We pad local actions with -1 as padding to get considered actions.
considered_actions = [
common_util.pad_sequence_to_length(action_list,
max_num_actions,
default_value=lambda: -1)
for action_list in local_actions
]
# action_embeddings: (group_size, num_embedded_actions, action_embedding_dim)
# action_mask: (group_size, num_embedded_actions)
action_embeddings, embedded_action_mask = self._get_action_embeddings(
state, global_actions_to_embed)
action_query = torch.cat([hidden_state, attended_sentence], dim=-1)
# (group_size, action_embedding_dim)
predicted_action_embedding = self._output_projection_layer(
action_query)
predicted_action_embedding = self._dropout(
torch.nn.functional.tanh(predicted_action_embedding))
if state.checklist_state[0] is not None:
embedding_addition = self._get_predicted_embedding_addition(state)
addition = embedding_addition * self._checklist_embedding_multiplier
predicted_action_embedding = predicted_action_embedding + addition
# We'll do a batch dot product here with `bmm`. We want `dot(predicted_action_embedding,
# action_embedding)` for each `action_embedding`, and we can get that efficiently with
# `bmm` and some squeezing.
# Shape: (group_size, num_embedded_actions)
action_logits = action_embeddings.bmm(
predicted_action_embedding.unsqueeze(-1)).squeeze(-1)
action_mask = embedded_action_mask.float()
if state.checklist_state[0] is not None:
# We will compute the logprobs and the checklists of potential next states together for
# efficiency.
logprobs, new_checklist_states = self._get_next_state_info_with_agenda(
state, considered_actions, action_logits, action_mask)
else:
logprobs = self._get_next_state_info_without_agenda(
state, considered_actions, action_logits, action_mask)
new_checklist_states = None
return self._compute_new_states(state, logprobs, hidden_state,
memory_cell, action_embeddings,
attended_sentence, considered_actions,
allowed_actions, new_checklist_states,
max_actions)
def _compute_new_states(cls,
state: NlvrDecoderState,
action_logprobs: List[List[Tuple[int,
torch.Tensor]]],
hidden_state: torch.Tensor,
memory_cell: torch.Tensor,
action_embeddings: torch.Tensor,
attended_sentence: torch.Tensor,
considered_actions: List[List[int]],
allowed_actions: List[Set[int]] = None,
new_checklist_states: List[
List[ChecklistState]] = None,
max_actions: int = None) -> List[NlvrDecoderState]:
"""
This method is very similar to ``WikiTabledDecoderStep._compute_new_states``.
The difference here is that we also keep track of checklists if they are passed to this
method.
"""
# batch_index -> group_index, action_index, checklist, score
states_info: Dict[int, List[Tuple[int, int, torch.Tensor,
torch.Tensor]]] = defaultdict(list)
if new_checklist_states is None:
# We do not have checklist states. Making a list of lists of Nones of the appropriate size for
# the zips below.
new_checklist_states = [[None for logprob in instance_logprobs]
for instance_logprobs in action_logprobs]
for group_index, instance_info in enumerate(
zip(state.batch_indices, action_logprobs,
new_checklist_states)):
batch_index, instance_logprobs, instance_new_checklist_states = instance_info
for (action_index,
score), checklist_state in zip(instance_logprobs,
instance_new_checklist_states):
states_info[batch_index].append(
(group_index, action_index, checklist_state, score))
new_states = []
for batch_index, instance_states_info in states_info.items():
batch_scores = torch.cat(
[info[-1] for info in instance_states_info])
_, sorted_indices = batch_scores.sort(-1, descending=True)
sorted_states_info = [
instance_states_info[i]
for i in sorted_indices.data.cpu().numpy()
]
allowed_states_info = []
for i, (group_index, action_index, _,
_) in enumerate(sorted_states_info):
action = considered_actions[group_index][action_index]
if allowed_actions is not None and action not in allowed_actions[
group_index]:
continue
allowed_states_info.append(sorted_states_info[i])
sorted_states_info = allowed_states_info
if max_actions is not None:
sorted_states_info = sorted_states_info[:max_actions]
for group_index, action_index, new_checklist_state, new_score in sorted_states_info:
# This is the actual index of the action from the original list of actions.
# We do not have to check whether it is the padding index because ``take_step``
# already took care of that.
action = considered_actions[group_index][action_index]
action_embedding = action_embeddings[group_index,
action_index, :]
new_action_history = state.action_history[group_index] + [
action
]
production_rule = state.possible_actions[batch_index][action][
0]
new_grammar_state = state.grammar_state[
group_index].take_action(production_rule)
new_rnn_state = RnnState(
hidden_state[group_index], memory_cell[group_index],
action_embedding, attended_sentence[group_index],
state.rnn_state[group_index].encoder_outputs,
state.rnn_state[group_index].encoder_output_mask)
new_state = NlvrDecoderState(
batch_indices=[batch_index],
action_history=[new_action_history],
score=[new_score],
rnn_state=[new_rnn_state],
grammar_state=[new_grammar_state],
action_embeddings=state.action_embeddings,
action_indices=state.action_indices,
possible_actions=state.possible_actions,
worlds=state.worlds,
label_strings=state.label_strings,
checklist_state=[new_checklist_state])
new_states.append(new_state)
return new_states
def forward(self, # type: ignore
sentence: Dict[str, torch.LongTensor],
worlds: List[List[NlvrWorld]],
actions: List[List[ProductionRuleArray]],
agenda: torch.LongTensor,
identifier: List[str] = None,
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 = [next(iter(sentence.values())).new_zeros(1, dtype=torch.float)
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 = checklist_target.new_zeros(checklist_target.size())
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)
if identifier is not None:
outputs['identifier'] = identifier
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
def take_step(self, # type: ignore
state: NlvrDecoderState,
max_actions: int = None,
allowed_actions: List[Set[int]] = None) -> List[NlvrDecoderState]:
"""
Given a ``NlvrDecoderState``, returns a list of next states that are sorted by their scores.
This method is very similar to ``WikiTablesDecoderStep._take_step``. The differences are
that depending on the type of supervision being used, we may not have a notion of
"allowed actions" here, and we do not perform entity linking here.
"""
# Outline here: first we'll construct the input to the decoder, which is a concatenation of
# an embedding of the decoder input (the last action taken) and an attention over the
# sentence. Then we'll update our decoder's hidden state given this input, and recompute
# an attention over the sentence given our new hidden state. We'll use a concatenation of
# the new hidden state, the new attention, and optionally the checklist balance to predict an
# output, then yield new states. We will compute and use a checklist balance when
# ``allowed_actions`` is None, with the assumption that the ``DecoderTrainer`` that is
# calling this method is trying to train a parser without logical form supervision.
# TODO (pradeep): Make the distinction between the two kinds of trainers in the way they
# call this method more explicit.
# Each new state corresponds to one valid action that can be taken from the current state,
# and they are ordered by model scores.
attended_sentence = torch.stack([rnn_state.attended_input for rnn_state in state.rnn_state])
hidden_state = torch.stack([rnn_state.hidden_state for rnn_state in state.rnn_state])
memory_cell = torch.stack([rnn_state.memory_cell for rnn_state in state.rnn_state])
previous_action_embedding = torch.stack([rnn_state.previous_action_embedding
for rnn_state in state.rnn_state])
# (group_size, decoder_input_dim)
decoder_input = self._input_projection_layer(torch.cat([attended_sentence,
previous_action_embedding], -1))
decoder_input = torch.nn.functional.tanh(decoder_input)
hidden_state, memory_cell = self._decoder_cell(decoder_input, (hidden_state, memory_cell))
hidden_state = self._dropout(hidden_state)
# (group_size, encoder_output_dim)
encoder_outputs = torch.stack([state.rnn_state[0].encoder_outputs[i] for i in state.batch_indices])
encoder_output_mask = torch.stack([state.rnn_state[0].encoder_output_mask[i] for i in state.batch_indices])
attended_sentence = self.attend_on_sentence(hidden_state, encoder_outputs, encoder_output_mask)
# We get global indices of actions to embed here. The following logic is similar to
# ``WikiTablesDecoderStep._get_actions_to_consider``, except that we do not have any actions
# to link.
valid_actions = state.get_valid_actions()
global_valid_actions: List[List[Tuple[int, int]]] = []
for batch_index, valid_action_list in zip(state.batch_indices, valid_actions):
global_valid_actions.append([])
for action_index in valid_action_list:
# state.action_indices is a dictionary that maps (batch_index, batch_action_index)
# to global_action_index
global_action_index = state.action_indices[(batch_index, action_index)]
global_valid_actions[-1].append((global_action_index, action_index))
global_actions_to_embed: List[List[int]] = []
local_actions: List[List[int]] = []
for global_action_list in global_valid_actions:
global_action_list.sort()
global_actions_to_embed.append([])
local_actions.append([])
for global_action_index, action_index in global_action_list:
global_actions_to_embed[-1].append(global_action_index)
local_actions[-1].append(action_index)
max_num_actions = max([len(action_list) for action_list in global_actions_to_embed])
# We pad local actions with -1 as padding to get considered actions.
considered_actions = [common_util.pad_sequence_to_length(action_list, max_num_actions,
default_value=lambda: -1)
for action_list in local_actions]
# action_embeddings: (group_size, num_embedded_actions, action_embedding_dim)
# action_mask: (group_size, num_embedded_actions)
action_embeddings, embedded_action_mask = self._get_action_embeddings(state,
global_actions_to_embed)
action_query = torch.cat([hidden_state, attended_sentence], dim=-1)
# (group_size, action_embedding_dim)
predicted_action_embedding = self._output_projection_layer(action_query)
predicted_action_embedding = self._dropout(torch.nn.functional.tanh(predicted_action_embedding))
if state.checklist_state[0] is not None:
embedding_addition = self._get_predicted_embedding_addition(state)
addition = embedding_addition * self._checklist_embedding_multiplier
predicted_action_embedding = predicted_action_embedding + addition
# We'll do a batch dot product here with `bmm`. We want `dot(predicted_action_embedding,
# action_embedding)` for each `action_embedding`, and we can get that efficiently with
# `bmm` and some squeezing.
# Shape: (group_size, num_embedded_actions)
action_logits = action_embeddings.bmm(predicted_action_embedding.unsqueeze(-1)).squeeze(-1)
action_mask = embedded_action_mask.float()
if state.checklist_state[0] is not None:
# We will compute the logprobs and the checklists of potential next states together for
# efficiency.
logprobs, new_checklist_states = self._get_next_state_info_with_agenda(state,
considered_actions,
action_logits,
action_mask)
else:
logprobs = self._get_next_state_info_without_agenda(state,
considered_actions,
action_logits,
action_mask)
new_checklist_states = None
return self._compute_new_states(state,
logprobs,
hidden_state,
memory_cell,
action_embeddings,
attended_sentence,
considered_actions,
allowed_actions,
new_checklist_states,
max_actions)