(symbols_to_logits_fn,
                  initial_ids,
                  beam_size,
                  decode_length,
                  vocab_size,
                  alpha,
                  eos_id,
                  states=None,
                  stop_early=True,
                  INF=1. * 1e20)
return finished_seq, finished_scores, states
Beam search with length penalties.
Beam search with length penalties.
Beam search with length penalties .
def beam_search(symbols_to_logits_fn,
                  initial_ids,
                  beam_size,
                  decode_length,
                  vocab_size,
                  alpha,
                  eos_id,
                  states=None,
                  stop_early=True,
                  INF=1. * 1e20):
    """Beam search with length penalties."""
    batch_size = utils.shape_list(initial_ids)[0]

    initial_log_probs = tf.constant([[0.] + [-INF] * (beam_size - 1)])
    # (batch_size, beam_size)
    alive_log_probs = tf.tile(initial_log_probs, [batch_size, 1])

    alive_seq = utils.expand_to_beam_size(initial_ids, beam_size)
    # (batch_size, beam_size, 1)
    alive_seq = tf.expand_dims(alive_seq, axis=2)
    if states:
      states = nest.map_structure(
          lambda state: utils.expand_to_beam_size(state, beam_size), states)
    else:
      states = {}

    # (batch_size, beam_size, 1)
    finished_seq = tf.zeros(utils.shape_list(alive_seq), tf.int32)
    # (batch_size, beam_size)
    finished_scores = tf.ones([batch_size, beam_size]) * -INF
    # (batch_size, beam_size)
    finished_flags = tf.zeros([batch_size, beam_size], tf.bool)

    def grow_finished(finished_seq, finished_scores, finished_flags, curr_seq,
                      curr_scores, curr_finished):
      """
        Given sequences and scores from finished sequence and current finished sequence
        , will gather the top k=beam size sequences to update finished seq.
      """
      # padding zero for finished seq
      finished_seq = tf.concat(
          [finished_seq,
           tf.zeros([batch_size, beam_size, 1], tf.int32)],
          axis=2)

      # mask unfinished curr seq
      curr_scores += (1. - tf.to_float(curr_finished)) * -INF

      # concatenating the sequences and scores along beam axis
      # (batch_size, 2xbeam_size, seq_len)
      curr_finished_seq = tf.concat([finished_seq, curr_seq], axis=1)
      curr_finished_scores = tf.concat([finished_scores, curr_scores], axis=1)
      curr_finished_flags = tf.concat([finished_flags, curr_finished], axis=1)
      return utils.compute_topk_scores_and_seq(curr_finished_seq,
                                               curr_finished_scores,
                                               curr_finished_scores,
                                               curr_finished_flags, beam_size,
                                               batch_size, "grow_finished")

    def grow_alive(curr_seq, curr_scores, curr_log_probs, curr_finished,
                   states):
      """Given sequences and scores, will gather the top k=beam size sequences."""
      curr_scores += tf.to_float(curr_finished) * -INF
      return utils.compute_topk_scores_and_seq(curr_seq, curr_scores,
                                               curr_log_probs, curr_finished,
                                               beam_size, batch_size,
                                               "grow_alive", states)

    def grow_topk(i, alive_seq, alive_log_probs, states):
      """Inner beam search loop."""
      flat_ids = tf.reshape(alive_seq, [batch_size * beam_size, -1])

      # (batch_size * beam_size, decoded_length)
      if states:
        flat_states = nest.map_structure(utils.merge_beam_dim, states)
        flat_logits, flat_states = symbols_to_logits_fn(flat_ids, i,
                                                        flat_states)
        states = nest.map_structure(
            lambda t: utils.unmerge_beam_dim(t, batch_size, beam_size),
            flat_states)
      else:
        flat_logits = symbols_to_logits_fn(flat_ids)

      logits = tf.reshape(flat_logits, [batch_size, beam_size, -1])
      candidate_log_probs = utils.log_prob_from_logits(logits)
      log_probs = candidate_log_probs + tf.expand_dims(alive_log_probs, axis=2)

      length_penalty = tf.pow(((5. + tf.to_float(i + 1)) / 6.), alpha)

      curr_scores = log_probs / length_penalty
      flat_curr_scores = tf.reshape(curr_scores, [-1, beam_size * vocab_size])

      topk_scores, topk_ids = tf.nn.top_k(flat_curr_scores, k=beam_size * 2)
      topk_log_probs = topk_scores * length_penalty

      topk_beam_index = topk_ids // vocab_size
      topk_ids %= vocab_size  # Unflatten the ids
      batch_pos = utils.compute_batch_indices(batch_size, beam_size * 2)
      topk_coordinates = tf.stack([batch_pos, topk_beam_index], axis=2)

      topk_seq = tf.gather_nd(alive_seq, topk_coordinates)
      if states:
        states = nest.map_structure(
            lambda state: tf.gather_nd(state, topk_coordinates), states)
      topk_seq = tf.concat([topk_seq, tf.expand_dims(topk_ids, axis=2)], axis=2)

      topk_finished = tf.equal(topk_ids, eos_id)

      return topk_seq, topk_log_probs, topk_scores, topk_finished, states

    def inner_loop(i, alive_seq, alive_log_probs, finished_seq, finished_scores,
                   finished_flags, states):
      """Inner beam search loop."""
      topk_seq, topk_log_probs, topk_scores, topk_finished, states = grow_topk(
          i, alive_seq, alive_log_probs, states)
      alive_seq, alive_log_probs, _, states = grow_alive(
          topk_seq, topk_scores, topk_log_probs, topk_finished, states)
      finished_seq, finished_scores, finished_flags, _ = grow_finished(
          finished_seq, finished_scores, finished_flags, topk_seq, topk_scores,
          topk_finished)

      return (i + 1, alive_seq, alive_log_probs, finished_seq, finished_scores,
              finished_flags, states)

    def _is_finished(i, unused_alive_seq, alive_log_probs, unused_finished_seq,
                     finished_scores, unused_finished_in_finished,
                     unused_states):
      """Checking termination condition.
      """
      max_length_penalty = tf.pow(((5. + tf.to_float(decode_length)) / 6.),
                                  alpha)
      lower_bound_alive_scores = alive_log_probs[:, 0] / max_length_penalty

      if not stop_early:
        lowest_score_of_finished_in_finished = tf.reduce_min(finished_scores)
      else:
        lowest_score_of_finished_in_finished = tf.reduce_max(
            finished_scores, axis=1)

      bound_is_met = tf.reduce_all(
          tf.greater(lowest_score_of_finished_in_finished,
                     lower_bound_alive_scores))

      return tf.logical_and(
          tf.less(i, decode_length), tf.logical_not(bound_is_met))

    inner_shape = tf.TensorShape([None, None, None])

    state_struc = nest.map_structure(utils.get_state_shape_invariants, states)
    (_, alive_seq, alive_log_probs, finished_seq, finished_scores,
     finished_flags, states) = tf.while_loop(
         _is_finished,
         inner_loop, [
             tf.constant(0), alive_seq, alive_log_probs, finished_seq,
             finished_scores, finished_flags, states
         ],
         shape_invariants=[
             tf.TensorShape([]), inner_shape,
             alive_log_probs.get_shape(), inner_shape,
             finished_scores.get_shape(),
             finished_flags.get_shape(), state_struc
         ],
         parallel_iterations=1,
         back_prop=False)

    alive_seq.set_shape((None, beam_size, None))
    finished_seq.set_shape((None, beam_size, None))
    finished_seq = tf.where(
        tf.reduce_any(finished_flags, 1), finished_seq, alive_seq)
    finished_scores = tf.where(
        tf.reduce_any(finished_flags, 1), finished_scores, alive_log_probs)
    return finished_seq, finished_scores, states
def beam_search ( symbols_to_logits_fn , initial_ids , beam_size , decode_length , vocab_size , alpha , eos_id , states = None , stop_early = True , INF = 1. * 1e20 ) : batch_size = utils . shape_list ( initial_ids ) [ 0 ] initial_log_probs = tf . constant ( [ [ 0. ] + [ - INF ] * ( beam_size - 1 ) ] ) # (batch_size, beam_size) alive_log_probs = tf . tile ( initial_log_probs , [ batch_size , 1 ] ) alive_seq = utils . expand_to_beam_size ( initial_ids , beam_size ) # (batch_size, beam_size, 1) alive_seq = tf . expand_dims ( alive_seq , axis = 2 ) if states : states = nest . map_structure ( lambda state : utils . expand_to_beam_size ( state , beam_size ) , states ) else : states = { } # (batch_size, beam_size, 1) finished_seq = tf . zeros ( utils . shape_list ( alive_seq ) , tf . int32 ) # (batch_size, beam_size) finished_scores = tf . ones ( [ batch_size , beam_size ] ) * - INF # (batch_size, beam_size) finished_flags = tf . zeros ( [ batch_size , beam_size ] , tf . bool ) def grow_finished ( finished_seq , finished_scores , finished_flags , curr_seq , curr_scores , curr_finished ) : """
        Given sequences and scores from finished sequence and current finished sequence
        , will gather the top k=beam size sequences to update finished seq.
      """ # padding zero for finished seq finished_seq = tf . concat ( [ finished_seq , tf . zeros ( [ batch_size , beam_size , 1 ] , tf . int32 ) ] , axis = 2 ) # mask unfinished curr seq curr_scores += ( 1. - tf . to_float ( curr_finished ) ) * - INF # concatenating the sequences and scores along beam axis # (batch_size, 2xbeam_size, seq_len) curr_finished_seq = tf . concat ( [ finished_seq , curr_seq ] , axis = 1 ) curr_finished_scores = tf . concat ( [ finished_scores , curr_scores ] , axis = 1 ) curr_finished_flags = tf . concat ( [ finished_flags , curr_finished ] , axis = 1 ) return utils . compute_topk_scores_and_seq ( curr_finished_seq , curr_finished_scores , curr_finished_scores , curr_finished_flags , beam_size , batch_size , "grow_finished" ) def grow_alive ( curr_seq , curr_scores , curr_log_probs , curr_finished , states ) : """Given sequences and scores, will gather the top k=beam size sequences.""" curr_scores += tf . to_float ( curr_finished ) * - INF return utils . compute_topk_scores_and_seq ( curr_seq , curr_scores , curr_log_probs , curr_finished , beam_size , batch_size , "grow_alive" , states ) def grow_topk ( i , alive_seq , alive_log_probs , states ) : """Inner beam search loop.""" flat_ids = tf . reshape ( alive_seq , [ batch_size * beam_size , - 1 ] ) # (batch_size * beam_size, decoded_length) if states : flat_states = nest . map_structure ( utils . merge_beam_dim , states ) flat_logits , flat_states = symbols_to_logits_fn ( flat_ids , i , flat_states ) states = nest . map_structure ( lambda t : utils . unmerge_beam_dim ( t , batch_size , beam_size ) , flat_states ) else : flat_logits = symbols_to_logits_fn ( flat_ids ) logits = tf . reshape ( flat_logits , [ batch_size , beam_size , - 1 ] ) candidate_log_probs = utils . log_prob_from_logits ( logits ) log_probs = candidate_log_probs + tf . expand_dims ( alive_log_probs , axis = 2 ) length_penalty = tf . pow ( ( ( 5. + tf . to_float ( i + 1 ) ) / 6. ) , alpha ) curr_scores = log_probs / length_penalty flat_curr_scores = tf . reshape ( curr_scores , [ - 1 , beam_size * vocab_size ] ) topk_scores , topk_ids = tf . nn . top_k ( flat_curr_scores , k = beam_size * 2 ) topk_log_probs = topk_scores * length_penalty topk_beam_index = topk_ids // vocab_size topk_ids %= vocab_size # Unflatten the ids batch_pos = utils . compute_batch_indices ( batch_size , beam_size * 2 ) topk_coordinates = tf . stack ( [ batch_pos , topk_beam_index ] , axis = 2 ) topk_seq = tf . gather_nd ( alive_seq , topk_coordinates ) if states : states = nest . map_structure ( lambda state : tf . gather_nd ( state , topk_coordinates ) , states ) topk_seq = tf . concat ( [ topk_seq , tf . expand_dims ( topk_ids , axis = 2 ) ] , axis = 2 ) topk_finished = tf . equal ( topk_ids , eos_id ) return topk_seq , topk_log_probs , topk_scores , topk_finished , states def inner_loop ( i , alive_seq , alive_log_probs , finished_seq , finished_scores , finished_flags , states ) : """Inner beam search loop.""" topk_seq , topk_log_probs , topk_scores , topk_finished , states = grow_topk ( i , alive_seq , alive_log_probs , states ) alive_seq , alive_log_probs , _ , states = grow_alive ( topk_seq , topk_scores , topk_log_probs , topk_finished , states ) finished_seq , finished_scores , finished_flags , _ = grow_finished ( finished_seq , finished_scores , finished_flags , topk_seq , topk_scores , topk_finished ) return ( i + 1 , alive_seq , alive_log_probs , finished_seq , finished_scores , finished_flags , states ) def _is_finished ( i , unused_alive_seq , alive_log_probs , unused_finished_seq , finished_scores , unused_finished_in_finished , unused_states ) : """Checking termination condition.
      """ max_length_penalty = tf . pow ( ( ( 5. + tf . to_float ( decode_length ) ) / 6. ) , alpha ) lower_bound_alive_scores = alive_log_probs [ : , 0 ] / max_length_penalty if not stop_early : lowest_score_of_finished_in_finished = tf . reduce_min ( finished_scores ) else : lowest_score_of_finished_in_finished = tf . reduce_max ( finished_scores , axis = 1 ) bound_is_met = tf . reduce_all ( tf . greater ( lowest_score_of_finished_in_finished , lower_bound_alive_scores ) ) return tf . logical_and ( tf . less ( i , decode_length ) , tf . logical_not ( bound_is_met ) ) inner_shape = tf . TensorShape ( [ None , None , None ] ) state_struc = nest . map_structure ( utils . get_state_shape_invariants , states ) ( _ , alive_seq , alive_log_probs , finished_seq , finished_scores , finished_flags , states ) = tf . while_loop ( _is_finished , inner_loop , [ tf . constant ( 0 ) , alive_seq , alive_log_probs , finished_seq , finished_scores , finished_flags , states ] , shape_invariants = [ tf . TensorShape ( [ ] ) , inner_shape , alive_log_probs . get_shape ( ) , inner_shape , finished_scores . get_shape ( ) , finished_flags . get_shape ( ) , state_struc ] , parallel_iterations = 1 , back_prop = False ) alive_seq . set_shape ( ( None , beam_size , None ) ) finished_seq . set_shape ( ( None , beam_size , None ) ) finished_seq = tf . where ( tf . reduce_any ( finished_flags , 1 ) , finished_seq , alive_seq ) finished_scores = tf . where ( tf . reduce_any ( finished_flags , 1 ) , finished_scores , alive_log_probs ) return finished_seq , finished_scores , states