MCPcopy Create free account
hub / github.com/NVIDIA/FasterTransformer / ft_decoding

Function ft_decoding

examples/tensorflow/decoding/utils/ft_decoding.py:43–167  ·  view source on GitHub ↗

Run the decoding with beam search by TensorFlow. Args: memory_tensor: A tf.tensor with shape [batch_size * beam_width, max(memory_sequence_length), encoder_hidden_dimension]. The results of encoder transformer layer. The rank must be 3.

(memory_tensor,
                        memory_sequence_length,
                        embedding_table,
                        decoding_vars,
                        decoding_args,
                        using_model_var=True,
                        checkpoint_filename=None)

Source from the content-addressed store, hash-verified

41 return ids, lengths
42
43def ft_decoding(memory_tensor,
44 memory_sequence_length,
45 embedding_table,
46 decoding_vars,
47 decoding_args,
48 using_model_var=True,
49 checkpoint_filename=None):
50 '''
51 Run the decoding with beam search by TensorFlow.
52
53 Args:
54 memory_tensor: A tf.tensor with shape [batch_size * beam_width, max(memory_sequence_length), encoder_hidden_dimension].
55 The results of encoder transformer layer. The rank must be 3.
56 Note that it must be extended by beam_width times.
57 memory_sequence_length: A tf.Tensor with shape [batch_size * beam_width], type tf.int.
58 The length of each sentence of results of encoder.
59 Note that it must be extended by beam_width times.
60 embedding_table: A tf.Tensor with shape [vocab_size, hidden_dimension].
61 The embedding table of embedding lookup for each step.
62 decoder_vars: A list of tf.Tensor. The variables for decoding. A list of model variables of TensorFlow model.
63 decoder_args: The arguments for decoding. The details are in the class "DecodingBeamsearchArgument" of common.py
64 using_model_var: A bool value. Using the model variables of TensorFlow or not.
65 The details are described in 'preprocess_decoder_var' function in the following.
66 checkpoint_filename: A string. The checkpoint file name of storing the values of model.
67 The details are described in 'preprocess_decoder_var' function in the following.
68 Outputs:
69 finalized_output_ids: A tf.Tensor with shape [batch_size, beam_width, max(sequence_lengths)], with tf.int type.
70 Finalized output_ids by beam search algorithm and parent_ids.
71 finalized_sequence_lengths: A tf.Tensor with shape [batch_size * beam_width], with int type.
72 Finalized sequence_lengths by beam search algorithm and parent_ids.
73 output_ids: A tf.Tensor with shape [batch_size, beam_width, max(sequence_lengths)], with tf.int type.
74 The results of decoding. It contains the id of token of vocabulary.
75 parent_ids: A tf.Tensor with shape [batch_size, beam_width, max(sequence_lengths)], with tf.int type.
76 The beam index of output ids for each step.
77 sequence_lengths: A tf.Tensor with shape [batch_size * beam_width], with int type.
78 '''
79
80 decoder_args = decoding_args.decoder_args
81 decoding_op_module = tf.load_op_library(os.path.join('./lib/libtf_decoding.so'))
82
83 extended_memory = tf.contrib.seq2seq.tile_batch(
84 memory_tensor, multiplier=decoder_args.beam_width)
85 extended_memory_sequence_length = tf.contrib.seq2seq.tile_batch(
86 memory_sequence_length, multiplier=decoder_args.beam_width)
87
88 position_encoder = SinusoidalPositionEncoder()
89 position_encoding_table = position_encoder._create_position_encoding_table(
90 decoding_args.max_seq_len, decoder_args.head_num * decoder_args.size_per_head, decoder_args.dtype)
91 # shape of position_encoding_table: [max_seq_len, hidden_dim]
92
93 cross_key_kernel_list = []
94 cross_value_kernel_list = []
95 cross_key_bias_list = []
96 cross_value_bias_list = []
97
98 var_dict = {}
99 for v in decoding_vars:
100 var_dict[v.name] = v

Callers 2

translateFunction · 0.90

Tested by

no test coverage detected