def fast_decode(encoder_output,
                encoder_decoder_attention_bias,
                symbols_to_logits_fn,
                hparams,
                decode_length,
                vocab_size,
                init_cache_fn=_init_transformer_cache,
                beam_size=1,
                top_beams=1,
                alpha=1.0,
                sos_id=0,
                eos_id=beam_search.EOS_ID,
                batch_size=None,
                force_decode_length=False,
                scope_prefix="body/",
                cache=None):
  """Given encoder output and a symbols to logits function, does fast decoding.

  Implements both greedy and beam search decoding, uses beam search iff
  beam_size > 1, otherwise beam search related arguments are ignored.

  Args:
    encoder_output: Output from encoder.
    encoder_decoder_attention_bias: a bias tensor for use in encoder-decoder
      attention
    symbols_to_logits_fn: Incremental decoding; function mapping triple `(ids,
      step, cache)` to symbol logits.
    hparams: run hyperparameters
    decode_length: an integer.  How many additional timesteps to decode.
    vocab_size: Output vocabulary size.
    init_cache_fn: Function that returns the initial cache dict.
    beam_size: number of beams.
    top_beams: an integer. How many of the beams to return.
    alpha: Float that controls the length penalty. larger the alpha, stronger
      the preference for longer translations.
    sos_id: End-of-sequence symbol in beam search.
    eos_id: End-of-sequence symbol in beam search.
    batch_size: an integer scalar - must be passed if there is no input
    force_decode_length: bool, whether to force the full decode length, or if
      False, stop when all beams hit eos_id.
    scope_prefix: str, prefix for decoder layer variable scopes.
    cache: cache dictionary for additional predictions.

  Returns:
      A dict of decoding results {
          "outputs": integer `Tensor` of decoded ids of shape
              [batch_size, <= decode_length] if top_beams == 1 or
              [batch_size, top_beams, <= decode_length] otherwise
          "scores": decoding log probs from the beam search,
              None if using greedy decoding (beam_size=1)
      }

    Raises:
      NotImplementedError: If beam size > 1 with partial targets.
  """
  if encoder_output is not None:
    batch_size = common_layers.shape_list(encoder_output)[0]

  cache = init_cache_fn(
      cache=cache,
      hparams=hparams,
      batch_size=batch_size,
      attention_init_length=0,
      encoder_output=encoder_output,
      encoder_decoder_attention_bias=encoder_decoder_attention_bias,
      scope_prefix=scope_prefix)

  if beam_size > 1:  # Beam Search
    initial_ids = sos_id * tf.ones([batch_size], dtype=tf.int32)
    decoded_ids, scores, cache = beam_search.beam_search(
        symbols_to_logits_fn,
        initial_ids,
        beam_size,
        decode_length,
        vocab_size,
        alpha,
        states=cache,
        eos_id=eos_id,
        stop_early=(top_beams == 1))

    if top_beams == 1:
      decoded_ids = decoded_ids[:, 0, 1:]
      scores = scores[:, 0]
    else:
      decoded_ids = decoded_ids[:, :top_beams, 1:]
      scores = scores[:, :top_beams]
  else:  # Greedy

    def inner_loop(i, hit_eos, next_id, decoded_ids, cache, log_prob):
      """One step of greedy decoding."""
      logits, cache = symbols_to_logits_fn(next_id, i, cache)
      log_probs = common_layers.log_prob_from_logits(logits)
      temperature = getattr(hparams, "sampling_temp", 0.0)
      keep_top = getattr(hparams, "sampling_keep_top_k", -1)
      if hparams.sampling_method == "argmax":
        temperature = 0.0
      next_id = common_layers.sample_with_temperature(
          logits, temperature, keep_top)
      hit_eos |= tf.equal(next_id, eos_id)

      log_prob_indices = tf.stack([tf.range(tf.to_int64(batch_size)), next_id],
                                  axis=1)
      log_prob += tf.gather_nd(log_probs, log_prob_indices)

      next_id = tf.expand_dims(next_id, axis=1)
      decoded_ids = tf.concat([decoded_ids, next_id], axis=1)
      return i + 1, hit_eos, next_id, decoded_ids, cache, log_prob

    def is_not_finished(i, hit_eos, *_):
      finished = i >= decode_length
      if not force_decode_length:
        finished |= tf.reduce_all(hit_eos)
      return tf.logical_not(finished)

    decoded_ids = tf.zeros([batch_size, 0], dtype=tf.int64)
    hit_eos = tf.fill([batch_size], False)
    next_id = sos_id * tf.ones([batch_size, 1], dtype=tf.int64)
    initial_log_prob = tf.zeros([batch_size], dtype=tf.float32)
    _, _, _, decoded_ids, cache, log_prob = tf.while_loop(
        is_not_finished,
        inner_loop, [
            tf.constant(0), hit_eos, next_id, decoded_ids, cache,
            initial_log_prob
        ],
        shape_invariants=[
            tf.TensorShape([]),
            tf.TensorShape([None]),
            tf.TensorShape([None, None]),
            tf.TensorShape([None, None]),
            nest.map_structure(beam_search.get_state_shape_invariants, cache),
            tf.TensorShape([None]),
        ])
    scores = log_prob

  return {"outputs": decoded_ids, "scores": scores, "cache": cache}
def fast_decode ( encoder_output , encoder_decoder_attention_bias , symbols_to_logits_fn , hparams , decode_length , vocab_size , init_cache_fn = _init_transformer_cache , beam_size = 1 , top_beams = 1 , alpha = 1.0 , sos_id = 0 , eos_id = beam_search . EOS_ID , batch_size = None , force_decode_length = False , scope_prefix = "body/" , cache = None ) : if encoder_output is not None : batch_size = common_layers . shape_list ( encoder_output ) [ 0 ] cache = init_cache_fn ( cache = cache , hparams = hparams , batch_size = batch_size , attention_init_length = 0 , encoder_output = encoder_output , encoder_decoder_attention_bias = encoder_decoder_attention_bias , scope_prefix = scope_prefix ) if beam_size > 1 : initial_ids = sos_id * tf . ones ( [ batch_size ] , dtype = tf . int32 ) decoded_ids , scores , cache = beam_search . beam_search ( symbols_to_logits_fn , initial_ids , beam_size , decode_length , vocab_size , alpha , states = cache , eos_id = eos_id , stop_early = ( top_beams == 1 ) ) if top_beams == 1 : decoded_ids = decoded_ids [ : , 0 , 1 : ] scores = scores [ : , 0 ] else : decoded_ids = decoded_ids [ : , : top_beams , 1 : ] scores = scores [ : , : top_beams ] else : def inner_loop ( i , hit_eos , next_id , decoded_ids , cache , log_prob ) : logits , cache = symbols_to_logits_fn ( next_id , i , cache ) log_probs = common_layers . log_prob_from_logits ( logits ) temperature = getattr ( hparams , "sampling_temp" , 0.0 ) keep_top = getattr ( hparams , "sampling_keep_top_k" , - 1 ) if hparams . sampling_method == "argmax" : temperature = 0.0 next_id = common_layers . sample_with_temperature ( logits , temperature , keep_top ) hit_eos |= tf . equal ( next_id , eos_id ) log_prob_indices = tf . stack ( [ tf . range ( tf . to_int64 ( batch_size ) ) , next_id ] , axis = 1 ) log_prob += tf . gather_nd ( log_probs , log_prob_indices ) next_id = tf . expand_dims ( next_id , axis = 1 ) decoded_ids = tf . concat ( [ decoded_ids , next_id ] , axis = 1 ) return i + 1 , hit_eos , next_id , decoded_ids , cache , log_prob def is_not_finished ( i , hit_eos , * _ ) : finished = i >= decode_length if not force_decode_length : finished |= tf . reduce_all ( hit_eos ) return tf . logical_not ( finished ) decoded_ids = tf . zeros ( [ batch_size , 0 ] , dtype = tf . int64 ) hit_eos = tf . fill ( [ batch_size ] , False ) next_id = sos_id * tf . ones ( [ batch_size , 1 ] , dtype = tf . int64 ) initial_log_prob = tf . zeros ( [ batch_size ] , dtype = tf . float32 ) _ , _ , _ , decoded_ids , cache , log_prob = tf . while_loop ( is_not_finished , inner_loop , [ tf . constant ( 0 ) , hit_eos , next_id , decoded_ids , cache , initial_log_prob ] , shape_invariants = [ tf . TensorShape ( [ ] ) , tf . TensorShape ( [ None ] ) , tf . TensorShape ( [ None , None ] ) , tf . TensorShape ( [ None , None ] ) , nest . map_structure ( beam_search . get_state_shape_invariants , cache ) , tf . TensorShape ( [ None ] ) , ] ) scores = log_prob return { "outputs" : decoded_ids , "scores" : scores , "cache" : cache }
Given encoder output and a symbols to logits function, does fast decoding.

  Implements both greedy and beam search decoding, uses beam search iff
  beam_size > 1, otherwise beam search related arguments are ignored.

  Args:
    encoder_output: Output from encoder.
    encoder_decoder_attention_bias: a bias tensor for use in encoder-decoder
      attention
    symbols_to_logits_fn: Incremental decoding; function mapping triple `(ids,
      step, cache)` to symbol logits.
    hparams: run hyperparameters
    decode_length: an integer.  How many additional timesteps to decode.
    vocab_size: Output vocabulary size.
    init_cache_fn: Function that returns the initial cache dict.
    beam_size: number of beams.
    top_beams: an integer. How many of the beams to return.
    alpha: Float that controls the length penalty. larger the alpha, stronger
      the preference for longer translations.
    sos_id: End-of-sequence symbol in beam search.
    eos_id: End-of-sequence symbol in beam search.
    batch_size: an integer scalar - must be passed if there is no input
    force_decode_length: bool, whether to force the full decode length, or if
      False, stop when all beams hit eos_id.
    scope_prefix: str, prefix for decoder layer variable scopes.
    cache: cache dictionary for additional predictions.

  Returns:
      A dict of decoding results {
          "outputs": integer `Tensor` of decoded ids of shape
              [batch_size, <= decode_length] if top_beams == 1 or
              [batch_size, top_beams, <= decode_length] otherwise
          "scores": decoding log probs from the beam search,
              None if using greedy decoding (beam_size=1)
      }

    Raises:
      NotImplementedError: If beam size > 1 with partial targets.