def forward(
self,
contextualized_input_batch: Tensor,
stacks: Tensor,
buffers: Tensor,
stack_lengths: Tensor,
buffer_lengths: Tensor,
padding: Optional[Tensor] = None,
sentence_features: Optional[torch.Tensor] = None,
sentence_ids: Optional[List[str]] = None,
**kwargs,
) -> Tensor:
stack_batch = lookup_tensors_for_indices(stacks,
contextualized_input_batch)
buffer_batch = lookup_tensors_for_indices(buffers,
contextualized_input_batch)
stack_batch = self.positional_encoder(stack_batch)
buffer_batch = self.positional_encoder(buffer_batch)
# Compute a representation of the stack / buffer as an weighted average based
# on the attention weights.
stack_batch_attention, _, stack_attention_energies = self.stack_attention(
keys=stack_batch, sequence_lengths=stack_lengths)
buffer_batch_attention, _, buffer_attention_energies = self.buffer_attention(
keys=buffer_batch, sequence_lengths=buffer_lengths)
if self.reporter:
self.reporter.log(
"buffer",
buffers,
buffer_lengths,
buffer_attention_energies,
sentence_features,
sentence_ids,
)
self.reporter.log(
"stack",
stacks,
stack_lengths,
stack_attention_energies,
sentence_features,
sentence_ids,
)
return torch.cat((stack_batch_attention, buffer_batch_attention),
dim=1)
def forward(
self,
contextualized_input_batch: Tensor,
stacks: Tensor,
buffers: Tensor,
stack_lengths: Tensor,
buffer_lengths: Tensor,
finished_tokens: Tensor,
finished_tokens_lengths: Tensor,
sentence_lengths: Tensor,
padding: Optional[Tensor] = None,
sentence_features: Optional[torch.Tensor] = None,
sentence_ids: Optional[List[str]] = None,
**kwargs,
) -> Tensor:
# Look-up the whole unpadded buffer and stack sequence
stack_keys = lookup_tensors_for_indices(stacks, contextualized_input_batch)
buffer_keys = lookup_tensors_for_indices(buffers, contextualized_input_batch)
finished_tokens_keys = lookup_tensors_for_indices(
finished_tokens, contextualized_input_batch
)
# Add a padding vector so that even empty sequences have at least one item
stack_keys = self._fix_empty_sequence(stack_keys)
buffer_keys = self._fix_empty_sequence(buffer_keys)
finished_tokens_keys = self._fix_empty_sequence(finished_tokens_keys)
# Add positional encoding
stack_keys = self.positional_encoder(stack_keys)
buffer_keys = self.positional_encoder(buffer_keys)
finished_tokens_keys = self.positional_encoder(finished_tokens_keys)
sentence_keys = self.positional_encoder(contextualized_input_batch)
# Run universal attention over the finished tokens
tokens_attention, _, tokens_attention_energies = self.finished_tokens_attention(
keys=finished_tokens_keys, sequence_lengths=finished_tokens_lengths
)
sentence_attention, _, sentence_attention_energies = self.sentence_attention(
queries=tokens_attention, keys=sentence_keys, sequence_lengths=sentence_lengths
)
# Compute a representation of the stack / buffer as an weighted average based
# on the attention weights.
stack_batch_attention, _, stack_attention_energies = self.stack_attention(
queries=sentence_attention, keys=stack_keys, sequence_lengths=stack_lengths
)
buffer_batch_attention, _, buffer_attention_energies = self.buffer_attention(
queries=sentence_attention, keys=buffer_keys, sequence_lengths=buffer_lengths
)
if self.reporter:
sentence_tokens = torch.arange(
contextualized_input_batch.size(1), device=self.device
).expand(contextualized_input_batch.size(0), contextualized_input_batch.size(1))
self.reporter.log(
"finished_tokens",
finished_tokens,
finished_tokens_lengths,
tokens_attention_energies,
sentence_features,
sentence_ids,
)
self.reporter.log(
"sentence",
sentence_tokens,
sentence_lengths,
sentence_attention_energies,
sentence_features,
sentence_ids,
)
self.reporter.log(
"stack",
stacks,
stack_lengths,
stack_attention_energies,
sentence_features,
sentence_ids,
)
self.reporter.log(
"buffer",
buffers,
buffer_lengths,
buffer_attention_energies,
sentence_features,
sentence_ids,
)
return torch.cat((stack_batch_attention, buffer_batch_attention), dim=1)
def forward(
self,
contextualized_input_batch: Tensor,
stacks: Tensor,
buffers: Tensor,
stack_lengths: Tensor,
buffer_lengths: Tensor,
padding: Optional[Tensor] = None,
sentence_features: Optional[torch.Tensor] = None,
sentence_ids: Optional[List[str]] = None,
**kwargs,
) -> Tensor:
# Look-up the whole unpadded buffer and stack sequence
stack_keys = lookup_tensors_for_indices(stacks, contextualized_input_batch)
buffer_keys = lookup_tensors_for_indices(buffers, contextualized_input_batch)
stack_keys = self._fix_empty_sequence(stack_keys)
buffer_keys = self._fix_empty_sequence(buffer_keys)
# Add positional encoding
stack_keys = self.positional_encoder(stack_keys)
buffer_keys = self.positional_encoder(buffer_keys)
# Take the first entry as query
stack_queries = buffer_keys.index_select(
dim=1, index=torch.zeros(1, dtype=torch.int64, device=self.device)
).squeeze(1)
buffer_queries = stack_keys.index_select(
dim=1, index=torch.zeros(1, dtype=torch.int64, device=self.device)
).squeeze(1)
# Compute a representation of the stack / buffer as an weighted average based
# on the attention weights.
stack_batch_attention, _, stack_attention_energies = self.stack_attention(
queries=stack_queries, keys=stack_keys, sequence_lengths=stack_lengths
)
buffer_batch_attention, _, buffer_attention_energies = self.buffer_attention(
queries=buffer_queries, keys=buffer_keys, sequence_lengths=buffer_lengths
)
if self.reporter:
self.reporter.log(
"stack",
stacks,
stack_lengths,
stack_attention_energies,
sentence_features,
sentence_ids,
)
self.reporter.log(
"buffer",
buffers,
buffer_lengths,
buffer_attention_energies,
sentence_features,
sentence_ids,
)
return torch.cat((stack_batch_attention, buffer_batch_attention), dim=1)
