def compute_topk_scores_and_seq(sequences, scores, scores_to_gather, flags,
beam_dim, prefix="default"):
"""Given sequences and scores, will gather the top k=beam size sequences.
This function is used to grow alive, and finished. It takes sequences,
scores, and flags, and returns the top k from sequences, scores_to_gather,
and flags based on the values in scores.
This method permits easy introspection using tfdbg. It adds two named ops
that are prefixed by `prefix`:
- _topk_seq: the tensor for topk_seq returned by this method.
- _topk_flags: the tensor for topk_finished_flags returned by this method.
Args:
sequences: Tensor of sequences that we need to gather from.
[batch_size, beam_size, seq_length]
scores: Tensor of scores for each sequence in sequences.
[batch_size, beam_size]. We will use these to compute the topk.
scores_to_gather: Tensor of scores for each sequence in sequences.
[batch_size, beam_size]. We will return the gathered scores from here.
Scores to gather is different from scores because for grow_alive, we will
need to return log_probs, while for grow_finished, we will need to return
the length penalized scores.
flags: Tensor of bools for sequences that say whether a sequence has reached
EOS or not
beam_dim: mtf.Dimension
prefix: an optional string
Returns:
Tuple of
(topk_seq [batch_size, beam_size, decode_length],
topk_gathered_scores [batch_size, beam_size],
topk_finished_flags[batch_size, beam_size],
selector)
"""
unused_batch_dim, old_beam_dim, unused_length_dim = sequences.shape.dims
topk_indices, _ = mtf.top_k(scores, old_beam_dim, beam_dim)
selector = mtf.one_hot(topk_indices, old_beam_dim, dtype=tf.float32)
# Gather up the highest scoring sequences.
# For each operation added, give it
# a concrete name to simplify observing these operations with tfdbg.
# Clients can capture these tensors by watching these node names.
def gather(tensor, name):
with tf.name_scope(prefix + name):
output_shape = mtf.Shape(
[beam_dim if d == old_beam_dim else d for d in tensor.shape.dims])
return mtf.gather(
tensor, topk_indices, old_beam_dim, output_shape=output_shape)
topk_seq = gather(sequences, "_seq")
topk_flags = gather(flags, "_flags")
topk_gathered_scores = gather(scores_to_gather, "_scores")
return topk_seq, topk_gathered_scores, topk_flags, selector
def grow_topk(i, alive_seq, alive_log_probs, states=None):
r"""Inner beam search loop.
This function takes the current alive sequences, and grows them to topk
sequences where k = 2*beam. We use 2*beam because, we could have beam_size
number of sequences that might hit <EOS> and there will be no alive
sequences to continue. With 2*beam_size, this will not happen. This relies
on the assumption the vocab size is > beam size. If this is true, we'll
have at least beam_size non <EOS> extensions if we extract the next top
2*beam words.
Length penalty is given by = (5+len(decode)/6) ^ -\alpha. Pls refer to
https://arxiv.org/abs/1609.08144.
Args:
i: loop index
alive_seq: Topk sequences decoded so far [batch, beam, length]
alive_log_probs: probabilities of these sequences. [batch, beam]
states: optional list of mtf.Tensor
Returns:
Tuple of
(Topk sequences extended by the next word,
The log probs of these sequences,
The scores with length penalty of these sequences,
Flags indicating which of these sequences have finished decoding,
list of transformed decoding states)
"""
logits, new_states = logits_fn(i, alive_seq, states)
batch_dim, beam_dim, vocab_dim = logits.shape.dims
# Convert logits to normalized log probs
candidate_log_probs = mtf.log_softmax(logits, vocab_dim)
# Multiply the probabilities by the current probabilities of the beam.
# (batch_size, beam_size, vocab_size) + (batch_size, beam_size, 1)
log_probs = candidate_log_probs + alive_log_probs
length_penalty = mtf.pow(((5. + mtf.cast(i + 1, logits.dtype)) / 6.),
alpha)
# scores have shape [batch, beam, vocab]
curr_scores = log_probs / length_penalty
# We find the top 2k sequences to make sure we get k alive sequences.
#
# TODO(noam): This is inefficient. We should separately compute the k
# finished sequences (previously alive sequences + EOS), and the top k new
# alive sequences.
double_beam = mtf.Dimension("double_beam", beam_dim.size * 2)
if use_tpu and layout is not None and mesh_shape is not None:
# Do some partial top-k-ing first locally to avoid communication.
# We reshape the logits from:
# [batch, beam, vocab] to
# [batch, beam, major_vocab, minor_vocab]
# We first reduce (locally) across the minor_vocab dimension. This makes
# the thing we need to broadcast smaller.
# This also enables our shortcut of only picking the top num_prefilter
# sequences per beam per major_vocab in the first pass.
major_vocab_size = mtf.tensor_dim_to_mesh_dim_size(
layout, mesh_shape, vocab_dim)
major_vocab = mtf.Dimension(vocab_dim.name, major_vocab_size)
minor_vocab = mtf.Dimension("minor_vocab",
vocab_dim.size // major_vocab_size)
curr_scores = mtf.reshape(
curr_scores, [batch_dim, beam_dim, major_vocab, minor_vocab])
prefilter = mtf.Dimension("prefilter", num_prefilter
or double_beam.size)
# shape = [batch_dim, beam_dim, major_vocab, prefilter]
top_scores, top_minor_vocab_ids = mtf.top_k(
curr_scores, reduced_dim=minor_vocab, k_dim=prefilter)
combined = mtf.Dimension(
"combined", beam_dim.size * major_vocab.size * prefilter.size)
top_scores = mtf.reshape(top_scores, [batch_dim, combined])
top_minor_vocab_ids = mtf.reshape(top_minor_vocab_ids,
[batch_dim, combined])
# shpae = [batch_dim, double_beam]
# ids are indices representing (beam, major_vocab, prefilter)
top_scores, top_combined_ids = mtf.top_k(top_scores,
reduced_dim=combined,
k_dim=double_beam)
top_minor_vocab_ids = mtf.gather(
top_minor_vocab_ids,
top_combined_ids,
combined,
output_shape=[batch_dim, double_beam])
top_beam_index = top_combined_ids // (major_vocab.size *
prefilter.size)
top_combined_ids -= top_beam_index * (major_vocab.size *
prefilter.size)
top_major_vocab_ids = top_combined_ids // prefilter.size
top_combined_ids -= top_major_vocab_ids * prefilter.size
top_ids = top_major_vocab_ids * minor_vocab.size + top_minor_vocab_ids
else:
beam_and_vocab_dim = mtf.Dimension("beam_and_vocab",
beam_dim.size * vocab_dim.size)
flat_shape = mtf.Shape([batch_dim, beam_and_vocab_dim])
# Flatten out (beam_size, vocab_size) probs into a list of possibilities
flat_curr_scores = mtf.reshape(curr_scores,
flat_shape,
name="flatten_scores")
top_scores, top_ids = mtf.top_k(flat_curr_scores,
reduced_dim=beam_and_vocab_dim,
k_dim=double_beam)
# Work out what beam the top probs are in.
top_beam_index = top_ids // vocab_dim.size
top_ids %= vocab_dim.size # Unflatten the ids
# Recovering the log probs because we will need to send them back
top_log_probs = top_scores * length_penalty
selector = mtf.one_hot(top_beam_index, beam_dim, dtype=tf.float32)
def my_gather(tensor):
return mtf.gather(tensor,
top_beam_index,
beam_dim,
output_shape=mtf.Shape([
double_beam if d == beam_dim else d
for d in tensor.shape.dims
]))
# Gather up the most probable 2*beams both for the ids and finished_in_alive
# bools
top_seq = my_gather(alive_seq)
# Append the most probable alive
top_seq += top_ids * mtf.one_hot(i, length_dim, dtype=tf.int32)
top_finished = mtf.equal(top_ids, eos_id)
return (top_seq, top_log_probs, top_scores, top_finished, new_states,
selector)
