Function decoding_body

examples/tensorflow/decoder/utils/decoding.py:97–198

(word_ids,
                  step,
                  memory,
                  memory_sequence_length,
                  my_cache,
                  op_self_cache,
                  op_mem_cache,
                  embedding_table,
                  decoding_args,
                  decoder_type,
                  sequence_lengths)

    return ids, lengths

def decoding_body(word_ids,
                  step,
                  memory,
                  memory_sequence_length,
                  my_cache,
                  op_self_cache,
                  op_mem_cache,
                  embedding_table,
                  decoding_args,
                  decoder_type,
                  sequence_lengths):
    
    decoder_args = decoding_args.decoder_args
    hidden_dim = decoder_args.hidden_dim
    k_init_range = decoder_args.kernel_init_range
    data_type = decoder_args.dtype
    
    batchxbeam = tf.shape(word_ids)[0]
    # [batch_size * beam_width, hidden_dim]
    inputs = tf.nn.embedding_lookup(embedding_table, word_ids)
    # [batch_size * beam_width, 1, hidden_dim]
    inputs = tf.expand_dims(inputs, 1)

    inputs *= hidden_dim**0.5
    position_encoder = SinusoidalPositionEncoder()
    if position_encoder is not None:
        position_encoding_table = position_encoder._create_position_encoding_table(decoding_args.max_seq_len, hidden_dim, data_type)
        position_encoding_val = position_encoding_table[step]
        position_encoding_val = tf.reshape(position_encoding_val, [1, 1, -1])
        position_encoding_val = tf.tile(position_encoding_val, [batchxbeam, 1, 1])
        inputs = inputs + position_encoding_val
        
    with tf.variable_scope("decoder", reuse=tf.AUTO_REUSE):
        tf_result = tf_decoder(decoder_args=decoder_args,
                                inputs=inputs,
                                memory=memory,
                                memory_sequence_length=memory_sequence_length,
                                step=step,
                                cache=my_cache)

        if decoder_type != 0:
            decoder_vars = tf.global_variables()
            decoder_vars_start_id = 0
            while decoder_vars_start_id < len(decoder_vars):
                if decoder_vars[decoder_vars_start_id].name.find("transformer/decoder/layer") != -1:
                    break
                decoder_vars_start_id += 1
            decoder_vars = decoder_vars[decoder_vars_start_id:]
            decoder_var_dict = {}
            for v in decoder_vars:
                decoder_var_dict[v.name] = v

            psuedo_input = []
            if decoder_type == 2:
                psuedo_input = tf_result
                
            op_result, op_self_cache, op_mem_cache = op_decoder(inputs,
                                                                memory,