def _map_entity_productions(linking_scores: torch.FloatTensor,
worlds: List[WikiTablesWorld],
actions: List[List[ProductionRuleArray]]) -> Tuple[torch.Tensor,
Dict[Tuple[int, int], int]]:
"""
Constructs a map from ``(batch_index, action_index)`` to ``(batch_index * entity_index)``.
That is, some actions correspond to terminal productions of entities from our table. We
need to find those actions and map them to their corresponding entity indices, where the
entity index is its position in the list of entities returned by the ``world``. This list
is what defines the second dimension of the ``linking_scores`` tensor, so we can use this
index to look up linking scores for each action in that tensor.
For easier processing later, the mapping that we return is `flattened` - we really want to
map ``(batch_index, action_index)`` to ``(batch_index, entity_index)``, but we are going to
have to use the result of this mapping to do ``index_selects`` on the ``linking_scores``
tensor. You can't do ``index_select`` with tuples, so we flatten ``linking_scores`` to
have shape ``(batch_size * num_entities, num_question_tokens)``, and return shifted indices
into this flattened tensor.
Parameters
----------
linking_scores : ``torch.Tensor``
A tensor representing linking scores between each table entity and each question token.
Has shape ``(batch_size, num_entities, num_question_tokens)``.
worlds : ``List[WikiTablesWorld]``
The ``World`` for each batch instance. The ``World`` contains a reference to the
``TableKnowledgeGraph`` that defines the set of entities in the linking.
actions : ``List[List[ProductionRuleArray]]``
The list of possible actions for each batch instance. Our action indices are defined
in terms of this list, so we'll find entity productions in this list and map them to
entity indices from the entity list we get from the ``World``.
Returns
-------
flattened_linking_scores : ``torch.Tensor``
A flattened version of ``linking_scores``, with shape ``(batch_size * num_entities,
num_question_tokens)``.
actions_to_entities : ``Dict[Tuple[int, int], int]``
A mapping from ``(batch_index, action_index)`` to ``(batch_size * num_entities)``,
representing which action indices correspond to which entity indices in the returned
``flattened_linking_scores`` tensor.
"""
batch_size, num_entities, num_question_tokens = linking_scores.size()
entity_map: Dict[Tuple[int, str], int] = {}
for batch_index, world in enumerate(worlds):
for entity_index, entity in enumerate(world.table_graph.entities):
entity_map[(batch_index, entity)] = batch_index * num_entities + entity_index
actions_to_entities: Dict[Tuple[int, int], int] = {}
for batch_index, action_list in enumerate(actions):
for action_index, action in enumerate(action_list):
if not action[0]:
# This action is padding.
continue
_, production = action[0].split(' -> ')
entity_index = entity_map.get((batch_index, production), None)
if entity_index is not None:
actions_to_entities[(batch_index, action_index)] = entity_index
flattened_linking_scores = linking_scores.view(batch_size * num_entities, num_question_tokens)
return flattened_linking_scores, actions_to_entities
def sequence_cross_entropy_with_logits(logits: torch.FloatTensor,
targets: torch.LongTensor,
weights: torch.FloatTensor,
batch_average: bool = True) -> torch.FloatTensor:
"""
Computes the cross entropy loss of a sequence, weighted with respect to
some user provided weights. Note that the weighting here is not the same as
in the :func:`torch.nn.CrossEntropyLoss()` criterion, which is weighting
classes; here we are weighting the loss contribution from particular elements
in the sequence. This allows loss computations for models which use padding.
Parameters
----------
logits : ``torch.FloatTensor``, required.
A ``torch.FloatTensor`` of size (batch_size, sequence_length, num_classes)
which contains the unnormalized probability for each class.
targets : ``torch.LongTensor``, required.
A ``torch.LongTensor`` of size (batch, sequence_length) which contains the
index of the true class for each corresponding step.
weights : ``torch.FloatTensor``, required.
A ``torch.FloatTensor`` of size (batch, sequence_length)
batch_average : bool, optional, (default = True).
A bool indicating whether the loss should be averaged across the batch,
or returned as a vector of losses per batch element.
Returns
-------
A torch.FloatTensor representing the cross entropy loss.
If ``batch_average == True``, the returned loss is a scalar.
If ``batch_average == False``, the returned loss is a vector of shape (batch_size,).
"""
# shape : (batch * sequence_length, num_classes)
logits_flat = logits.view(-1, logits.size(-1))
# shape : (batch * sequence_length, num_classes)
log_probs_flat = torch.nn.functional.log_softmax(logits_flat)
# shape : (batch * max_len, 1)
targets_flat = targets.view(-1, 1).long()
# Contribution to the negative log likelihood only comes from the exact indices
# of the targets, as the target distributions are one-hot. Here we use torch.gather
# to extract the indices of the num_classes dimension which contribute to the loss.
# shape : (batch * sequence_length, 1)
negative_log_likelihood_flat = - torch.gather(log_probs_flat, dim=1, index=targets_flat)
# shape : (batch, sequence_length)
negative_log_likelihood = negative_log_likelihood_flat.view(*targets.size())
# shape : (batch, sequence_length)
negative_log_likelihood = negative_log_likelihood * weights.float()
# shape : (batch_size,)
per_batch_loss = negative_log_likelihood.sum(1) / (weights.sum(1).float() + 1e-13)
if batch_average:
num_non_empty_sequences = ((weights.sum(1) > 0).float().sum() + 1e-13)
return per_batch_loss.sum() / num_non_empty_sequences
return per_batch_loss
def _compute_span_pair_embeddings(self,
top_span_embeddings: torch.FloatTensor,
antecedent_embeddings: torch.FloatTensor,
antecedent_offsets: torch.FloatTensor):
"""
Computes an embedding representation of pairs of spans for the pairwise scoring function
to consider. This includes both the original span representations, the element-wise
similarity of the span representations, and an embedding representation of the distance
between the two spans.
Parameters
----------
top_span_embeddings : ``torch.FloatTensor``, required.
Embedding representations of the top spans. Has shape
(batch_size, num_spans_to_keep, embedding_size).
antecedent_embeddings : ``torch.FloatTensor``, required.
Embedding representations of the antecedent spans we are considering
for each top span. Has shape
(batch_size, num_spans_to_keep, max_antecedents, embedding_size).
antecedent_offsets : ``torch.IntTensor``, required.
The offsets between each top span and its antecedent spans in terms
of spans we are considering. Has shape (1, max_antecedents).
Returns
-------
span_pair_embeddings : ``torch.FloatTensor``
Embedding representation of the pair of spans to consider. Has shape
(batch_size, num_spans_to_keep, max_antecedents, embedding_size)
"""
# Shape: (batch_size, num_spans_to_keep, max_antecedents, embedding_size)
target_embeddings = top_span_embeddings.unsqueeze(2).expand_as(antecedent_embeddings)
# Shape: (1, max_antecedents, embedding_size)
antecedent_distance_embeddings = self._distance_embedding(
util.bucket_values(antecedent_offsets,
num_total_buckets=self._num_distance_buckets))
# Shape: (1, 1, max_antecedents, embedding_size)
antecedent_distance_embeddings = antecedent_distance_embeddings.unsqueeze(0)
expanded_distance_embeddings_shape = (antecedent_embeddings.size(0),
antecedent_embeddings.size(1),
antecedent_embeddings.size(2),
antecedent_distance_embeddings.size(-1))
# Shape: (batch_size, num_spans_to_keep, max_antecedents, embedding_size)
antecedent_distance_embeddings = antecedent_distance_embeddings.expand(*expanded_distance_embeddings_shape)
# Shape: (batch_size, num_spans_to_keep, max_antecedents, embedding_size)
span_pair_embeddings = torch.cat([target_embeddings,
antecedent_embeddings,
antecedent_embeddings * target_embeddings,
antecedent_distance_embeddings], -1)
return span_pair_embeddings
def show(cls,
W: torch.FloatTensor,
b: torch.FloatTensor):
desc = 'y ='
W = W.view(-1)
for i, w in enumerate(W):
desc += ' {:+.2f} x^{}'.format(w, len(W) - i)
desc += ' {:+.2f}'.format(b[0])
return desc
def __init__(self, vocab_size, embedding_size, weight: torch.FloatTensor = None,
trainable=True, scale_by_embedding_size=False):
super(Embedding, self).__init__()
self.output_dim = embedding_size
self.scale_by_embedding_size = scale_by_embedding_size
if weight is None:
weight = torch.FloatTensor(vocab_size, embedding_size)
self.weight = torch.nn.Parameter(weight, requires_grad=trainable)
torch.nn.init.xavier_uniform_(self.weight)
else:
assert weight.size() == (vocab_size, embedding_size)
self.weight = torch.nn.Parameter(weight, requires_grad=trainable)
def __init__(self,
num_embeddings: int,
embedding_dim: int,
projection_dim: int = None,
weight: torch.FloatTensor = None,
padding_index: int = None,
trainable: bool = True,
max_norm: float = None,
norm_type: float = 2.,
scale_grad_by_freq: bool = False,
sparse: bool = False) -> None:
super(Embedding, self).__init__()
self.num_embeddings = num_embeddings
self.padding_index = padding_index
self.max_norm = max_norm
self.norm_type = norm_type
self.scale_grad_by_freq = scale_grad_by_freq
self.sparse = sparse
self.output_dim = projection_dim or embedding_dim
if weight is None:
weight = torch.FloatTensor(num_embeddings, embedding_dim)
self.weight = torch.nn.Parameter(weight, requires_grad=trainable)
torch.nn.init.xavier_uniform_(self.weight)
else:
if weight.size() != (num_embeddings, embedding_dim):
raise ConfigurationError("A weight matrix was passed with contradictory embedding shapes.")
self.weight = torch.nn.Parameter(weight, requires_grad=trainable)
if self.padding_index is not None:
self.weight.data[self.padding_index].fill_(0)
if projection_dim:
self._projection = torch.nn.Linear(embedding_dim, projection_dim)
else:
self._projection = None
def _get_linking_probabilities(self,
worlds: List[WikiTablesWorld],
linking_scores: torch.FloatTensor,
question_mask: torch.LongTensor,
entity_type_dict: Dict[int, int]) -> torch.FloatTensor:
"""
Produces the probability of an entity given a question word and type. The logic below
separates the entities by type since the softmax normalization term sums over entities
of a single type.
Parameters
----------
worlds : ``List[WikiTablesWorld]``
linking_scores : ``torch.FloatTensor``
Has shape (batch_size, num_question_tokens, num_entities).
question_mask: ``torch.LongTensor``
Has shape (batch_size, num_question_tokens).
entity_type_dict : ``Dict[int, int]``
This is a mapping from ((batch_index * num_entities) + entity_index) to entity type id.
Returns
-------
batch_probabilities : ``torch.FloatTensor``
Has shape ``(batch_size, num_question_tokens, num_entities)``.
Contains all the probabilities for an entity given a question word.
"""
_, num_question_tokens, num_entities = linking_scores.size()
batch_probabilities = []
for batch_index, world in enumerate(worlds):
all_probabilities = []
num_entities_in_instance = 0
# NOTE: The way that we're doing this here relies on the fact that entities are
# implicitly sorted by their types when we sort them by name, and that numbers come
# before "fb:cell", and "fb:cell" comes before "fb:row". This is not a great
# assumption, and could easily break later, but it should work for now.
for type_index in range(self._num_entity_types):
# This index of 0 is for the null entity for each type, representing the case where a
# word doesn't link to any entity.
entity_indices = [0]
entities = world.table_graph.entities
for entity_index, _ in enumerate(entities):
if entity_type_dict[batch_index * num_entities + entity_index] == type_index:
entity_indices.append(entity_index)
if len(entity_indices) == 1:
# No entities of this type; move along...
continue
# We're subtracting one here because of the null entity we added above.
num_entities_in_instance += len(entity_indices) - 1
# We separate the scores by type, since normalization is done per type. There's an
# extra "null" entity per type, also, so we have `num_entities_per_type + 1`. We're
# selecting from a (num_question_tokens, num_entities) linking tensor on _dimension 1_,
# so we get back something of shape (num_question_tokens,) for each index we're
# selecting. All of the selected indices together then make a tensor of shape
# (num_question_tokens, num_entities_per_type + 1).
indices = linking_scores.new_tensor(entity_indices, dtype=torch.long)
entity_scores = linking_scores[batch_index].index_select(1, indices)
# We used index 0 for the null entity, so this will actually have some values in it.
# But we want the null entity's score to be 0, so we set that here.
entity_scores[:, 0] = 0
# No need for a mask here, as this is done per batch instance, with no padding.
type_probabilities = torch.nn.functional.softmax(entity_scores, dim=1)
all_probabilities.append(type_probabilities[:, 1:])
# We need to add padding here if we don't have the right number of entities.
if num_entities_in_instance != num_entities:
zeros = linking_scores.new_zeros(num_question_tokens,
num_entities - num_entities_in_instance)
all_probabilities.append(zeros)
# (num_question_tokens, num_entities)
probabilities = torch.cat(all_probabilities, dim=1)
batch_probabilities.append(probabilities)
batch_probabilities = torch.stack(batch_probabilities, dim=0)
return batch_probabilities * question_mask.unsqueeze(-1).float()
def forward(self, # pylint: disable=arguments-differ
inputs: torch.FloatTensor,
batch_lengths: List[int],
initial_state: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
"""
Parameters
----------
inputs : ``torch.FloatTensor``, required.
A tensor of shape (batch_size, num_timesteps, input_size)
to apply the LSTM over.
batch_lengths : ``List[int]``, required.
A list of length batch_size containing the lengths of the sequences in batch.
initial_state : ``Tuple[torch.Tensor, torch.Tensor]``, optional, (default = None)
A tuple (state, memory) representing the initial hidden state and memory
of the LSTM. The ``state`` has shape (1, batch_size, hidden_size) and the
``memory`` has shape (1, batch_size, cell_size).
Returns
-------
output_accumulator : ``torch.FloatTensor``
The outputs of the LSTM for each timestep. A tensor of shape
(batch_size, max_timesteps, hidden_size) where for a given batch
element, all outputs past the sequence length for that batch are
zero tensors.
final_state : ``Tuple[``torch.FloatTensor, torch.FloatTensor]``
A tuple (state, memory) representing the initial hidden state and memory
of the LSTM. The ``state`` has shape (1, batch_size, hidden_size) and the
``memory`` has shape (1, batch_size, cell_size).
"""
batch_size = inputs.size()[0]
total_timesteps = inputs.size()[1]
output_accumulator = inputs.new_zeros(batch_size, total_timesteps, self.hidden_size)
if initial_state is None:
full_batch_previous_memory = inputs.new_zeros(batch_size, self.cell_size)
full_batch_previous_state = inputs.new_zeros(batch_size, self.hidden_size)
else:
full_batch_previous_state = initial_state[0].squeeze(0)
full_batch_previous_memory = initial_state[1].squeeze(0)
current_length_index = batch_size - 1 if self.go_forward else 0
if self.recurrent_dropout_probability > 0.0 and self.training:
dropout_mask = get_dropout_mask(self.recurrent_dropout_probability,
full_batch_previous_state)
else:
dropout_mask = None
for timestep in range(total_timesteps):
# The index depends on which end we start.
index = timestep if self.go_forward else total_timesteps - timestep - 1
# What we are doing here is finding the index into the batch dimension
# which we need to use for this timestep, because the sequences have
# variable length, so once the index is greater than the length of this
# particular batch sequence, we no longer need to do the computation for
# this sequence. The key thing to recognise here is that the batch inputs
# must be _ordered_ by length from longest (first in batch) to shortest
# (last) so initially, we are going forwards with every sequence and as we
# pass the index at which the shortest elements of the batch finish,
# we stop picking them up for the computation.
if self.go_forward:
while batch_lengths[current_length_index] <= index:
current_length_index -= 1
# If we're going backwards, we are _picking up_ more indices.
else:
# First conditional: Are we already at the maximum number of elements in the batch?
# Second conditional: Does the next shortest sequence beyond the current batch
# index require computation use this timestep?
while current_length_index < (len(batch_lengths) - 1) and \
batch_lengths[current_length_index + 1] > index:
current_length_index += 1
# Actually get the slices of the batch which we
# need for the computation at this timestep.
# shape (batch_size, cell_size)
previous_memory = full_batch_previous_memory[0: current_length_index + 1].clone()
# Shape (batch_size, hidden_size)
previous_state = full_batch_previous_state[0: current_length_index + 1].clone()
# Shape (batch_size, input_size)
timestep_input = inputs[0: current_length_index + 1, index]
# Do the projections for all the gates all at once.
# Both have shape (batch_size, 4 * cell_size)
projected_input = self.input_linearity(timestep_input)
projected_state = self.state_linearity(previous_state)
# Main LSTM equations using relevant chunks of the big linear
# projections of the hidden state and inputs.
input_gate = torch.sigmoid(projected_input[:, (0 * self.cell_size):(1 * self.cell_size)] +
projected_state[:, (0 * self.cell_size):(1 * self.cell_size)])
forget_gate = torch.sigmoid(projected_input[:, (1 * self.cell_size):(2 * self.cell_size)] +
projected_state[:, (1 * self.cell_size):(2 * self.cell_size)])
memory_init = torch.tanh(projected_input[:, (2 * self.cell_size):(3 * self.cell_size)] +
projected_state[:, (2 * self.cell_size):(3 * self.cell_size)])
output_gate = torch.sigmoid(projected_input[:, (3 * self.cell_size):(4 * self.cell_size)] +
projected_state[:, (3 * self.cell_size):(4 * self.cell_size)])
memory = input_gate * memory_init + forget_gate * previous_memory
# Here is the non-standard part of this LSTM cell; first, we clip the
# memory cell, then we project the output of the timestep to a smaller size
# and again clip it.
if self.memory_cell_clip_value:
# pylint: disable=invalid-unary-operand-type
memory = torch.clamp(memory, -self.memory_cell_clip_value, self.memory_cell_clip_value)
# shape (current_length_index, cell_size)
pre_projection_timestep_output = output_gate * torch.tanh(memory)
# shape (current_length_index, hidden_size)
timestep_output = self.state_projection(pre_projection_timestep_output)
if self.state_projection_clip_value:
# pylint: disable=invalid-unary-operand-type
timestep_output = torch.clamp(timestep_output,
-self.state_projection_clip_value,
self.state_projection_clip_value)
# Only do dropout if the dropout prob is > 0.0 and we are in training mode.
if dropout_mask is not None:
timestep_output = timestep_output * dropout_mask[0: current_length_index + 1]
# We've been doing computation with less than the full batch, so here we create a new
# variable for the the whole batch at this timestep and insert the result for the
# relevant elements of the batch into it.
full_batch_previous_memory = full_batch_previous_memory.clone()
full_batch_previous_state = full_batch_previous_state.clone()
full_batch_previous_memory[0:current_length_index + 1] = memory
full_batch_previous_state[0:current_length_index + 1] = timestep_output
output_accumulator[0:current_length_index + 1, index] = timestep_output
# Mimic the pytorch API by returning state in the following shape:
# (num_layers * num_directions, batch_size, ...). As this
# LSTM cell cannot be stacked, the first dimension here is just 1.
final_state = (full_batch_previous_state.unsqueeze(0),
full_batch_previous_memory.unsqueeze(0))
return output_accumulator, final_state
def sequence_cross_entropy_with_logits(logits: torch.FloatTensor,
targets: torch.LongTensor,
weights: torch.FloatTensor,
batch_average: bool = True,
label_smoothing: float = None) -> torch.FloatTensor:
"""
Computes the cross entropy loss of a sequence, weighted with respect to
some user provided weights. Note that the weighting here is not the same as
in the :func:`torch.nn.CrossEntropyLoss()` criterion, which is weighting
classes; here we are weighting the loss contribution from particular elements
in the sequence. This allows loss computations for models which use padding.
Parameters
----------
logits : ``torch.FloatTensor``, required.
A ``torch.FloatTensor`` of size (batch_size, sequence_length, num_classes)
which contains the unnormalized probability for each class.
targets : ``torch.LongTensor``, required.
A ``torch.LongTensor`` of size (batch, sequence_length) which contains the
index of the true class for each corresponding step.
weights : ``torch.FloatTensor``, required.
A ``torch.FloatTensor`` of size (batch, sequence_length)
batch_average : bool, optional, (default = True).
A bool indicating whether the loss should be averaged across the batch,
or returned as a vector of losses per batch element.
label_smoothing : ``float``, optional (default = None)
Whether or not to apply label smoothing to the cross-entropy loss.
For example, with a label smoothing value of 0.2, a 4 class classifcation
target would look like ``[0.05, 0.05, 0.85, 0.05]`` if the 3rd class was
the correct label.
Returns
-------
A torch.FloatTensor representing the cross entropy loss.
If ``batch_average == True``, the returned loss is a scalar.
If ``batch_average == False``, the returned loss is a vector of shape (batch_size,).
"""
# shape : (batch * sequence_length, num_classes)
logits_flat = logits.view(-1, logits.size(-1))
# shape : (batch * sequence_length, num_classes)
log_probs_flat = torch.nn.functional.log_softmax(logits_flat, dim=-1)
# shape : (batch * max_len, 1)
targets_flat = targets.view(-1, 1).long()
if label_smoothing is not None and label_smoothing > 0.0:
num_classes = logits.size(-1)
smoothing_value = label_smoothing / num_classes
# Fill all the correct indices with 1 - smoothing value.
one_hot_targets = torch.zeros_like(log_probs_flat).scatter_(-1, targets_flat, 1.0 - label_smoothing)
smoothed_targets = one_hot_targets + smoothing_value
negative_log_likelihood_flat = - log_probs_flat * smoothed_targets
negative_log_likelihood_flat = negative_log_likelihood_flat.sum(-1, keepdim=True)
else:
# Contribution to the negative log likelihood only comes from the exact indices
# of the targets, as the target distributions are one-hot. Here we use torch.gather
# to extract the indices of the num_classes dimension which contribute to the loss.
# shape : (batch * sequence_length, 1)
negative_log_likelihood_flat = - torch.gather(log_probs_flat, dim=1, index=targets_flat)
# shape : (batch, sequence_length)
negative_log_likelihood = negative_log_likelihood_flat.view(*targets.size())
# shape : (batch, sequence_length)
negative_log_likelihood = negative_log_likelihood * weights.float()
# shape : (batch_size,)
per_batch_loss = negative_log_likelihood.sum(1) / (weights.sum(1).float() + 1e-13)
if batch_average:
num_non_empty_sequences = ((weights.sum(1) > 0).float().sum() + 1e-13)
return per_batch_loss.sum() / num_non_empty_sequences
return per_batch_loss
def _viterbi_decode(self, emissions: torch.FloatTensor,
mask: torch.ByteTensor) -> List[List[int]]:
# emissions: (seq_length, batch_size, num_tags)
# mask: (seq_length, batch_size)
assert emissions.dim() == 3 and mask.dim() == 2
assert emissions.shape[:2] == mask.shape
assert emissions.size(2) == self.num_tags
assert mask[0].all()
seq_length, batch_size = mask.shape
# Start transition and first emission
# shape: (batch_size, num_tags)
score = self.start_transitions + emissions[0]
history = []
# score is a tensor of size (batch_size, num_tags) where for every batch,
# value at column j stores the score of the best tag sequence so far that ends
# with tag j
# history saves where the best tags candidate transitioned from; this is used
# when we trace back the best tag sequence
# Viterbi algorithm recursive case: we compute the score of the best tag sequence
# for every possible next tag
for i in range(1, seq_length):
# Broadcast viterbi score for every possible next tag
# shape: (batch_size, num_tags, 1)
broadcast_score = score.unsqueeze(2)
# Broadcast emission score for every possible current tag
# shape: (batch_size, 1, num_tags)
broadcast_emission = emissions[i].unsqueeze(1)
# Compute the score tensor of size (batch_size, num_tags, num_tags) where
# for each sample, entry at row i and column j stores the score of the best
# tag sequence so far that ends with transitioning from tag i to tag j and emitting
# shape: (batch_size, num_tags, num_tags)
next_score = broadcast_score + self.transitions + broadcast_emission
# Find the maximum score over all possible current tag
# shape: (batch_size, num_tags)
next_score, indices = next_score.max(dim=1)
# Set score to the next score if this timestep is valid (mask == 1)
# and save the index that produces the next score
# shape: (batch_size, num_tags)
score = torch.where(mask[i].unsqueeze(1), next_score, score)
history.append(indices)
# End transition score
# shape: (batch_size, num_tags)
score += self.end_transitions
# Now, compute the best path for each sample
# shape: (batch_size,)
seq_ends = mask.long().sum(dim=0) - 1
best_tags_list = []
for idx in range(batch_size):
# Find the tag which maximizes the score at the last timestep; this is our best tag
# for the last timestep
_, best_last_tag = score[idx].max(dim=0)
best_tags = [best_last_tag.item()]
# We trace back where the best last tag comes from, append that to our best tag
# sequence, and trace it back again, and so on
for hist in reversed(history[:seq_ends[idx]]):
best_last_tag = hist[idx][best_tags[-1]]
best_tags.append(best_last_tag.item())
# Reverse the order because we start from the last timestep
best_tags.reverse()
best_tags_list.append(best_tags)
return best_tags_list
def forward(self, x: torch.FloatTensor) -> torch.FloatTensor: # noqa: D102
value = x.abs().pow(self.p).sum(dim=self.dim).mean()
if not self.normalize:
return value
dim = torch.as_tensor(x.shape[-1], dtype=torch.float, device=x.device)
return value / dim
def imshow(x: torch.FloatTensor):
x = x[0] if len(x.shape) == 4 else x
x = x.permute(1, 2, 0)
plt.imshow(x)
def _verify_initialization(self,
x: torch.FloatTensor) -> None: # noqa: D102
xn = x.norm(dim=-1)
assert torch.allclose(xn, torch.ones_like(xn))
def conve_interaction(
h: torch.FloatTensor,
r: torch.FloatTensor,
t: torch.FloatTensor,
t_bias: torch.FloatTensor,
input_channels: int,
embedding_height: int,
embedding_width: int,
hr2d: nn.Module,
hr1d: nn.Module,
) -> torch.FloatTensor:
"""Evaluate the ConvE interaction function.
:param h: shape: (batch_size, num_heads, 1, 1, dim)
The head representations.
:param r: shape: (batch_size, 1, num_relations, 1, dim)
The relation representations.
:param t: shape: (batch_size, 1, 1, num_tails, dim)
The tail representations.
:param t_bias: shape: (batch_size, 1, 1, num_tails, 1)
The tail entity bias.
:param input_channels:
The number of input channels.
:param embedding_height:
The height of the reshaped embedding.
:param embedding_width:
The width of the reshaped embedding.
:param hr2d:
The first module, transforming the 2D stacked head-relation "image".
:param hr1d:
The second module, transforming the 1D flattened output of the 2D module.
:return: shape: (batch_size, num_heads, num_relations, num_tails)
The scores.
"""
# repeat if necessary, and concat head and relation, batch_size', num_input_channels, 2*height, width
# with batch_size' = batch_size * num_heads * num_relations
x = broadcast_cat(
[
h.view(*h.shape[:-1], input_channels, embedding_height,
embedding_width),
r.view(*r.shape[:-1], input_channels, embedding_height,
embedding_width),
],
dim=-2,
).view(-1, input_channels, 2 * embedding_height, embedding_width)
# batch_size', num_input_channels, 2*height, width
x = hr2d(x)
# batch_size', num_output_channels * (2 * height - kernel_height + 1) * (width - kernel_width + 1)
x = x.view(-1, numpy.prod(x.shape[-3:]))
x = hr1d(x)
# reshape: (batch_size', embedding_dim) -> (b, h, r, 1, d)
x = x.view(-1, h.shape[1], r.shape[2], 1, h.shape[-1])
# For efficient calculation, each of the convolved [h, r] rows has only to be multiplied with one t row
# output_shape: (batch_size, num_heads, num_relations, num_tails)
t = t.transpose(-1, -2)
x = (x @ t).squeeze(dim=-2)
# add bias term
return x + t_bias.squeeze(dim=-1)
def convkb_interaction(
h: torch.FloatTensor,
r: torch.FloatTensor,
t: torch.FloatTensor,
conv: nn.Conv2d,
activation: nn.Module,
hidden_dropout: nn.Dropout,
linear: nn.Linear,
) -> torch.FloatTensor:
r"""Evaluate the ConvKB interaction function.
.. math::
W_L drop(act(W_C \ast ([h; r; t]) + b_C)) + b_L
:param h: shape: (batch_size, num_heads, 1, 1, dim)
The head representations.
:param r: shape: (batch_size, 1, num_relations, 1, dim)
The relation representations.
:param t: shape: (batch_size, 1, 1, num_tails, dim)
The tail representations.
:param conv:
The 3x1 convolution.
:param activation:
The activation function.
:param hidden_dropout:
The dropout layer applied to the hidden activations.
:param linear:
The final linear layer.
:return: shape: (batch_size, num_heads, num_relations, num_tails)
The scores.
"""
# decompose convolution for faster computation in 1-n case
num_filters = conv.weight.shape[0]
assert conv.weight.shape == (num_filters, 1, 1, 3)
# compute conv(stack(h, r, t))
# prepare input shapes for broadcasting
# (b, h, r, t, 1, d)
h = h.unsqueeze(dim=-2)
r = r.unsqueeze(dim=-2)
t = t.unsqueeze(dim=-2)
# conv.weight.shape = (C_out, C_in, kernel_size[0], kernel_size[1])
# here, kernel_size = (1, 3), C_in = 1, C_out = num_filters
# -> conv_head, conv_rel, conv_tail shapes: (num_filters,)
# reshape to (1, 1, 1, 1, f, 1)
conv_head, conv_rel, conv_tail, conv_bias = [
c.view(1, 1, 1, 1, num_filters, 1)
for c in list(conv.weight[:, 0, 0, :].t()) + [conv.bias]
]
# convolve -> output.shape: (*, embedding_dim, num_filters)
h = conv_head @ h
r = conv_rel @ r
t = conv_tail @ t
x = tensor_sum(conv_bias, h, r, t)
x = activation(x)
# Apply dropout, cf. https://github.com/daiquocnguyen/ConvKB/blob/master/model.py#L54-L56
x = hidden_dropout(x)
# Linear layer for final scores; use flattened representations, shape: (b, h, r, t, d * f)
x = x.view(*x.shape[:-2], -1)
x = linear(x)
return x.squeeze(dim=-1)
def generate_positive_labels(size: int, smooth: bool):
if smooth:
return to_gpu_if_available(FloatTensor(size).uniform_(0.9, 1))
else:
return to_gpu_if_available(ones(size))
def forward(self, # pylint: disable=arguments-differ
inputs: torch.FloatTensor,
batch_lengths: List[int],
initial_state: Optional[Tuple[torch.Tensor, torch.Tensor]] = None):
"""
Parameters
----------
inputs : ``torch.FloatTensor``, required.
A tensor of shape (batch_size, num_timesteps, input_size)
to apply the LSTM over.
"For multidirectional input, transpose the position of each input first. With forward tag.
If we don't fine tune the params, we only use the final states, then we can transpose the final states
back to the original position."
batch_lengths : ``List[int]``, required.
A list of length batch_size containing the lengths of the sequences in batch.
initial_state : ``Tuple[torch.Tensor, torch.Tensor]``, optional, (default = None)
A tuple (state, memory) representing the initial hidden state and memory
of the LSTM. The ``state`` has shape (1, batch_size, hidden_size) and the
``memory`` has shape (1, batch_size, cell_size).
Returns
-------
output_accumulator : ``torch.FloatTensor``
The outputs of the LSTM for each timestep. A tensor of shape
(batch_size, max_timesteps, hidden_size) where for a given batch
element, all outputs past the sequence length for that batch are
zero tensors.
final_state : ``Tuple[``torch.FloatTensor, torch.FloatTensor]``
A tuple (state, memory) representing the initial hidden state and memory
of the LSTM. The ``state`` has shape (1, batch_size, hidden_size) and the
``memory`` has shape (1, batch_size, cell_size).
"""
batch_size = inputs.size()[0]
total_timesteps = inputs.size()[1]
output_accumulator = inputs.new_zeros(batch_size, total_timesteps, self.hidden_size)
if initial_state is None:
full_batch_previous_memory = inputs.new_zeros(batch_size, self.cell_size)
full_batch_previous_state = inputs.new_zeros(batch_size, self.hidden_size)
else:
full_batch_previous_state = initial_state[0].squeeze(0)
full_batch_previous_memory = initial_state[1].squeeze(0)
current_length_index = batch_size - 1 if self.go_forward else 0
if self.recurrent_dropout_probability > 0.0 and self.training:
dropout_mask = get_dropout_mask(self.recurrent_dropout_probability,
full_batch_previous_state)
else:
dropout_mask = None
for timestep in range(total_timesteps):
# The index depends on which end we start.
index = timestep if self.go_forward else total_timesteps - timestep - 1
# index to indicate forward or backward
# What we are doing here is finding the index into the batch dimension
# which we need to use for this timestep, because the sequences have
# variable length, so once the index is greater than the length of this
# particular batch sequence, we no longer need to do the computation for
# this sequence. The key thing to recognise here is that the batch inputs
# must be _ordered_ by length from longest (first in batch) to shortest
# (last) so initially, we are going forwards with every sequence and as we
# pass the index at which the shortest elements of the batch finish,
# we stop picking them up for the computation.
if self.go_forward:
while batch_lengths[current_length_index] <= index:
current_length_index -= 1
# If we're going backwards, we are _picking up_ more indices.
else:
# First conditional: Are we already at the maximum number of elements in the batch?
# Second conditional: Does the next shortest sequence beyond the current batch
# index require computation use this timestep?
while current_length_index < (len(batch_lengths) - 1) and \
batch_lengths[current_length_index + 1] > index:
current_length_index += 1
# Actually get the slices of the batch which we
# need for the computation at this timestep.
# shape (batch_size, cell_size)
previous_memory = full_batch_previous_memory[0: current_length_index + 1].clone()
# Shape (batch_size, hidden_size)
previous_state = full_batch_previous_state[0: current_length_index + 1].clone()
# Shape (batch_size, input_size)
timestep_input = inputs[0: current_length_index + 1, index]
# Do the projections for all the gates all at once.
# Both have shape (batch_size, 4 * cell_size)
projected_input = self.input_linearity(timestep_input)
projected_state = self.state_linearity(previous_state)
# Main LSTM equations using relevant chunks of the big linear
# projections of the hidden state and inputs.
input_gate = torch.sigmoid(projected_input[:, (0 * self.cell_size):(1 * self.cell_size)] +
projected_state[:, (0 * self.cell_size):(1 * self.cell_size)])
forget_gate = torch.sigmoid(projected_input[:, (1 * self.cell_size):(2 * self.cell_size)] +
projected_state[:, (1 * self.cell_size):(2 * self.cell_size)])
memory_init = torch.tanh(projected_input[:, (2 * self.cell_size):(3 * self.cell_size)] +
projected_state[:, (2 * self.cell_size):(3 * self.cell_size)])
output_gate = torch.sigmoid(projected_input[:, (3 * self.cell_size):(4 * self.cell_size)] +
projected_state[:, (3 * self.cell_size):(4 * self.cell_size)])
memory = input_gate * memory_init + forget_gate * previous_memory
# Here is the non-standard part of this LSTM cell; first, we clip the
# memory cell, then we project the output of the timestep to a smaller size
# and again clip it.
if self.memory_cell_clip_value:
# pylint: disable=invalid-unary-operand-type
memory = torch.clamp(memory, -self.memory_cell_clip_value, self.memory_cell_clip_value)
# shape (current_length_index, cell_size)
pre_projection_timestep_output = output_gate * torch.tanh(memory)
# shape (current_length_index, hidden_size)
timestep_output = self.state_projection(pre_projection_timestep_output)
if self.state_projection_clip_value:
# pylint: disable=invalid-unary-operand-type
timestep_output = torch.clamp(timestep_output,
-self.state_projection_clip_value,
self.state_projection_clip_value)
# Only do dropout if the dropout prob is > 0.0 and we are in training mode.
if dropout_mask is not None:
timestep_output = timestep_output * dropout_mask[0: current_length_index + 1]
# We've been doing computation with less than the full batch, so here we create a new
# variable for the the whole batch at this timestep and insert the result for the
# relevant elements of the batch into it.
full_batch_previous_memory = full_batch_previous_memory.clone()
full_batch_previous_state = full_batch_previous_state.clone()
full_batch_previous_memory[0:current_length_index + 1] = memory
full_batch_previous_state[0:current_length_index + 1] = timestep_output
output_accumulator[0:current_length_index + 1, index] = timestep_output
# Mimic the pytorch API by returning state in the following shape:
# (num_layers * num_directions, batch_size, ...). As this
# LSTM cell cannot be stacked, the first dimension here is just 1.
final_state = (full_batch_previous_state.unsqueeze(0),
full_batch_previous_memory.unsqueeze(0))
return output_accumulator, final_state
def forward(self, input):
assert input.dim() == 2
self.input = input
return max(input, FloatTensor([0]))
def forward(self,
sequence_tensor: torch.FloatTensor,
span_indices: torch.LongTensor,
sequence_mask: torch.LongTensor = None,
span_indices_mask: torch.LongTensor = None,
dep_masks: torch.IntTensor = None,
span_labels: torch.IntTensor = None) -> torch.FloatTensor:
# both of shape (batch_size, num_spans, 1)
span_starts, span_ends = span_indices.split(1, dim=-1)
# Shape: (batch_size, num_spans, 1)
has_span_mask = (span_starts > -1).float()
# These span widths are off by 1, because the span ends are `inclusive`.
span_widths = span_ends - span_starts
# We need to know the maximum span width so we can
# generate indices to extract the spans from the sequence tensor.
# These indices will then get masked below, such that if the length
# of a given span is smaller than the max, the rest of the values
# are masked.
max_batch_span_width = span_widths.max().item() + 1
# Shape: (1, 1, max_batch_span_width)
max_span_range_indices = util.get_range_vector(
max_batch_span_width,
util.get_device_of(sequence_tensor)).view(1, 1, -1)
# Shape: (batch_size, num_spans, max_batch_span_width)
# This is a broadcasted comparison - for each span we are considering,
# we are creating a range vector of size max_span_width, but masking values
# which are greater than the actual length of the span.
#
# We're using <= here (and for the mask below) because the span ends are
# inclusive, so we want to include indices which are equal to span_widths rather
# than using it as a non-inclusive upper bound.
# Shape: (batch_size, num_spans, max_batch_span_width)
span_mask = (max_span_range_indices <= span_widths).float()
span_mask = span_mask * has_span_mask
raw_span_indices = span_starts + max_span_range_indices
# We also don't want to include span indices which are less than zero,
# which happens because some spans near the beginning of the sequence
# have an end index < max_batch_span_width, so we add this to the mask here.
# span_mask = span_mask * (raw_span_indices >= 0).float()
span_indices = torch.nn.functional.relu(
raw_span_indices.float()).long()
span_indices = span_indices * span_mask.long()
# Shape: (batch_size * num_spans * max_batch_span_width)
flat_span_indices = flatten_and_batch_shift_indices(
span_indices, sequence_tensor.size(1))
# Shape: (batch_size * num_spans, max_batch_span_width, embedding_dim)
batch_size, num_spans, max_batch_span_width = span_indices.size()
span_embeddings = util.batched_index_select(sequence_tensor, span_indices, flat_span_indices)\
.view(batch_size * num_spans, max_batch_span_width, -1)
# Shape: (batch_size * num_spans, max_batch_sent_length, embedding_dim)
# attended_span_embeddings = self._span_self_attentive_encoder(span_embeddings,
# span_mask.view(batch_size * num_spans, -1))
attended_span_embeddings = span_embeddings
# Shape: (batch_size * num_spans, embedding_dim)
max_pooled_span_emb = torch.max(attended_span_embeddings, dim=1)[0]
# Shape: (batch_size * num_spans, 2)
gate_logit = self.modeled_gate(max_pooled_span_emb)
gate = F.softmax(gate_logit, dim=-1)
gate = torch.chunk(gate, 2, dim=-1)[1].view(batch_size * num_spans, -1)
# positive_labels = span_labels.long() * has_span_mask.long().squeeze()
# negative_labels = (1 - span_labels.long()) * has_span_mask.long().squeeze()
# positive_num = torch.sum(positive_labels)
# negative_num = torch.sum(negative_labels)
# print(positive_labels, negative_labels)
# print(positive_num, negative_num)
# print(torch.sum(gate.float().view(-1) * positive_labels.float().view(-1)) / positive_num.float())
# print(torch.sum(gate.float().view(-1) * negative_labels.float().view(-1)) / negative_num.float())
# print(torch.sum(positive_labels.view(-1) * gate.view(-1)))
dumb_gate = torch.full_like(
gate.view(batch_size * num_spans, -1).float(), 0.3)
new_gate = torch.cat([dumb_gate, gate], dim=-1)
# gate = (gate >= 0.3).long()
# attended_span_embeddings = attended_span_embeddings * gate.unsqueeze(-1).float()
loss = F.nll_loss(F.log_softmax(gate_logit,
dim=-1).view(batch_size * num_spans,
-1),
span_labels.long().view(batch_size * num_spans),
ignore_index=-1)
print('\n strong sup loss', loss)
# We split the context into sentences and now we want to reconstruct the context using the sentences
# Shape: (batch_size, num_sentences, max_batch_sent_length, embedding_dim) ->
# (batch_size, max_batch_context_length, embedding_dim )
recovered_indices = []
for slice_embs, m in zip(
attended_span_embeddings.view(batch_size, num_spans,
max_batch_span_width, -1),
span_mask.byte()):
# print('selected_shape:', torch.masked_select(slice_embs, m.unsqueeze(-1)).shape)
# print(torch.masked_select(slice_embs, m.unsqueeze(-1)).view(-1, 200).shape)
recovered_indices.append(
torch.masked_select(slice_embs, m.unsqueeze(-1)))
recovered_context_representation = torch.nn.utils.rnn.pad_sequence(
recovered_indices, batch_first=True)
recovered_context_representation = recovered_context_representation.view(
batch_size, sequence_tensor.size(1), -1)
return recovered_context_representation, loss, new_gate
def _interaction_function(
self,
h: torch.FloatTensor,
t: torch.FloatTensor,
r_indices: Optional[torch.LongTensor] = None,
) -> torch.FloatTensor:
#: Prepare h: (b, e, d) -> (b, e, 1, 1, d)
h_for_w = h.unsqueeze(dim=-2).unsqueeze(dim=-2)
#: Prepare t: (b, e, d) -> (b, e, 1, d, 1)
t_for_w = t.unsqueeze(dim=-2).unsqueeze(dim=-1)
#: Prepare w: (R, k, d, d) -> (b, k, d, d) -> (b, 1, k, d, d)
w_r = self.w.index_select(dim=0, index=r_indices).unsqueeze(dim=1)
# h.T @ W @ t, shape: (b, e, k, 1, 1)
hwt = (h_for_w @ w_r @ t_for_w)
#: reduce (b, e, k, 1, 1) -> (b, e, k)
hwt = hwt.squeeze(dim=-1).squeeze(dim=-1)
#: Prepare vh: (R, k, d) -> (b, k, d) -> (b, 1, k, d)
vh_r = self.vh.index_select(dim=0, index=r_indices).unsqueeze(dim=1)
#: Prepare h: (b, e, d) -> (b, e, d, 1)
h_for_v = h.unsqueeze(dim=-1)
# V_h @ h, shape: (b, e, k, 1)
vhh = vh_r @ h_for_v
#: reduce (b, e, k, 1) -> (b, e, k)
vhh = vhh.squeeze(dim=-1)
#: Prepare vt: (R, k, d) -> (b, k, d) -> (b, 1, k, d)
vt_r = self.vt.index_select(dim=0, index=r_indices).unsqueeze(dim=1)
#: Prepare t: (b, e, d) -> (b, e, d, 1)
t_for_v = t.unsqueeze(dim=-1)
# V_t @ t, shape: (b, e, k, 1)
vtt = vt_r @ t_for_v
#: reduce (b, e, k, 1) -> (b, e, k)
vtt = vtt.squeeze(dim=-1)
#: Prepare b: (R, k) -> (b, k) -> (b, 1, k)
b = self.b.index_select(dim=0, index=r_indices).unsqueeze(dim=1)
# a = f(h.T @ W @ t + Vh @ h + Vt @ t + b), shape: (b, e, k)
pre_act = hwt + vhh + vtt + b
act = self.non_linearity(pre_act)
# prepare u: (R, k) -> (b, k) -> (b, 1, k, 1)
u = self.u.index_select(dim=0, index=r_indices).unsqueeze(dim=1).unsqueeze(dim=-1)
# prepare act: (b, e, k) -> (b, e, 1, k)
act = act.unsqueeze(dim=-2)
# compute score, shape: (b, e, 1, 1)
score = act @ u
# reduce
score = score.squeeze(dim=-1).squeeze(dim=-1)
return score
def __init__(self, in_features, out_features):
super(linear, self).__init__()
self.weight = Parameter(FloatTensor(in_features, out_features))
self.register_parameter('bias', None)
stdv = 1. / sqrt(self.weight.size(1))
self.weight.data.uniform_(-stdv, stdv)
((idx + 1), len(train_table)))
print("[INFO] load train_x successfully, train_x length :", len(train_x))
# load model
import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = MFrnnVGG()
model.to(device)
model.eval()
model.load_train_pretrain()
# eval
from torch import FloatTensor
from torch.autograd import Variable
with open(os.path.join(output_dir, 'p2_result.txt'), 'w+') as file:
for idx, x in enumerate(train_x):
try:
x = np.transpose(x, (0, 3, 1, 2))
x = Variable(FloatTensor(x)).to(device)
pred = model(x).cpu().detach().numpy()
pred = pred.argmax().item()
file.write(str(pred) + '\n')
except Exception as e:
print(idx, e)
file.write('1\n')
# output
def forward(
self,
X: torch.FloatTensor,
edge_index: Union[torch.LongTensor, List[torch.LongTensor]],
) -> torch.FloatTensor:
"""
Making a forward pass with the ASTGCN block.
Arg types:
* **X** (PyTorch Float Tensor) - Node features for T time periods, with shape (B, N_nodes, F_in, T_in).
* **edge_index** (LongTensor): Edge indices, can be an array of a list of Tensor arrays, depending on whether edges change over time.
Return types:
* **X** (PyTorch Float Tensor) - Hidden state tensor for all nodes, with shape (B, N_nodes, nb_time_filter, T_out).
"""
batch_size, num_of_vertices, num_of_features, num_of_timesteps = X.shape
X_tilde = self._temporal_attention(X)
X_tilde = torch.matmul(X.reshape(batch_size, -1, num_of_timesteps), X_tilde)
X_tilde = X_tilde.reshape(
batch_size, num_of_vertices, num_of_features, num_of_timesteps
)
X_tilde = self._spatial_attention(X_tilde)
if not isinstance(edge_index, list):
data = Data(
edge_index=edge_index, edge_attr=None, num_nodes=num_of_vertices
)
if self._normalization != "sym":
lambda_max = LaplacianLambdaMax()(data).lambda_max
else:
lambda_max = None
X_hat = []
for t in range(num_of_timesteps):
X_hat.append(
torch.unsqueeze(
self._chebconv_attention(
X[:, :, :, t], edge_index, X_tilde, lambda_max=lambda_max
),
-1,
)
)
X_hat = F.relu(torch.cat(X_hat, dim=-1))
else:
X_hat = []
for t in range(num_of_timesteps):
data = Data(
edge_index=edge_index[t], edge_attr=None, num_nodes=num_of_vertices
)
if self._normalization != "sym":
lambda_max = LaplacianLambdaMax()(data).lambda_max
else:
lambda_max = None
X_hat.append(
torch.unsqueeze(
self._chebconv_attention(
X[:, :, :, t], edge_index[t], X_tilde, lambda_max=lambda_max
),
-1,
)
)
X_hat = F.relu(torch.cat(X_hat, dim=-1))
X_hat = self._time_convolution(X_hat.permute(0, 2, 1, 3))
X = self._residual_convolution(X.permute(0, 2, 1, 3))
X = self._layer_norm(F.relu(X + X_hat).permute(0, 3, 2, 1))
X = X.permute(0, 2, 3, 1)
return X
def compute_loss(self,
task_logits: torch.FloatTensor,
labels: torch.LongTensor,
task_ind,
task_num_labels,
batch_size,
seq_len,
state: dict,
) -> None:
"""
Computes the loss for a single batch and a single task.
Args:
task_logits:
labels:
task_ind:
task_num_labels:
batch_size:
seq_len:
state:
"""
if task_num_labels == 1:
# We are doing regression
loss_fct = MSELoss()
task_loss = loss_fct(task_logits.view(-1), labels.view(-1))
else:
if not self.class_weights[task_ind] is None:
class_weights = torch.FloatTensor(self.class_weights[task_ind]).to(self.device)
else:
class_weights = None
loss_fct = CrossEntropyLoss(weight=class_weights)
if self.relations[task_ind]:
task_labels = labels[:, 0, state['task_label_ind']:state['task_label_ind'] + seq_len, :]
state['task_label_ind'] += seq_len
task_loss = loss_fct(task_logits.permute(0, 3, 1, 2),
task_labels.type(torch.LongTensor).to(labels.device))
elif self.tagger[task_ind]:
# in cases where we are only given a single task the HF code will have one fewer dimension in the labels, so just add a dummy dimension to make our indexing work:
if labels.ndim == 3:
labels = labels.unsqueeze(1)
task_labels = labels[:, 0, state['task_label_ind'], :]
state['task_label_ind'] += 1
task_loss = loss_fct(task_logits.view(-1, task_num_labels),
task_labels.reshape([batch_size * seq_len, ]).type(torch.LongTensor).to(
labels.device))
else:
if labels.ndim == 2:
task_labels = labels[:, 0]
elif labels.ndim == 3:
task_labels = labels[:, 0, task_ind]
else:
raise NotImplementedError(
'Have not implemented the case where a classification task '
'is part of an MTL setup with relations and sequence tagging')
state['task_label_ind'] += 1
task_loss = loss_fct(task_logits, task_labels.type(torch.LongTensor).to(labels.device))
if state['loss'] is None:
state['loss'] = task_loss
else:
task_weight = 1.0 if task_ind + 1 < len(self.num_labels) else self.final_task_weight
state['loss'] += (task_weight * task_loss)
def forward(
self, # pylint: disable=arguments-differ
span_embeddings: torch.FloatTensor,
span_mask: torch.LongTensor,
num_spans_to_keep: int
) -> Tuple[torch.FloatTensor, torch.LongTensor, torch.LongTensor,
torch.FloatTensor]:
"""
Extracts the top-k scoring spans with respect to the scorer. We additionally return
the indices of the top-k in their original order, not ordered by score, so that we
can rely on the ordering to consider the previous k spans as antecedents for each
span later.
Parameters
----------
span_embeddings : ``torch.FloatTensor``, required.
A tensor of shape (batch_size, num_spans, embedding_size), representing
the set of embedded span representations.
span_mask : ``torch.LongTensor``, required.
A tensor of shape (batch_size, num_spans), denoting unpadded elements
of ``span_embeddings``.
num_spans_to_keep : ``int``, required.
The number of spans to keep when pruning.
Returns
-------
top_span_embeddings : ``torch.FloatTensor``
The span representations of the top-k scoring spans.
Has shape (batch_size, num_spans_to_keep, embedding_size).
top_span_mask : ``torch.LongTensor``
The coresponding mask for ``top_span_embeddings``.
Has shape (batch_size, num_spans_to_keep).
top_span_indices : ``torch.IntTensor``
The indices of the top-k scoring spans into the original ``span_embeddings``
tensor. This is returned because it can be useful to retain pointers to
the original spans, if each span is being scored by multiple distinct
scorers, for instance. Has shape (batch_size, num_spans_to_keep).
top_span_scores : ``torch.FloatTensor``
The values of the top-k scoring spans.
Has shape (batch_size, num_spans_to_keep, 1).
"""
span_mask = span_mask.unsqueeze(-1)
num_spans = span_embeddings.size(1)
# Shape: (batch_size, num_spans, 1)
span_scores = self._scorer(span_embeddings)
if span_scores.size(-1) != 1 or span_scores.dim() != 3:
raise ValueError(
f"The scorer passed to SpanPruner must produce a tensor of shape"
f"(batch_size, num_spans, 1), but found shape {span_scores.size()}"
)
# Make sure that we don't select any masked spans by
# setting their scores to be -inf.
span_scores += span_mask.log()
# Shape: (batch_size, num_spans_to_keep, 1)
_, top_span_indices = span_scores.topk(num_spans_to_keep, 1)
# Now we order the selected indices in increasing order with
# respect to their indices (and hence, with respect to the
# order they originally appeared in the ``span_embeddings`` tensor).
top_span_indices, _ = torch.sort(top_span_indices, 1)
# Shape: (batch_size, num_spans_to_keep)
top_span_indices = top_span_indices.squeeze(-1)
# Shape: (batch_size * num_spans_to_keep)
# torch.index_select only accepts 1D indices, but here
# we need to select spans for each element in the batch.
flat_top_span_indices = util.flatten_and_batch_shift_indices(
top_span_indices, num_spans)
# Shape: (batch_size, num_spans_to_keep, embedding_size)
top_span_embeddings = util.batched_index_select(
span_embeddings, top_span_indices, flat_top_span_indices)
# Shape: (batch_size, num_spans_to_keep)
top_span_mask = util.batched_index_select(span_mask, top_span_indices,
flat_top_span_indices)
# Shape: (batch_size, num_spans_to_keep, 1)
top_span_scores = util.batched_index_select(span_scores,
top_span_indices,
flat_top_span_indices)
return top_span_embeddings, top_span_mask.squeeze(
-1), top_span_indices, top_span_scores
def _viterbi_decode(self, emissions: torch.FloatTensor,
mask: torch.ByteTensor) -> List[List[int]]:
# emissions: (seq_length, batch_size, num_tags)
# mask: (seq_length, batch_size)
assert emissions.dim() == 3 and mask.dim() == 2
assert emissions.shape[:2] == mask.shape
assert emissions.size(2) == self.num_tags
assert mask[0].all()
seq_length, batch_size = mask.shape
# self.start_transitions start 到其他tag(不包含end)的得分
score = self.start_transitions + emissions[0]
history = []
# for i in range(1,seq_length):
#
# # shape : (batch_size,num_tag,1)
# broadcast_score = score.unsqueeze(dim=2)
#
# # shape: (batch_size,1,num_tags)
# broadcast_emissions = emissions[i].unsqueeze(1)
#
# next_score = broadcast_score + self.transitions + broadcast_emissions
#
# next_score = torch.logsumexp(next_score,dim = 1)
#
# score = torch.where(mask[i].unsqueeze(1),next_score,score)
for i in range(1, seq_length):
broadcast_score = score.unsqueeze(2)
broadcast_emission = emissions[i].unsqueeze(1)
next_score = broadcast_score + self.transitions + broadcast_emission
next_score, indices = next_score.max(dim=1)
score = torch.where(mask[i].unsqueeze(1), next_score, score)
history.append(indices)
score += self.end_transitions
seq_ends = mask.long().sum(dim=0) - 1
best_tags_list = []
for idx in range(batch_size):
_, best_last_tag = score[idx].max(dim=0)
best_tags = [best_last_tag.item()]
# history[:seq_ends[idx]].shape (seq_ends[idx])
for hist in reversed(history[:seq_ends[idx]]):
best_last_tag = hist[idx][best_tags[-1]]
best_tags.append(best_last_tag.item())
best_tags.reverse()
best_tags_list.append(best_tags)
return best_tags_list
def analyze_prob(
log_prob: torch.
FloatTensor, # SHAPE: (batch_size, num_temp, vocab_size)
label_index: torch.LongTensor, # SHAPE: (batch_size)
output: bool = False,
method: str = 'all'):
topk = 5
show_num = 5
vocab_size = log_prob.size(-1)
prob = log_prob.exp()
# SHAPE: (batch_size, num_temp, topk)
prob_top, prob_top_ind = torch.topk(input=prob, k=topk, dim=-1)
# SHAPE: (batch_size, num_temp)
correct_mask = prob_top_ind[:, :, 0].eq(label_index.view(-1, 1))
# SHAPE: (batch_size, num_temp)
prob_gap = prob_top[:, :, 0] - prob_top[:, :, 1]
prob_abs = prob_top[:, :, 0]
c_prob_top = prob_top.masked_select(correct_mask.unsqueeze(-1)).view(
-1, topk)
inc_prob_top = prob_top.masked_select(~correct_mask.unsqueeze(-1)).view(
-1, topk)
if method == 'all':
## overall statistics
# SHAPE: (None)
c_prob_gap = prob_gap.masked_select(correct_mask)
c_prob_abs = prob_abs.masked_select(correct_mask)
# SHAPE: (None, topk)
inc_prob_gap = prob_gap.masked_select(~correct_mask)
inc_prob_abs = prob_abs.masked_select(~correct_mask)
num_c, num_inc = c_prob_gap.size(0), inc_prob_gap.size(0)
c_gap = c_prob_gap.sum().item()
inc_gap = inc_prob_gap.sum().item()
c_abs = c_prob_abs.sum().item()
inc_abs = inc_prob_abs.sum().item()
elif method == 'sample':
## sample-wise statistics
correct_mask = correct_mask.float()
incorrect_mask = 1 - correct_mask
# only consider samples with both correct and incorrect templates
correct_mask = correct_mask * incorrect_mask.max(-1, keepdim=True)[0]
incorrect_mask = incorrect_mask * correct_mask.max(-1, keepdim=True)[0]
c_gap = (prob_gap * correct_mask).max(-1)[0].sum().item()
c_abs = (prob_abs * correct_mask).max(-1)[0].sum().item()
inc_gap = (prob_gap * incorrect_mask).max(-1)[0].sum().item()
inc_abs = (prob_abs * incorrect_mask).max(-1)[0].sum().item()
num_c = correct_mask.max(-1)[0].sum()
num_inc = incorrect_mask.max(-1)[0].sum()
else:
raise Exception
if output:
print('#correct temp {}, #incorrect temp {}'.format(num_c, num_inc))
print('correct')
print(c_prob_top[np.random.choice(num_c,
min(show_num, num_c),
replace=False)])
print('incorrect')
print(inc_prob_top[np.random.choice(num_inc,
min(show_num, num_inc),
replace=False)])
return (c_gap, c_abs, num_c), (inc_gap, inc_abs, num_inc)
def _get_linking_probabilities(
self, worlds: List[WikiTablesLanguage],
linking_scores: torch.FloatTensor, question_mask: torch.LongTensor,
entity_type_dict: Dict[int, int]) -> torch.FloatTensor:
"""
Produces the probability of an entity given a question word and type. The logic below
separates the entities by type since the softmax normalization term sums over entities
of a single type.
Parameters
----------
worlds : ``List[WikiTablesLanguage]``
linking_scores : ``torch.FloatTensor``
Has shape (batch_size, num_question_tokens, num_entities).
question_mask: ``torch.LongTensor``
Has shape (batch_size, num_question_tokens).
entity_type_dict : ``Dict[int, int]``
This is a mapping from ((batch_index * num_entities) + entity_index) to entity type id.
Returns
-------
batch_probabilities : ``torch.FloatTensor``
Has shape ``(batch_size, num_question_tokens, num_entities)``.
Contains all the probabilities for an entity given a question word.
"""
_, num_question_tokens, num_entities = linking_scores.size()
batch_probabilities = []
for batch_index, world in enumerate(worlds):
all_probabilities = []
num_entities_in_instance = 0
# NOTE: The way that we're doing this here relies on the fact that entities are
# implicitly sorted by their types when we sort them by name, and that numbers come
# before "date_column:", followed by "number_column:", "string:", and "string_column:".
# This is not a great assumption, and could easily break later, but it should work for now.
for type_index in range(self._num_entity_types):
# This index of 0 is for the null entity for each type, representing the case where a
# word doesn't link to any entity.
entity_indices = [0]
entities = world.table_graph.entities
for entity_index, _ in enumerate(entities):
if entity_type_dict[batch_index * num_entities +
entity_index] == type_index:
entity_indices.append(entity_index)
if len(entity_indices) == 1:
# No entities of this type; move along...
continue
# We're subtracting one here because of the null entity we added above.
num_entities_in_instance += len(entity_indices) - 1
# We separate the scores by type, since normalization is done per type. There's an
# extra "null" entity per type, also, so we have `num_entities_per_type + 1`. We're
# selecting from a (num_question_tokens, num_entities) linking tensor on _dimension 1_,
# so we get back something of shape (num_question_tokens,) for each index we're
# selecting. All of the selected indices together then make a tensor of shape
# (num_question_tokens, num_entities_per_type + 1).
indices = linking_scores.new_tensor(entity_indices,
dtype=torch.long)
entity_scores = linking_scores[batch_index].index_select(
1, indices)
# We used index 0 for the null entity, so this will actually have some values in it.
# But we want the null entity's score to be 0, so we set that here.
entity_scores[:, 0] = 0
# No need for a mask here, as this is done per batch instance, with no padding.
type_probabilities = torch.nn.functional.softmax(entity_scores,
dim=1)
all_probabilities.append(type_probabilities[:, 1:])
# We need to add padding here if we don't have the right number of entities.
if num_entities_in_instance != num_entities:
zeros = linking_scores.new_zeros(
num_question_tokens,
num_entities - num_entities_in_instance)
all_probabilities.append(zeros)
# (num_question_tokens, num_entities)
probabilities = torch.cat(all_probabilities, dim=1)
batch_probabilities.append(probabilities)
batch_probabilities = torch.stack(batch_probabilities, dim=0)
return batch_probabilities * question_mask.unsqueeze(-1).float()
def forward(inputs_x: torch.FloatTensor):
if inputs_x.size() == (1, 1, 1):
return None
else:
raise Exception('Mock')
def forward(self,
sequence_tensor: torch.FloatTensor,
span_indices: torch.LongTensor,
sequence_mask: torch.LongTensor = None,
span_indices_mask: torch.LongTensor = None) -> torch.FloatTensor:
# Both of shape (batch_size, sequence_length, embedding_size / 2)
forward_sequence, backward_sequence = sequence_tensor.split(int(self._input_dim / 2), dim=-1)
forward_sequence = forward_sequence.contiguous()
backward_sequence = backward_sequence.contiguous()
# shape (batch_size, num_spans)
span_starts, span_ends = [index.squeeze(-1) for index in span_indices.split(1, dim=-1)]
if span_indices_mask is not None:
span_starts = span_starts * span_indices_mask
span_ends = span_ends * span_indices_mask
# We want `exclusive` span starts, so we remove 1 from the forward span starts
# as the AllenNLP ``SpanField`` is inclusive.
# shape (batch_size, num_spans)
exclusive_span_starts = span_starts - 1
# shape (batch_size, num_spans, 1)
start_sentinel_mask = (exclusive_span_starts == -1).long().unsqueeze(-1)
# We want `exclusive` span ends for the backward direction
# (so that the `start` of the span in that direction is exlusive), so
# we add 1 to the span ends as the AllenNLP ``SpanField`` is inclusive.
exclusive_span_ends = span_ends + 1
if sequence_mask is not None:
# shape (batch_size)
sequence_lengths = util.get_lengths_from_binary_sequence_mask(sequence_mask)
else:
# shape (batch_size), filled with the sequence length size of the sequence_tensor.
sequence_lengths = util.ones_like(sequence_tensor[:, 0, 0]).long() * sequence_tensor.size(1)
# shape (batch_size, num_spans, 1)
end_sentinel_mask = (exclusive_span_ends == sequence_lengths.unsqueeze(-1)).long().unsqueeze(-1)
# As we added 1 to the span_ends to make them exclusive, which might have caused indices
# equal to the sequence_length to become out of bounds, we multiply by the inverse of the
# end_sentinel mask to erase these indices (as we will replace them anyway in the block below).
# The same argument follows for the exclusive span start indices.
exclusive_span_ends = exclusive_span_ends * (1 - end_sentinel_mask.squeeze(-1))
exclusive_span_starts = exclusive_span_starts * (1 - start_sentinel_mask.squeeze(-1))
# We'll check the indices here at runtime, because it's difficult to debug
# if this goes wrong and it's tricky to get right.
if (exclusive_span_starts < 0).any() or (exclusive_span_ends > sequence_lengths.unsqueeze(-1)).any():
raise ValueError(f"Adjusted span indices must lie inside the length of the sequence tensor, "
f"but found: exclusive_span_starts: {exclusive_span_starts}, "
f"exclusive_span_ends: {exclusive_span_ends} for a sequence tensor with lengths "
f"{sequence_lengths}.")
# Forward Direction: start indices are exclusive. Shape (batch_size, num_spans, input_size / 2)
forward_start_embeddings = util.batched_index_select(forward_sequence, exclusive_span_starts)
# Forward Direction: end indices are inclusive, so we can just use span_ends.
# Shape (batch_size, num_spans, input_size / 2)
forward_end_embeddings = util.batched_index_select(forward_sequence, span_ends)
# Backward Direction: The backward start embeddings use the `forward` end
# indices, because we are going backwards.
# Shape (batch_size, num_spans, input_size / 2)
backward_start_embeddings = util.batched_index_select(backward_sequence, exclusive_span_ends)
# Backward Direction: The backward end embeddings use the `forward` start
# indices, because we are going backwards.
# Shape (batch_size, num_spans, input_size / 2)
backward_end_embeddings = util.batched_index_select(backward_sequence, span_starts)
if self._use_sentinels:
# If we're using sentinels, we need to replace all the elements which were
# outside the dimensions of the sequence_tensor with either the start sentinel,
# or the end sentinel.
float_end_sentinel_mask = end_sentinel_mask.float()
float_start_sentinel_mask = start_sentinel_mask.float()
forward_start_embeddings = forward_start_embeddings * (1 - float_start_sentinel_mask) \
+ float_start_sentinel_mask * self._start_sentinel
backward_start_embeddings = backward_start_embeddings * (1 - float_end_sentinel_mask) \
+ float_end_sentinel_mask * self._end_sentinel
# Now we combine the forward and backward spans in the manner specified by the
# respective combinations and concatenate these representations.
# Shape (batch_size, num_spans, forward_combination_dim)
forward_spans = util.combine_tensors(self._forward_combination,
[forward_start_embeddings, forward_end_embeddings])
# Shape (batch_size, num_spans, backward_combination_dim)
backward_spans = util.combine_tensors(self._backward_combination,
[backward_start_embeddings, backward_end_embeddings])
# Shape (batch_size, num_spans, forward_combination_dim + backward_combination_dim)
span_embeddings = torch.cat([forward_spans, backward_spans], -1)
if self._span_width_embedding is not None:
# Embed the span widths and concatenate to the rest of the representations.
if self._bucket_widths:
span_widths = util.bucket_values(span_ends - span_starts,
num_total_buckets=self._num_width_embeddings)
else:
span_widths = span_ends - span_starts
span_width_embeddings = self._span_width_embedding(span_widths)
return torch.cat([span_embeddings, span_width_embeddings], -1)
if span_indices_mask is not None:
return span_embeddings * span_indices_mask.float().unsqueeze(-1)
return span_embeddings
def forward(self, x: torch.FloatTensor) -> torch.FloatTensor: # noqa: D102
value = x.norm(p=self.p, dim=self.dim).mean()
if not self.normalize:
return value
return value / _get_expected_norm(p=self.p, d=x.shape[-1])
def affines_mat(matrices=None):
if matrices is None: matrices = []
matrices = [FloatTensor(m) for m in matrices if m is not None]
return reduce(torch.matmul, matrices, torch.eye(3))
def forward(
self,
sequence_tensor: torch.FloatTensor,
span_indices: torch.LongTensor,
sequence_mask: torch.LongTensor = None,
span_indices_mask: torch.LongTensor = None) -> torch.FloatTensor:
# Both of shape (batch_size, sequence_length, embedding_size / 2)
forward_sequence, backward_sequence = sequence_tensor.split(int(
self._input_dim / 2),
dim=-1)
forward_sequence = forward_sequence.contiguous()
backward_sequence = backward_sequence.contiguous()
# shape (batch_size, num_spans)
span_starts, span_ends = [
index.squeeze(-1) for index in span_indices.split(1, dim=-1)
]
if span_indices_mask is not None:
span_starts = span_starts * span_indices_mask
span_ends = span_ends * span_indices_mask
# We want `exclusive` span starts, so we remove 1 from the forward span starts
# as the AllenNLP ``SpanField`` is inclusive.
# shape (batch_size, num_spans)
exclusive_span_starts = span_starts - 1
# shape (batch_size, num_spans, 1)
start_sentinel_mask = (
exclusive_span_starts == -1).long().unsqueeze(-1)
# We want `exclusive` span ends for the backward direction
# (so that the `start` of the span in that direction is exlusive), so
# we add 1 to the span ends as the AllenNLP ``SpanField`` is inclusive.
exclusive_span_ends = span_ends + 1
if sequence_mask is not None:
# shape (batch_size)
sequence_lengths = util.get_lengths_from_binary_sequence_mask(
sequence_mask)
else:
# shape (batch_size), filled with the sequence length size of the sequence_tensor.
sequence_lengths = (
torch.ones_like(sequence_tensor[:, 0, 0], dtype=torch.long) *
sequence_tensor.size(1))
# shape (batch_size, num_spans, 1)
end_sentinel_mask = (
exclusive_span_ends >=
sequence_lengths.unsqueeze(-1)).long().unsqueeze(-1)
# As we added 1 to the span_ends to make them exclusive, which might have caused indices
# equal to the sequence_length to become out of bounds, we multiply by the inverse of the
# end_sentinel mask to erase these indices (as we will replace them anyway in the block below).
# The same argument follows for the exclusive span start indices.
exclusive_span_ends = exclusive_span_ends * (
1 - end_sentinel_mask.squeeze(-1))
exclusive_span_starts = exclusive_span_starts * (
1 - start_sentinel_mask.squeeze(-1))
# We'll check the indices here at runtime, because it's difficult to debug
# if this goes wrong and it's tricky to get right.
if (exclusive_span_starts < 0).any() or (
exclusive_span_ends > sequence_lengths.unsqueeze(-1)).any():
raise ValueError(
f"Adjusted span indices must lie inside the length of the sequence tensor, "
f"but found: exclusive_span_starts: {exclusive_span_starts}, "
f"exclusive_span_ends: {exclusive_span_ends} for a sequence tensor with lengths "
f"{sequence_lengths}.")
# Forward Direction: start indices are exclusive. Shape (batch_size, num_spans, input_size / 2)
forward_start_embeddings = util.batched_index_select(
forward_sequence, exclusive_span_starts)
# Forward Direction: end indices are inclusive, so we can just use span_ends.
# Shape (batch_size, num_spans, input_size / 2)
forward_end_embeddings = util.batched_index_select(
forward_sequence, span_ends)
# Backward Direction: The backward start embeddings use the `forward` end
# indices, because we are going backwards.
# Shape (batch_size, num_spans, input_size / 2)
backward_start_embeddings = util.batched_index_select(
backward_sequence, exclusive_span_ends)
# Backward Direction: The backward end embeddings use the `forward` start
# indices, because we are going backwards.
# Shape (batch_size, num_spans, input_size / 2)
backward_end_embeddings = util.batched_index_select(
backward_sequence, span_starts)
if self._use_sentinels:
# If we're using sentinels, we need to replace all the elements which were
# outside the dimensions of the sequence_tensor with either the start sentinel,
# or the end sentinel.
float_end_sentinel_mask = end_sentinel_mask.float()
float_start_sentinel_mask = start_sentinel_mask.float()
forward_start_embeddings = forward_start_embeddings * (1 - float_start_sentinel_mask) \
+ float_start_sentinel_mask * self._start_sentinel
backward_start_embeddings = backward_start_embeddings * (1 - float_end_sentinel_mask) \
+ float_end_sentinel_mask * self._end_sentinel
# Now we combine the forward and backward spans in the manner specified by the
# respective combinations and concatenate these representations.
# Shape (batch_size, num_spans, forward_combination_dim)
forward_spans = util.combine_tensors(
self._forward_combination,
[forward_start_embeddings, forward_end_embeddings])
# Shape (batch_size, num_spans, backward_combination_dim)
backward_spans = util.combine_tensors(
self._backward_combination,
[backward_start_embeddings, backward_end_embeddings])
# Shape (batch_size, num_spans, forward_combination_dim + backward_combination_dim)
span_embeddings = torch.cat([forward_spans, backward_spans], -1)
if self._span_width_embedding is not None:
# Embed the span widths and concatenate to the rest of the representations.
if self._bucket_widths:
span_widths = util.bucket_values(
span_ends - span_starts,
num_total_buckets=self._num_width_embeddings)
else:
span_widths = span_ends - span_starts
span_width_embeddings = self._span_width_embedding(span_widths)
return torch.cat([span_embeddings, span_width_embeddings], -1)
if span_indices_mask is not None:
return span_embeddings * span_indices_mask.float().unsqueeze(-1)
return span_embeddings
def to_tensor(data):
return [FloatTensor(point) for point in data]
def forward(self,
sequence_tensor: torch.FloatTensor,
span_indices: torch.LongTensor,
sequence_mask: torch.LongTensor = None,
span_indices_mask: torch.LongTensor = None) -> torch.FloatTensor:
# both of shape (batch_size, num_spans, 1)
span_starts, span_ends = span_indices.split(1, dim=-1)
# shape (batch_size, num_spans, 1)
# These span widths are off by 1, because the span ends are `inclusive`.
span_widths = span_ends - span_starts
# We need to know the maximum span width so we can
# generate indices to extract the spans from the sequence tensor.
# These indices will then get masked below, such that if the length
# of a given span is smaller than the max, the rest of the values
# are masked.
max_batch_span_width = span_widths.max().item() + 1
# shape (batch_size, sequence_length, 1)
global_attention_logits = self._global_attention(sequence_tensor)
# Shape: (1, 1, max_batch_span_width)
max_span_range_indices = util.get_range_vector(max_batch_span_width,
util.get_device_of(sequence_tensor)).view(1, 1, -1)
# Shape: (batch_size, num_spans, max_batch_span_width)
# This is a broadcasted comparison - for each span we are considering,
# we are creating a range vector of size max_span_width, but masking values
# which are greater than the actual length of the span.
#
# We're using <= here (and for the mask below) because the span ends are
# inclusive, so we want to include indices which are equal to span_widths rather
# than using it as a non-inclusive upper bound.
span_mask = (max_span_range_indices <= span_widths).float()
raw_span_indices = span_ends - max_span_range_indices
# We also don't want to include span indices which are less than zero,
# which happens because some spans near the beginning of the sequence
# have an end index < max_batch_span_width, so we add this to the mask here.
span_mask = span_mask * (raw_span_indices >= 0).float()
span_indices = torch.nn.functional.relu(raw_span_indices.float()).long()
# Shape: (batch_size * num_spans * max_batch_span_width)
flat_span_indices = util.flatten_and_batch_shift_indices(span_indices, sequence_tensor.size(1))
# Shape: (batch_size, num_spans, max_batch_span_width, embedding_dim)
span_embeddings = util.batched_index_select(sequence_tensor, span_indices, flat_span_indices)
# Shape: (batch_size, num_spans, max_batch_span_width)
span_attention_logits = util.batched_index_select(global_attention_logits,
span_indices,
flat_span_indices).squeeze(-1)
# Shape: (batch_size, num_spans, max_batch_span_width)
span_attention_weights = util.masked_softmax(span_attention_logits, span_mask)
# Do a weighted sum of the embedded spans with
# respect to the normalised attention distributions.
# Shape: (batch_size, num_spans, embedding_dim)
attended_text_embeddings = util.weighted_sum(span_embeddings, span_attention_weights)
if span_indices_mask is not None:
# Above we were masking the widths of spans with respect to the max
# span width in the batch. Here we are masking the spans which were
# originally passed in as padding.
return attended_text_embeddings * span_indices_mask.unsqueeze(-1).float()
return attended_text_embeddings
def forward(
self,
x: torch.FloatTensor,
edge_index: torch.LongTensor,
spatial_attention: torch.FloatTensor,
edge_weight: OptTensor = None,
batch: OptTensor = None,
lambda_max: OptTensor = None,
) -> torch.FloatTensor:
"""
Making a forward pass of the ChebConv Attention layer.
Arg types:
* x (PyTorch Float Tensor) - Node features for T time periods, with shape (B, N_nodes, F_in).
* edge_index (Tensor array) - Edge indices.
* spatial_attention (PyTorch Float Tensor) - Spatial attention weights, with shape (B, N_nodes, N_nodes).
* edge_weight (PyTorch Float Tensor, optional) - Edge weights corresponding to edge indices.
* batch (PyTorch Tensor, optional) - Batch labels for each edge.
* lambda_max (optional, but mandatory if normalization is None) - Largest eigenvalue of Laplacian.
Return types:
* out (PyTorch Float Tensor) - Hidden state tensor for all nodes, with shape (B, N_nodes, F_out).
"""
if self._normalization != "sym" and lambda_max is None:
raise ValueError(
"You need to pass `lambda_max` to `forward() in`"
"case the normalization is non-symmetric."
)
if lambda_max is None:
lambda_max = torch.tensor(2.0, dtype=x.dtype, device=x.device)
if not isinstance(lambda_max, torch.Tensor):
lambda_max = torch.tensor(lambda_max, dtype=x.dtype, device=x.device)
assert lambda_max is not None
edge_index, norm = self.__norm__(
edge_index,
x.size(self.node_dim),
edge_weight,
self._normalization,
lambda_max,
dtype=x.dtype,
batch=batch,
)
row, col = edge_index
Att_norm = norm * spatial_attention[:, row, col]
num_nodes = x.size(self.node_dim)
TAx_0 = torch.matmul(
(torch.eye(num_nodes).to(edge_index.device) * spatial_attention).permute(
0, 2, 1
),
x,
)
out = torch.matmul(TAx_0, self._weight[0])
edge_index_transpose = edge_index[[1, 0]]
if self._weight.size(0) > 1:
TAx_1 = self.propagate(
edge_index_transpose, x=TAx_0, norm=Att_norm, size=None
)
out = out + torch.matmul(TAx_1, self._weight[1])
for k in range(2, self._weight.size(0)):
TAx_2 = self.propagate(edge_index_transpose, x=TAx_1, norm=norm, size=None)
TAx_2 = 2.0 * TAx_2 - TAx_0
out = out + torch.matmul(TAx_2, self._weight[k])
TAx_0, TAx_1 = TAx_1, TAx_2
if self._bias is not None:
out += self._bias
return out
def __getitem__(self, index):
if (self.return_label):
return (FloatTensor(self.data[index]), self.label[index])
else:
return FloatTensor(self.data[index])
def __init__(self):
super(DistanceAdj, self).__init__()
self.sigma = Parameter(FloatTensor(1))
self.sigma.data.fill_(0.1)
def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
top_k = min(max(self.top_k, self.min_tokens_to_keep), scores.size(-1)) # Safety check
# Remove all tokens with a probability less than the last token of the top-k
indices_to_remove = scores < torch.topk(scores, top_k)[0][..., -1, None]
scores = scores.masked_fill(indices_to_remove, self.filter_value)
return scores
def _viterbi_decode(self, emissions: torch.FloatTensor,
mask: torch.FloatTensor) -> torch.Tensor:
seq_len = emissions.shape[1]
mask = mask.to(torch.uint8)
log_prob = emissions[:, 0].clone()
log_prob += self.transitions[
self.start_tag, :self.start_tag].unsqueeze(0)
# At each step, we need to keep track of the total score, as if this step
# was the last valid step.
end_scores = log_prob + self.transitions[:self.start_tag,
self.end_tag].unsqueeze(0)
best_scores_list = []
# If the element has only token, empty tensor in best_paths helps
# torch.cat() from crashing
best_paths_list = [GetTensor(torch.Tensor().long())]
best_scores_list.append(end_scores.unsqueeze(1))
for idx in range(1, seq_len):
broadcast_emissions = emissions[:, idx].unsqueeze(1)
broadcast_transmissions = self.transitions[:self.start_tag, :self.
start_tag].unsqueeze(0)
broadcast_log_prob = log_prob.unsqueeze(2)
score = broadcast_emissions + broadcast_transmissions + broadcast_log_prob
max_scores, max_score_indices = torch.max(score, 1)
best_paths_list.append(max_score_indices.unsqueeze(1))
# Storing the scores incase this was the last step.
end_scores = max_scores + self.transitions[:self.start_tag, self.
end_tag].unsqueeze(0)
best_scores_list.append(end_scores.unsqueeze(1))
log_prob = max_scores
best_scores = torch.cat(best_scores_list, 1).float()
best_paths = torch.cat(best_paths_list, 1)
_, max_indices_from_scores = torch.max(best_scores, 2)
valid_index_tensor = GetTensor(torch.tensor(0)).long()
if self.ignore_index == Padding.DEFAULT_LABEL_PAD_IDX:
# No label for padding, so use 0 index.
padding_tensor = valid_index_tensor
else:
padding_tensor = GetTensor(torch.tensor(self.ignore_index)).long()
# Label for the last position is always based on the index with max score
# For illegal timesteps, we set as ignore_index
labels = max_indices_from_scores[:, seq_len - 1]
labels = self._mask_tensor(labels, 1 - mask[:, seq_len - 1],
padding_tensor)
all_labels = labels.unsqueeze(1).long()
# For Viterbi decoding, we start at the last position and go towards first
for idx in range(seq_len - 2, -1, -1):
# There are two ways to obtain labels for tokens at a particular position.
# Option 1: Use the labels obtained from the previous position to index
# the path in present position. This is used for all positions except
# last position in the sequence.
# Option 2: Find the indices with maximum scores obtained during
# viterbi decoding. This is used for the token at the last position
# For option 1 need to convert invalid indices to 0 so that lookups
# dont fail.
indices_for_lookup = all_labels[:, -1].clone()
indices_for_lookup = self._mask_tensor(
indices_for_lookup,
indices_for_lookup == self.ignore_index,
valid_index_tensor,
)
# Option 1 is used here when previous timestep (idx+1) was valid.
indices_from_prev_pos = (best_paths[:, idx, :].gather(
1,
indices_for_lookup.view(-1, 1).long()).squeeze(1))
indices_from_prev_pos = self._mask_tensor(indices_from_prev_pos,
(1 - mask[:, idx + 1]),
padding_tensor)
# Option 2 is used when last timestep was not valid which means idx+1
# is the last position in the sequence.
indices_from_max_scores = max_indices_from_scores[:, idx]
indices_from_max_scores = self._mask_tensor(
indices_from_max_scores, mask[:, idx + 1], padding_tensor)
# We need to combine results from 1 and 2 as rows in a batch can have
# sequences of varying lengths
labels = torch.where(
indices_from_max_scores == self.ignore_index,
indices_from_prev_pos,
indices_from_max_scores,
)
# Set to ignore_index if present state is not valid.
labels = self._mask_tensor(labels, (1 - mask[:, idx]),
padding_tensor)
all_labels = torch.cat((all_labels, labels.view(-1, 1).long()), 1)
return torch.flip(all_labels, [1])
def forward(self, # pylint: disable=arguments-differ
embeddings: torch.FloatTensor,
mask: torch.LongTensor,
num_items_to_keep: int) -> Tuple[torch.FloatTensor, torch.LongTensor,
torch.LongTensor, torch.FloatTensor]:
"""
Extracts the top-k scoring items with respect to the scorer. We additionally return
the indices of the top-k in their original order, not ordered by score, so that downstream
components can rely on the original ordering (e.g., for knowing what spans are valid
antecedents in a coreference resolution model).
Parameters
----------
embeddings : ``torch.FloatTensor``, required.
A tensor of shape (batch_size, num_items, embedding_size), containing an embedding for
each item in the list that we want to prune.
mask : ``torch.LongTensor``, required.
A tensor of shape (batch_size, num_items), denoting unpadded elements of
``embeddings``.
num_items_to_keep : ``int``, required.
The number of items to keep when pruning.
Returns
-------
top_embeddings : ``torch.FloatTensor``
The representations of the top-k scoring items.
Has shape (batch_size, num_items_to_keep, embedding_size).
top_mask : ``torch.LongTensor``
The corresponding mask for ``top_embeddings``.
Has shape (batch_size, num_items_to_keep).
top_indices : ``torch.IntTensor``
The indices of the top-k scoring items into the original ``embeddings``
tensor. This is returned because it can be useful to retain pointers to
the original items, if each item is being scored by multiple distinct
scorers, for instance. Has shape (batch_size, num_items_to_keep).
top_item_scores : ``torch.FloatTensor``
The values of the top-k scoring items.
Has shape (batch_size, num_items_to_keep, 1).
"""
mask = mask.unsqueeze(-1)
num_items = embeddings.size(1)
# Shape: (batch_size, num_items, 1)
scores = self._scorer(embeddings)
if scores.size(-1) != 1 or scores.dim() != 3:
raise ValueError(f"The scorer passed to Pruner must produce a tensor of shape"
f"(batch_size, num_items, 1), but found shape {scores.size()}")
# Make sure that we don't select any masked items by setting their scores to be very
# negative. These are logits, typically, so -1e20 should be plenty negative.
scores = util.replace_masked_values(scores, mask, -1e20)
# Shape: (batch_size, num_items_to_keep, 1)
_, top_indices = scores.topk(num_items_to_keep, 1)
# Now we order the selected indices in increasing order with
# respect to their indices (and hence, with respect to the
# order they originally appeared in the ``embeddings`` tensor).
top_indices, _ = torch.sort(top_indices, 1)
# Shape: (batch_size, num_items_to_keep)
top_indices = top_indices.squeeze(-1)
# Shape: (batch_size * num_items_to_keep)
# torch.index_select only accepts 1D indices, but here
# we need to select items for each element in the batch.
flat_top_indices = util.flatten_and_batch_shift_indices(top_indices, num_items)
# Shape: (batch_size, num_items_to_keep, embedding_size)
top_embeddings = util.batched_index_select(embeddings, top_indices, flat_top_indices)
# Shape: (batch_size, num_items_to_keep)
top_mask = util.batched_index_select(mask, top_indices, flat_top_indices)
# Shape: (batch_size, num_items_to_keep, 1)
top_scores = util.batched_index_select(scores, top_indices, flat_top_indices)
return top_embeddings, top_mask.squeeze(-1), top_indices, top_scores
def batch_shuffle_single_gpu(cls, x: torch.FloatTensor):
idx_shuffle = torch.randperm(x.size(0)).to(x.device)
idx_unshuffle = torch.argsort(idx_shuffle)
return x[idx_shuffle], idx_unshuffle
