def _joint_likelihood(
self, logits: torch.Tensor, tags: torch.Tensor, mask: torch.BoolTensor
) -> torch.Tensor:
"""
Computes the numerator term for the log-likelihood, which is just score(inputs, tags)
"""
batch_size, sequence_length, _ = logits.data.shape
# Transpose batch size and sequence dimensions:
logits = logits.transpose(0, 1).contiguous()
mask = mask.transpose(0, 1).contiguous()
tags = tags.transpose(0, 1).contiguous()
start_index = 1 if (~mask[0]).all() else 0 # to skip [CLS] used in BERT
# Start with the transition scores from start_tag to the first tag in each input
if self.include_start_end_transitions:
score = self.start_transitions.index_select(0, tags[start_index] * mask[start_index])
else:
score = 0.0
# Add up the scores for the observed transitions and all the inputs but the last
current_valid_tag = tags[start_index] * mask[start_index]
for i in range(start_index, sequence_length-1):
# Each is shape (batch_size,)
current_tag, next_tag = tags[i], tags[i + 1]
if i > start_index: # preserve the tag for subwords
current_valid_tag = torch.where(mask[i] > 0, current_tag, current_valid_tag)
# The scores for transitioning from current_tag to next_tag
transition_score = self.transitions[current_valid_tag.view(-1), next_tag.view(-1) * mask[i + 1]]
# The score for using current_tag
emit_score = logits[i].gather(1, (current_tag * mask[i]).view(batch_size, 1)).squeeze(1)
# Include transition score if next element is unmasked,
# input_score if this element is unmasked.
score = score + transition_score * mask[i + 1] + emit_score * mask[i]
# Transition from last state to "stop" state. To start with, we need to find the last tag
# for each instance.
# last_tag_index = mask.sum(0).long() - 1
# last_tags = tags.gather(0, last_tag_index.view(1, batch_size)).squeeze(0)
last_tags = current_valid_tag
# Compute score of transitioning to `stop_tag` from each "last tag".
if self.include_start_end_transitions:
last_transition_score = self.end_transitions.index_select(0, last_tags)
else:
last_transition_score = 0.0
# Add the last input if it's not masked.
# last_inputs = logits[-1] # (batch_size, num_tags)
# last_input_score = last_inputs.gather(1, last_tags.view(-1, 1)) # (batch_size, 1)
# last_input_score = last_input_score.squeeze() # (batch_size,)
score = score + last_transition_score # + last_input_score * mask[-1]
return score
def _joint_likelihood(self, logits: torch.Tensor, tags: torch.Tensor,
mask: torch.BoolTensor) -> torch.Tensor:
"""
Computes the numerator term for the log-likelihood, which is just score(inputs, tags)
"""
batch_size, sequence_length, _ = logits.shape
# Transpose batch size and sequence dimensions:
logits = logits.transpose(0, 1)
mask = mask.transpose(0, 1).float()
tags = tags.transpose(0, 1)
# Start with the transition scores from start_tag to the first tag in each input
if self.include_start_end_transitions:
score = self.start_transitions.index_select(0, tags[0])
else:
score = 0.0
# Add up the scores for the observed transitions and all the inputs but the last
for i in range(sequence_length - 1):
# Each is shape (batch_size,)
current_tag, next_tag = tags[i], tags[i + 1]
# The scores for transitioning from current_tag to next_tag
transition_score = self.transitions[current_tag, next_tag]
# The score for using current_tag
emit_score = logits[i].gather(1, current_tag.view(batch_size,
1)).squeeze()
# Include t
score = score + transition_score * mask[i +
1] + emit_score * mask[i]
# Transition from last state to "stop" state. To start with, we need to find the last tag
# for each instance.
last_tag_index = mask.sum(0).long() - 1
last_tags = tags.gather(0, last_tag_index.view(1,
batch_size)).squeeze(0)
# Compute score of transitioning to `stop_tag` from each "last tag".
if self.include_start_end_transitions:
last_transition_score = self.end_transitions.index_select(
0, last_tags)
else:
last_transition_score = 0.0
# Add the last input if it's not masked.
last_inputs = logits[-1] # (batch_size, num_tags)
last_input_score = last_inputs.gather(1, last_tags.view(
-1, 1)).squeeze() # (batch_size)
# last_input_score = last_input_score.squeeze() # (batch_size,)
score = score + last_transition_score + last_input_score * mask[-1]
return score
def _input_likelihood(self, logits: torch.Tensor, mask: torch.BoolTensor) -> torch.Tensor:
"""
Computes the (batch_size,) denominator term for the log-likelihood, which is the
sum of the likelihoods across all possible state sequences.
"""
batch_size, sequence_length, num_tags = logits.size()
# Transpose batch size and sequence dimensions
mask = mask.transpose(0, 1).contiguous()
logits = logits.transpose(0, 1).contiguous()
start_index = 1 if (~mask[0]).all() else 0 # to skip [CLS] used in BERT
# Initial alpha is the (batch_size, num_tags) tensor of likelihoods combining the
# transitions to the initial states and the logits for the first timestep.
if self.include_start_end_transitions:
alpha = self.start_transitions.view(1, num_tags) + logits[start_index]
else:
alpha = logits[start_index]
# For each i we compute logits for the transitions from timestep i-1 to timestep i.
# We do so in a (batch_size, num_tags, num_tags) tensor where the axes are
# (instance, current_tag, next_tag)
for i in range(start_index, sequence_length):
# The emit scores are for time i ("next_tag") so we broadcast along the current_tag axis.
emit_scores = logits[i].view(batch_size, 1, num_tags)
# Transition scores are (current_tag, next_tag) so we broadcast along the instance axis.
transition_scores = self.transitions.view(1, num_tags, num_tags)
# Alpha is for the current_tag, so we broadcast along the next_tag axis.
broadcast_alpha = alpha.view(batch_size, num_tags, 1)
# Add all the scores together and logexp over the current_tag axis.
inner = broadcast_alpha + emit_scores + transition_scores
# In valid positions (mask == True) we want to take the logsumexp over the current_tag dimension
# of `inner`. Otherwise (mask == False) we want to retain the previous alpha.
alpha = logsumexp(inner, 1) * mask[i].view(batch_size, 1) + alpha * (
~mask[i]
).view(batch_size, 1)
# Every sequence needs to end with a transition to the stop_tag.
if self.include_start_end_transitions:
stops = alpha + self.end_transitions.view(1, num_tags)
else:
stops = alpha
# Finally we log_sum_exp along the num_tags dim, result is (batch_size,)
return logsumexp(stops)
