# Path: simple_model.py
# class S2SAttentionModel(S2SModel):
#     def get_decoder_cell(self, encoder_state, encoder_outputs, decoder_cell,
#                          num_units, source_sequence_length):
#         attention_states = encoder_outputs
#         attention_mechanism = tf.contrib.seq2seq.LuongAttention(
#             num_units, attention_states,
#             memory_sequence_length=source_sequence_length)
#         decoder_cell = tf.contrib.seq2seq.AttentionWrapper(
#             decoder_cell, attention_mechanism,
#             attention_layer_size=num_units)
# 
#         initial_state = decoder_cell.zero_state(dtype=tf.float32,
#                                                 batch_size=self.batch_size)
#         initial_state = initial_state.clone(cell_state=encoder_state)
# 
#         return decoder_cell, initial_state


from simple_model import S2SAttentionModel
import tensorflow as tf
        decoder_cell = tf.nn.rnn_cell.BasicLSTMCell(num_units)
        decoder_cell, encoder_state = self.get_decoder_cell(
            encoder_state, encoder_outputs, decoder_cell,
            num_units, source_sequence_length)

        ones = tf.ones(shape=self.batch_size, dtype=tf.int32)
        lengths = ones * self.max_output_length
        helper = tf.contrib.seq2seq.TrainingHelper(decoder_emb_inp, lengths)
        decoder = tf.contrib.seq2seq.BasicDecoder(decoder_cell, helper,
                                                  encoder_state)
        outputs, final_state, final_sequence_lengths = tf.contrib.seq2seq.dynamic_decode(
            decoder)
        logits = outputs.rnn_output

        score = self.ff(logits, 1, 32, self.U)
        return score

    def create_model(self):
        self.x = tf.placeholder(tf.int64, shape=(None, self.max_input_length,), name='x')
        self.y = tf.placeholder(tf.int64, shape=(None, self.max_output_length,), name='y')
        self.x_len = tf.placeholder(tf.int32, shape=(None,), name='x_len')
        self.y_len = tf.placeholder(tf.int32, shape=(None,), name='y_len')
        self.noise_weight = tf.placeholder(tf.float32, shape=(), name='noise_weight')
        self.batch_size = tf.shape(self.x_len)

        with tf.variable_scope("compute_switch"):
            switch_score = self.get_embeddings(self.x, 1)
            self.switch = tf.nn.sigmoid(
                switch_score / self.args.switch_temperature)
        # masks

# Path: simple_model.py
# class S2SAttentionModel(S2SModel):
#     def get_decoder_cell(self, encoder_state, encoder_outputs, decoder_cell,
#                          num_units, source_sequence_length):
#         attention_states = encoder_outputs
#         attention_mechanism = tf.contrib.seq2seq.LuongAttention(
#             num_units, attention_states,
#             memory_sequence_length=source_sequence_length)
#         decoder_cell = tf.contrib.seq2seq.AttentionWrapper(
#             decoder_cell, attention_mechanism,
#             attention_layer_size=num_units)
# 
#         initial_state = decoder_cell.zero_state(dtype=tf.float32,
#                                                 batch_size=self.batch_size)
#         initial_state = initial_state.clone(cell_state=encoder_state)
# 
#         return decoder_cell, initial_state


# Path: extended_model.py
from simple_model import S2SAttentionModel
import tensorflow as tf

        decoder_cell = tf.nn.rnn_cell.BasicLSTMCell(num_units)
        decoder_cell, encoder_state = self.get_decoder_cell(
            encoder_state, encoder_outputs, decoder_cell,
            num_units, source_sequence_length)

        ones = tf.ones(shape=self.batch_size, dtype=tf.int32)
        lengths = ones * self.max_output_length
        helper = tf.contrib.seq2seq.TrainingHelper(decoder_emb_inp, lengths)
        decoder = tf.contrib.seq2seq.BasicDecoder(decoder_cell, helper,
                                                  encoder_state)
        outputs, final_state, final_sequence_lengths = tf.contrib.seq2seq.dynamic_decode(
            decoder)
        logits = outputs.rnn_output

        score = self.ff(logits, 1, 32, self.U)
        return score

    def create_model(self):
        self.x = tf.placeholder(tf.int64, shape=(None, self.max_input_length,), name='x')
        self.y = tf.placeholder(tf.int64, shape=(None, self.max_output_length,), name='y')
        self.x_len = tf.placeholder(tf.int32, shape=(None,), name='x_len')
        self.y_len = tf.placeholder(tf.int32, shape=(None,), name='y_len')
        self.noise_weight = tf.placeholder(tf.float32, shape=(), name='noise_weight')
        self.batch_size = tf.shape(self.x_len)

        with tf.variable_scope("compute_switch"):
            switch_score = self.get_embeddings(self.x, 1)
            self.switch = tf.nn.sigmoid(
                switch_score / self.args.switch_temperature)
        # masks

        self.target_mask_float = tf.sequence_mask(