def beam_decoder(features, mode, vocab, encoder_outputs, hps): """Beam search decoder. Args: features: Dictionary of input Tensors. mode: train or eval. Keys from tf.estimator.ModeKeys. vocab: A list of strings of words in the vocabulary. encoder_outputs: output tensors from the encoder hps: Hyperparams. Returns: Decoder outputs """ assert mode is not tf.estimator.ModeKeys.TRAIN, "Not using beam in training." embeddings = encoder_outputs.embeddings mem_input = encoder_outputs.mem_input batch_size = tensor_utils.shape(features["src_len"], 0) src_len, src_inputs = features["src_len"], features["src_inputs"] src_mask = tf.sequence_mask(src_len, tf.shape(src_inputs)[1]) if hps.att_neighbor: neighbor_len, neighbor_inputs \ = features["neighbor_len"], features["neighbor_inputs"] neighbor_mask = tf.sequence_mask(neighbor_len, tf.shape(neighbor_inputs)[1]) inputs = tf.concat([src_inputs, neighbor_inputs], 1) # lens = src_len + neighbor_len mask = tf.concat([src_mask, neighbor_mask], axis=1) lens = tf.shape(mask)[1] * tf.ones([batch_size], tf.int32) else: inputs = features["src_inputs"] lens = features["src_len"] mask = src_mask sparse_inputs = None float_mask = tf.cast(mask, dtype=tf.float32) if hps.use_copy: sparse_inputs = sparse_map(tf.expand_dims(inputs, axis=2), tf.expand_dims(float_mask, axis=2), vocab.size()) sparse_inputs = sparse_tile_batch(sparse_inputs, multiplier=hps.beam_width) tiled_mask = tf.contrib.seq2seq.tile_batch(mask, multiplier=hps.beam_width) inputs = tf.contrib.seq2seq.tile_batch(inputs, multiplier=hps.beam_width) lens = tf.contrib.seq2seq.tile_batch(lens, multiplier=hps.beam_width) def _beam_decode(cell): """Beam decode.""" with tf.variable_scope("beam_decoder"): initial_state = cell.zero_state(batch_size=batch_size * hps.beam_width, dtype=tf.float32) if hps.use_bridge: h_state, c_state = encoder_outputs.states h_state = tf.contrib.seq2seq.tile_batch( h_state, multiplier=hps.beam_width) c_state = tf.contrib.seq2seq.tile_batch( c_state, multiplier=hps.beam_width) initial_cell_state = tf.contrib.rnn.LSTMStateTuple( h_state, c_state) initial_state = initial_state.clone( cell_state=(initial_cell_state, )) decoder = tf.contrib.seq2seq.BeamSearchDecoder( cell=cell, embedding=embeddings, start_tokens=tf.fill([batch_size], vocab.word2id(data.START_DECODING)), end_token=vocab.word2id(data.STOP_DECODING), initial_state=initial_state, beam_width=hps.beam_width, length_penalty_weight=hps.length_norm, coverage_penalty_weight=hps.cp) with tf.variable_scope("dynamic_decode", reuse=tf.AUTO_REUSE): decoder_outputs, _, decoder_len = tf.contrib.seq2seq.dynamic_decode( decoder=decoder, maximum_iterations=hps.max_dec_steps) return decoder_outputs, decoder_len # [batch_size*beam_width, src_len, encoder_dim] att_context = tf.contrib.seq2seq.tile_batch(encoder_outputs.att_context, multiplier=hps.beam_width) copy_context = None if hps.use_copy: copy_context = tf.contrib.seq2seq.tile_batch( encoder_outputs.copy_context, multiplier=hps.beam_width) with tf.variable_scope("attention", reuse=tf.AUTO_REUSE): if hps.att_type == "luong": attention = tf.contrib.seq2seq.LuongAttention( num_units=hps.decoder_dim, memory=att_context, memory_sequence_length=lens) elif hps.att_type == "bahdanau": attention = tf.contrib.seq2seq.BahdanauAttention( num_units=hps.decoder_dim, memory=att_context, memory_sequence_length=lens) elif hps.att_type == "hyper": attention = HyperAttention(num_units=hps.decoder_dim, mem_input=mem_input, hps=hps, memory=att_context, use_beam=True, memory_sequence_length=lens) elif hps.att_type == "my": attention = MyAttention(num_units=hps.decoder_dim, memory=att_context, memory_sequence_length=lens, mask=tiled_mask) with tf.variable_scope("rnn_decoder", reuse=tf.AUTO_REUSE): decoder_cell = get_rnn_cell(mode=mode, hps=hps, input_dim=hps.decoder_dim + hps.emb_dim, num_units=hps.decoder_dim, num_layers=hps.num_decoder_layers, dropout=hps.decoder_drop, mem_input=mem_input, use_beam=True, cell_type=hps.rnn_cell) with tf.variable_scope("attention_wrapper", reuse=tf.AUTO_REUSE): decoder_cell = tf.contrib.seq2seq.AttentionWrapper( decoder_cell, attention, attention_layer_size=hps.decoder_dim, alignment_history=hps.use_copy) with tf.variable_scope("output_projection", reuse=tf.AUTO_REUSE): weights = tf.transpose(embeddings) if hps.tie_embedding else None hidden_dim = hps.emb_dim if hps.tie_embedding else hps.decoder_dim decoder_cell = OutputWrapper( decoder_cell, num_layers=hps.num_mlp_layers, hidden_dim=hidden_dim, output_size=vocab.size() if hps.tie_embedding else hps.output_size, weights=weights, dropout=hps.out_drop, use_copy=hps.use_copy, encoder_emb=copy_context, sparse_inputs=sparse_inputs, mask=tf.cast(tiled_mask, dtype=tf.float32), hps=hps, mode=mode, reuse=tf.AUTO_REUSE) decoder_outputs, decoder_len = _beam_decode(decoder_cell) return DecoderOutputs(decoder_outputs=decoder_outputs, decoder_len=decoder_len)
def basic_decoder(features, mode, vocab, encoder_outputs, hps): """Decoder. Args: features: Dictionary of input Tensors. mode: train or eval. Keys from tf.estimator.ModeKeys. vocab: A list of strings of words in the vocabulary. encoder_outputs: output tensors from the encoder hps: Hyperparams. Returns: Decoder outputs """ embeddings = encoder_outputs.embeddings mem_input = encoder_outputs.mem_input batch_size = tensor_utils.shape(mem_input, 0) src_len, src_inputs = features["src_len"], features["src_inputs"] src_mask = tf.sequence_mask(src_len, tf.shape(src_inputs)[1]) if hps.att_neighbor: neighbor_len, neighbor_inputs \ = features["neighbor_len"], features["neighbor_inputs"] neighbor_mask = tf.sequence_mask(neighbor_len, tf.shape(neighbor_inputs)[1]) inputs = tf.concat([src_inputs, neighbor_inputs], 1) lens = src_len + neighbor_len mask = tf.concat([src_mask, neighbor_mask], axis=1) else: inputs = features["src_inputs"] lens = features["src_len"] mask = src_mask sparse_inputs = None float_mask = tf.cast(mask, dtype=tf.float32) if hps.use_copy: sparse_inputs = sparse_map(tf.expand_dims(inputs, axis=2), tf.expand_dims(float_mask, axis=2), vocab.size()) # [batch_size, dec_len] decoder_inputs = features["decoder_inputs"] # [batch_size, dec_len, emb_dim] decoder_input_emb = tf.nn.embedding_lookup(embeddings, decoder_inputs) if mode == tf.estimator.ModeKeys.TRAIN and hps.emb_drop > 0.: decoder_input_emb = tf.nn.dropout(decoder_input_emb, keep_prob=1.0 - hps.emb_drop) def _decode(cell, helper): """Decode function. Args: cell: rnn cell helper: a helper instance from tf.contrib.seq2seq. Returns: decoded outputs and lengths. """ with tf.variable_scope("decoder"): initial_state = cell.zero_state(batch_size, tf.float32) if hps.use_bridge: h_state, c_state = encoder_outputs.states initial_cell_state = tf.contrib.rnn.LSTMStateTuple( h_state, c_state) initial_state = initial_state.clone( cell_state=(initial_cell_state, )) decoder = tf.contrib.seq2seq.BasicDecoder( cell=cell, helper=helper, initial_state=initial_state) with tf.variable_scope("dynamic_decode", reuse=tf.AUTO_REUSE): decoder_outputs, _, decoder_len = tf.contrib.seq2seq.dynamic_decode( decoder=decoder, maximum_iterations=hps.max_dec_steps) return decoder_outputs, decoder_len att_context = encoder_outputs.att_context with tf.variable_scope("attention"): if hps.att_type == "luong": attention = tf.contrib.seq2seq.LuongAttention( num_units=hps.decoder_dim, memory=att_context, memory_sequence_length=lens) elif hps.att_type == "bahdanau": attention = tf.contrib.seq2seq.BahdanauAttention( num_units=hps.decoder_dim, memory=att_context, memory_sequence_length=lens) elif hps.att_type == "hyper": attention = HyperAttention(num_units=hps.decoder_dim, mem_input=mem_input, hps=hps, memory=att_context, use_beam=False, memory_sequence_length=lens) elif hps.att_type == "my": attention = MyAttention(num_units=hps.decoder_dim, memory=att_context, memory_sequence_length=lens, mask=mask) with tf.variable_scope("rnn_decoder"): decoder_cell = get_rnn_cell(mode=mode, hps=hps, input_dim=hps.decoder_dim + hps.emb_dim, num_units=hps.decoder_dim, num_layers=hps.num_decoder_layers, dropout=hps.decoder_drop, mem_input=mem_input, use_beam=False, cell_type=hps.rnn_cell) with tf.variable_scope("attention_wrapper"): decoder_cell = tf.contrib.seq2seq.AttentionWrapper( decoder_cell, attention, attention_layer_size=hps.decoder_dim, alignment_history=hps.use_copy) with tf.variable_scope("output_projection"): weights = tf.transpose(embeddings) if hps.tie_embedding else None hidden_dim = hps.emb_dim if hps.tie_embedding else hps.decoder_dim decoder_cell = OutputWrapper( decoder_cell, num_layers=hps.num_mlp_layers, hidden_dim=hidden_dim, output_size=hps.output_size, # output_size=vocab.size(), weights=weights, dropout=hps.out_drop, use_copy=hps.use_copy, encoder_emb=encoder_outputs.copy_context, sparse_inputs=sparse_inputs, mask=float_mask, mode=mode, hps=hps) if mode == tf.estimator.ModeKeys.TRAIN: if hps.sampling_probability > 0.: helper = tf.contrib.seq2seq.ScheduledEmbeddingTrainingHelper( inputs=decoder_input_emb, sequence_length=features["decoder_len"], embedding=embeddings, sampling_probability=hps.sampling_probability) else: helper = tf.contrib.seq2seq.TrainingHelper(decoder_input_emb, features["decoder_len"]) decoder_outputs, _ = _decode(decoder_cell, helper=helper) return DecoderOutputs(decoder_outputs=decoder_outputs, decoder_len=features["decoder_len"]), None # used to compute loss teacher_helper = tf.contrib.seq2seq.TrainingHelper(decoder_input_emb, features["decoder_len"]) teacher_decoder_outputs, _ = _decode(decoder_cell, helper=teacher_helper) helper = tf.contrib.seq2seq.GreedyEmbeddingHelper( embedding=embeddings, start_tokens=tf.fill([batch_size], vocab.word2id(data.START_DECODING)), end_token=vocab.word2id(data.STOP_DECODING)) decoder_outputs, decoder_len = _decode(decoder_cell, helper=helper) return DecoderOutputs( decoder_outputs=decoder_outputs, decoder_len=decoder_len), \ DecoderOutputs( decoder_outputs=teacher_decoder_outputs, decoder_len=features["decoder_len"] )