def _is_finished(i, unused_alive_seq, alive_log_probs, unused_finished_seq, finished_scores, finished_in_finished, *unused_states): """Checking termination condition. We terminate when we decoded up to decode_length or the lowest scoring item in finished has a greater score that the highest prob item in alive divided by the max length penalty Args: i: loop index alive_log_probs: probabilities of the beams. [batch_size, beam_size] finished_scores: scores for each of these sequences. [batch_size, beam_size] finished_in_finished: finished bools for each of these sequences. [batch_size, beam_size] Returns: Bool. """ # TODO(noam): support a different decode length... # decode_length = mtf.constant(mesh, length_dim.size, dtype=tf.int32) # del alive_log_probs, finished_scores, finished_in_finished # return mtf.less(i, length_dim.size) if not stop_early: return mtf.less(i, decode_length) max_length_penalty = mtf.pow( ((5. + mtf.cast(decode_length, finished_scores.dtype)) / 6.), alpha) # The best possible score of the most likely alive sequence. lower_bound_alive_scores = mtf.gather( alive_log_probs, mtf.constant(mesh, 0, dtype=tf.int32), beam_dim) / max_length_penalty # Now to compute the lowest score of a finished sequence in finished # If the sequence isn't finished, we multiply it's score by 0. since # scores are all -ve, taking the min will give us the score of the lowest # finished item. lowest_score_of_finished_in_finished = mtf.reduce_min( finished_scores * mtf.cast(finished_in_finished, finished_scores.dtype), reduced_dim=beam_dim) # If none of the sequences have finished, then the min will be 0 and # we have to replace it by -ve INF if it is. The score of any seq in alive # will be much higher than -ve INF and the termination condition will not # be met. lowest_score_of_finished_in_finished += ((1. - mtf.cast( mtf.reduce_any(finished_in_finished, reduced_dim=beam_dim), finished_scores.dtype)) * -INF) bound_is_met = mtf.reduce_all( mtf.greater(lowest_score_of_finished_in_finished, lower_bound_alive_scores)) return mtf.logical_and(mtf.less(i, decode_length), mtf.logical_not(bound_is_met))
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)