text: # Copyright 2017 Google Inc. All Rights Reserved.
  # Copyright 2019 Korrawe Karunratanakul
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  
  """Basic sequence-to-sequence model with dynamic RNN support."""
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import abc
  
  import tensorflow as tf
  
  from tensorflow.python.layers import core as layers_core
  from tensorflow.python.ops import array_ops
  
  from . import model_helper
  from . import ion_beam_search_decoder
  from . import input_config as aa_conf
  from .utils import iterator_utils
  from .utils import misc_utils as utils
  
  utils.check_tensorflow_version()
  
  __all__ = ["BaseModel", "Model"]
  
  
  class BaseModel(object):
    """Sequence-to-sequence base class.
    """
  
    def __init__(self,
                 hparams,
                 mode,
                 iterator,
                 #source_vocab_table,
                 aa_weight_table, ########## <<<<<<<<<<<<<<<
                 target_vocab_table,
                 reverse_target_vocab_table=None,
                 scope=None,
                 extra_args=None):
      """Create the model.
  
      Args:
        hparams: Hyperparameter configurations.
        mode: TRAIN | EVAL | INFER
        iterator: Dataset Iterator that feeds data.
        source_vocab_table: Lookup table mapping source words to ids.
        target_vocab_table: Lookup table mapping target words to ids.
        reverse_target_vocab_table: Lookup table mapping ids to target words. Only
          required in INFER mode. Defaults to None.
        scope: scope of the model.
        extra_args: model_helper.ExtraArgs, for passing customizable functions.
  
      """
      assert isinstance(iterator, iterator_utils.BatchedInput)
      self.iterator = iterator
      self.mode = mode
      self.tgt_vocab_table = target_vocab_table
      self.aa_weight_table = aa_weight_table
  
      self.tgt_vocab_size = hparams.tgt_vocab_size
      self.num_gpus = hparams.num_gpus
      self.time_major = hparams.time_major
  
      # extra_args: to make it flexible for adding external customizable code
      self.single_cell_fn = None
      if extra_args:
        self.single_cell_fn = extra_args.single_cell_fn
  
      # Set num layers
      self.num_decoder_layers = hparams.num_decoder_layers
      assert self.num_decoder_layers
  
      # Set num residual layers
      if hasattr(hparams, "num_residual_layers"):  # compatible common_test_utils
        self.num_decoder_residual_layers = hparams.num_residual_layers
      else:
        self.num_decoder_residual_layers = hparams.num_decoder_residual_layers
  
      # Initializer
      initializer = model_helper.get_initializer(
          hparams.init_op, hparams.random_seed, hparams.init_weight)
      tf.get_variable_scope().set_initializer(initializer)
  
      # Embeddings
      self.init_embeddings(hparams, scope)
      self.batch_size = tf.size(self.iterator.source_spectrum_length) 
      
      # Projection
      with tf.variable_scope(scope or "build_network"):
        
        with tf.variable_scope("posible_suffix_table"), tf.device("/gpu:0"):
          self.suffix_dp_table = tf.get_variable("suffix_dp_table", [aa_conf.max_dp_array_size], initializer = tf.zeros_initializer, dtype=tf.bool)
          
        with tf.variable_scope("decoder/output_projection"):
          self.output_layer = layers_core.Dense(
            hparams.tgt_vocab_size, use_bias=False, name="output_projection")
  
        with tf.variable_scope("decoder/cnn_input_candidate"):
          with tf.device("/gpu:0"):
            self.input_cnn_1 = tf.layers.Conv1D(
              filters=32,
              kernel_size=1,
              padding="same",
              data_format='channels_last',
              activation=tf.nn.relu,
              name="cnn_input_candidate_1")
            self.input_cnn_2 = tf.layers.Conv1D(
              filters=64,
              kernel_size=1,
              padding="same",
              data_format='channels_last',
              activation=tf.nn.relu,
              name="cnn_input_candidate_2")
            self.input_dense_1 = layers_core.Dense(512, activation=tf.nn.relu, name="input_dense_1")
            self.input_dense_2 = layers_core.Dense(512, activation=tf.nn.relu, name="input_dense_2")
      
  
      ## Train graph
      res = self.build_graph(hparams, scope=scope)
      
      # edited until here
      
  
      if self.mode == tf.contrib.learn.ModeKeys.TRAIN:
        self.train_loss = res[1]
        self.word_count = tf.reduce_sum(self.iterator.target_sequence_length)
      elif self.mode == tf.contrib.learn.ModeKeys.EVAL:
        self.eval_loss = res[1]
        self.eval_logits = res[0] ##############
        self.target_output = self.iterator.target_output ##### 
      elif self.mode == tf.contrib.learn.ModeKeys.INFER:
        self.infer_logits, _, self.final_context_state, self.sample_id, self.step_log_probs = res
        self.sample_words = reverse_target_vocab_table.lookup(
            tf.to_int64(self.sample_id))
  
      if self.mode != tf.contrib.learn.ModeKeys.INFER:
        ## Count the number of predicted words for compute ppl.
        self.predict_count = tf.reduce_sum(
            self.iterator.target_sequence_length)
  
      self.global_step = tf.Variable(0, trainable=False)
      params = tf.trainable_variables()
  
      # Gradients and SGD update operation for training the model.
      # Arrage for the embedding vars to appear at the beginning.
      if self.mode == tf.contrib.learn.ModeKeys.TRAIN:
        self.learning_rate = tf.constant(hparams.learning_rate)
        # warm-up
        self.learning_rate = self._get_learning_rate_warmup(hparams)
        # decay
        self.learning_rate = self._get_learning_rate_decay(hparams)
  
        # Optimizer
        if hparams.optimizer == "sgd":
          opt = tf.train.GradientDescentOptimizer(self.learning_rate)
          tf.summary.scalar("lr", self.learning_rate)
        elif hparams.optimizer == "adam":
          opt = tf.train.AdamOptimizer(self.learning_rate)
  
        # Gradients
        gradients = tf.gradients(
            self.train_loss,
            params,
            colocate_gradients_with_ops=hparams.colocate_gradients_with_ops)
  
        clipped_grads, grad_norm_summary, grad_norm = model_helper.gradient_clip(
            gradients, max_gradient_norm=hparams.max_gradient_norm)
        self.grad_norm = grad_norm
  
        self.update = opt.apply_gradients(
            zip(clipped_grads, params), global_step=self.global_step)
  
        '''
        # Add histograms for variables, gradients and gradient norms.
        for gradient, variable in zip(gradients, params):
          if isinstance(gradient, tf.IndexedSlices):
            grad_values = gradient.values
          else:
            grad_values = gradient
          tf.summary.histogram(variable.name, variable)
          tf.summary.histogram(variable.name + "/gradients", grad_values)
        
        # Add histograms for variables, gradients and gradient norms. 
        for gradient, variable in zip(clipped_grads, params):
          if isinstance(gradient, tf.IndexedSlices):
            grad_values = gradient.values
          else:
            grad_values = gradient
          tf.summary.histogram(variable.name + "/gradient_norm", grad_values)
        '''
        # Summary
        self.train_summary = tf.summary.merge([
            tf.summary.scalar("lr", self.learning_rate),
            tf.summary.scalar("train_loss", self.train_loss),
        ] + grad_norm_summary)
  
      if self.mode == tf.contrib.learn.ModeKeys.INFER:
        self.infer_summary = self._get_infer_summary(hparams)
  
      # Saver
      self.saver = tf.train.Saver(
          tf.global_variables(), max_to_keep=hparams.num_keep_ckpts)
  
      # Print trainable variables
      utils.print_out("# Trainable variables")
      for param in params:
        utils.print_out("  %s, %s, %s" % (param.name, str(param.get_shape()),
                                          param.op.device))
      
  
    def _get_learning_rate_warmup(self, hparams):
      """Get learning rate warmup."""
      warmup_steps = hparams.warmup_steps
      warmup_scheme = hparams.warmup_scheme
      utils.print_out("  learning_rate=%g, warmup_steps=%d, warmup_scheme=%s" %
                      (hparams.learning_rate, warmup_steps, warmup_scheme))
  
      if warmup_scheme == "t2t":
        warmup_factor = tf.exp(tf.log(0.01) / warmup_steps)
        inv_decay = warmup_factor**(
            tf.to_float(warmup_steps - self.global_step))
      else:
        raise ValueError("Unknown warmup scheme %s" % warmup_scheme)
  
      return tf.cond(
          self.global_step < hparams.warmup_steps,
          lambda: inv_decay * self.learning_rate,
          lambda: self.learning_rate,
          name="learning_rate_warump_cond")
  
    def _get_learning_rate_decay(self, hparams):
      """Get learning rate decay."""
      if hparams.decay_scheme in ["luong5", "luong10", "luong234"]:
        decay_factor = 0.5
        if hparams.decay_scheme == "luong5":
          start_decay_step = int(hparams.num_train_steps / 2)
          decay_times = 5
        elif hparams.decay_scheme == "luong10":
          start_decay_step = int(hparams.num_train_steps / 2)
          decay_times = 10
        elif hparams.decay_scheme == "luong234":
          start_decay_step = int(hparams.num_train_steps * 2 / 3)
          decay_times = 4
        remain_steps = hparams.num_train_steps - start_decay_step
        decay_steps = int(remain_steps / decay_times)
      elif not hparams.decay_scheme:  # no decay
        start_decay_step = hparams.num_train_steps
        decay_steps = 0
        decay_factor = 1.0
      elif hparams.decay_scheme:
        raise ValueError("Unknown decay scheme %s" % hparams.decay_scheme)
      utils.print_out("  decay_scheme=%s, start_decay_step=%d, decay_steps %d, "
                      "decay_factor %g" % (hparams.decay_scheme,
                                           start_decay_step,
                                           decay_steps,
                                           decay_factor))
  
      return tf.cond(
          self.global_step < start_decay_step,
          lambda: self.learning_rate,
          lambda: tf.train.exponential_decay(
              self.learning_rate,
              (self.global_step - start_decay_step),
              decay_steps, decay_factor, staircase=True),
          name="learning_rate_decay_cond")
  
    def init_embeddings(self, hparams, scope):
      """Init embeddings."""
      self.embedding_decoder = (
          model_helper.create_emb_for_decoder(
              tgt_vocab_size=self.tgt_vocab_size,
              tgt_embed_size=hparams.embed_size,
              num_partitions=hparams.num_embeddings_partitions,
              tgt_vocab_file=hparams.tgt_vocab_file,
              tgt_embed_file=hparams.tgt_embed_file,
              scope=scope,))
  
    def train(self, sess):
      assert self.mode == tf.contrib.learn.ModeKeys.TRAIN
      return sess.run([self.update,
                       self.train_loss,
                       self.predict_count,
                       self.train_summary,
                       self.global_step,
                       self.word_count,
                       self.batch_size,
                       self.grad_norm,
                       self.learning_rate])
  
    def eval(self, sess):
      assert self.mode == tf.contrib.learn.ModeKeys.EVAL
      return sess.run([self.eval_loss,
                       self.predict_count,
                       self.batch_size])
    
    def sample_eval(self, sess):
      assert self.mode == tf.contrib.learn.ModeKeys.EVAL
      return sess.run([self.eval_logits,
                       self.target_output,
                       self.eval_loss,
                       self.predict_count,
                       self.batch_size])
  
    def build_graph(self, hparams, scope=None):
      """Subclass must implement this method.
  
      Creates a sequence-to-sequence model with dynamic RNN decoder API.
      Args:
        hparams: Hyperparameter configurations.
        scope: VariableScope for the created subgraph; default "dynamic_seq2seq".
  
      Returns:
        A tuple of the form (logits, loss, final_context_state),
        where:
          logits: float32 Tensor [batch_size x num_decoder_symbols].
          loss: the total loss / batch_size.
          final_context_state: The final state of decoder RNN.
  
      Raises:
        ValueError: if encoder_type differs from mono and bi, or
          attention_option is not (luong | scaled_luong |
          bahdanau | normed_bahdanau).
      """
      utils.print_out("# creating %s graph ..." % self.mode)
      dtype = tf.float32
  
      with tf.variable_scope(scope or "dynamic_seq2seq", dtype=dtype):
        # Encoder
        encoder_outputs, encoder_state = self._build_encoder(hparams)
  
        ## Decoder
        logits, sample_id, final_context_state, step_log_probs = self._build_decoder(
            encoder_outputs, encoder_state, hparams)
  
        ## Loss
        if self.mode != tf.contrib.learn.ModeKeys.INFER:
          with tf.device(model_helper.get_device_str(0, #self.num_encoder_layers - 1,
                                                     self.num_gpus)):
            loss = self._compute_loss(logits)
        else:
          loss = None
  
        return logits, loss, final_context_state, sample_id, step_log_probs
  
    @abc.abstractmethod
    def _build_encoder(self, hparams):
      """Subclass must implement this.
  
      Build and run an RNN encoder.
  
      Args:
        hparams: Hyperparameters configurations.
  
      Returns:
        A tuple of encoder_outputs and encoder_state.
      """
      pass
  
    def _build_encoder_cell(self, hparams, num_layers, num_residual_layers,
                            base_gpu=0):
      """Build a multi-layer RNN cell that can be used by encoder."""
  
      return model_helper.create_rnn_cell(
          unit_type=hparams.unit_type,
          num_units=hparams.num_units,
          num_layers=num_layers,
          num_residual_layers=num_residual_layers,
          forget_bias=hparams.forget_bias,
          dropout=hparams.dropout,
          num_gpus=hparams.num_gpus,
          mode=self.mode,
          base_gpu=base_gpu,
          single_cell_fn=self.single_cell_fn)
  
    def _get_infer_maximum_iterations(self, hparams):
      """Maximum decoding steps at inference time."""
      maximum_iterations = hparams.tgt_max_len_infer
      utils.print_out("  decoding maximum_iterations %d" % maximum_iterations)
      return maximum_iterations
    
    def decoder_input_model_fn(self, aa_spec):
      with tf.device("/gpu:0"):
        spec_shape = tf.shape(aa_spec)
        aa_spec = tf.reshape(aa_spec, tf.concat(([spec_shape[0] * spec_shape[1]], spec_shape[2:]), 0))
        interested_regions_num = 8 * aa_conf.vocab_size # + aa_conf.vocab_size
        aa_spec.set_shape([None, interested_regions_num, aa_conf.full_aa_window])
  
        # apply cnn to candidate input, this layer produce output of size [time * batch_size, filters=32, 1]
        decoder_aa_spec_inp = self.input_cnn_1(aa_spec) # [time * batch_size, 160, filters=32]
        decoder_aa_spec_inp = self.input_cnn_2(decoder_aa_spec_inp) # [time * batch_size, 160, filters=32]
  
        decoder_aa_spec_inp = tf.contrib.layers.flatten(decoder_aa_spec_inp) # [time * batch_size, interested_regions_num * filters]
        aa_feature_size = tf.shape(decoder_aa_spec_inp)[1]
  
        # reshape back to [time, batch_size, interested_regions_num * filters]
        decoder_aa_spec_inp = tf.reshape(decoder_aa_spec_inp, [spec_shape[0], spec_shape[1], aa_feature_size])
        decoder_aa_spec_inp.set_shape([None, None, interested_regions_num * 64]) # 32
  
        # apply dense layer -> [time, batch_size, 512]
        decoder_aa_spec_inp = self.input_dense_1(decoder_aa_spec_inp)
        # decoder_aa_spec_inp = tf.layers.dropout(inputs=decoder_aa_spec_inp, rate=0.2)
        decoder_aa_spec_inp = self.input_dense_2(decoder_aa_spec_inp)
        # decoder_aa_spec_inp = tf.layers.dropout(inputs=decoder_aa_spec_inp, rate=0.2)
        # decoder_aa_spec_inp = self.input_dense_3(decoder_aa_spec_inp)
      return decoder_aa_spec_inp
    
    def _build_decoder(self, encoder_outputs, encoder_state, hparams):
      """Build and run a RNN decoder with a final projection layer.
  
      Args:
        encoder_outputs: The outputs of encoder for every time step.
        encoder_state: The final state of the encoder.
        hparams: The Hyperparameters configurations.
  
      Returns:
        A tuple of final logits and final decoder state:
          logits: size [time, batch_size, vocab_size] when time_major=True.
      """
      tgt_sos_id = tf.cast(self.tgt_vocab_table.lookup(tf.constant(hparams.sos)),
                           tf.int32)
      tgt_eos_id = tf.cast(self.tgt_vocab_table.lookup(tf.constant(hparams.eos)),
                           tf.int32)
      iterator = self.iterator
  
      # maximum_iteration: The maximum decoding steps.
      maximum_iterations = self._get_infer_maximum_iterations(hparams)
  
      ## Decoder.
      with tf.variable_scope("decoder") as decoder_scope:
        cell, decoder_initial_state = self._build_decoder_cell(
            hparams, encoder_outputs, encoder_state,
            iterator.source_spectrum_length)
        
        # self.init_state_dec = decoder_initial_state
        #decoder_initial_state = tf.transpose(decoder_initial_state)
        ## Train or eval
        if self.mode != tf.contrib.learn.ModeKeys.INFER:
          # decoder_emp_inp: [max_time, batch_size, num_units]
          target_input = iterator.target_input
          aa_spec = iterator.aa_spectrum
          if self.time_major:
            target_input = tf.transpose(target_input)
            aa_spec = tf.transpose(aa_spec, perm=[1,0,2,3])
          decoder_emb_inp = tf.nn.embedding_lookup(
              self.embedding_decoder, target_input)
          
          # reshape input from [time, batch_size, 160, full_input_window] to [time * batch_size, 160, full_input_window]
          spec_shape = tf.shape(aa_spec)
          aa_spec = tf.reshape(aa_spec, tf.concat(([spec_shape[0] * spec_shape[1]], spec_shape[2:]), 0))
          interested_regions_num = 8 * aa_conf.vocab_size # + aa_conf.vocab_size
          aa_spec.set_shape([None, interested_regions_num, aa_conf.full_aa_window])
          
          # apply cnn to candidate input, this layer produce output of size
          decoder_aa_spec_inp = self.input_cnn_1(aa_spec) # [time * batch_size, filters=32, full_window_size]
          decoder_aa_spec_inp = self.input_cnn_2(decoder_aa_spec_inp) # [time * batch_size, filters=32, full_window_size]
    
          decoder_aa_spec_inp = tf.contrib.layers.flatten(decoder_aa_spec_inp) # [time * batch_size, filters * full_window_size]
          aa_feature_size = tf.shape(decoder_aa_spec_inp)[1]
          
          # reshape back to [time, batch_size, filters=32]
          decoder_aa_spec_inp = tf.reshape(decoder_aa_spec_inp, [spec_shape[0], spec_shape[1], aa_feature_size])
          decoder_aa_spec_inp.set_shape([None, None, interested_regions_num * 64]) # 32 * aa_conf.full_aa_window])
          
          # apply dense layer -> [time, batch_size, 512]
          decoder_aa_spec_inp = self.input_dense_1(decoder_aa_spec_inp)
          decoder_aa_spec_inp = self.input_dense_2(decoder_aa_spec_inp)
  
          combine_inp = tf.concat([decoder_emb_inp,decoder_aa_spec_inp], 2)
          
          # Helper
          helper = tf.contrib.seq2seq.TrainingHelper(
              # decoder_emb_inp, iterator.target_sequence_length,
              combine_inp, iterator.target_sequence_length,
              time_major=self.time_major)
  
          # Decoder
          my_decoder = tf.contrib.seq2seq.BasicDecoder(
              cell,
              helper,
              decoder_initial_state,)
  
          # Dynamic decoding
          outputs, final_context_state, _ = tf.contrib.seq2seq.dynamic_decode(
              my_decoder,
              output_time_major=self.time_major,
              swap_memory=True,
              scope=decoder_scope)
  
          sample_id = outputs.sample_id
          step_log_probs = tf.no_op()
          logits = self.output_layer(outputs.rnn_output)
  
        ## Inference
        else:
          beam_width = hparams.beam_width
          length_penalty_weight = hparams.length_penalty_weight
          start_tokens = tf.fill([self.batch_size], tgt_sos_id)
          end_token = tgt_eos_id
  
          if beam_width > 0:
            peptide_mass = iterator.peptide_mass
            my_decoder = ion_beam_search_decoder.IonBeamSearchDecoder(
                cell=cell,
                embedding=self.embedding_decoder,
                start_tokens=start_tokens,
                end_token=end_token,
                initial_state=decoder_initial_state,
                beam_width=beam_width,
                source_spectrum=iterator.source,
                peptide_mass=peptide_mass,
                suffix_dp_table=self.suffix_dp_table,
                aa_weight_table=self.aa_weight_table,
                decoder_input_model_fn=self.decoder_input_model_fn,
                output_layer=self.output_layer,
                length_penalty_weight=length_penalty_weight)
          else:
            # Helper
            sampling_temperature = hparams.sampling_temperature
            if sampling_temperature > 0.0:
              helper = tf.contrib.seq2seq.SampleEmbeddingHelper(
                  self.embedding_decoder, start_tokens, end_token,
                  softmax_temperature=sampling_temperature,
                  seed=hparams.random_seed)
            else:
              helper = tf.contrib.seq2seq.GreedyEmbeddingHelper(
                  self.embedding_decoder, start_tokens, end_token)
  
            # Decoder
            my_decoder = tf.contrib.seq2seq.BasicDecoder(
                cell,
                helper,
                decoder_initial_state,
                output_layer=self.output_layer  # applied per timestep
            )
  
          # Dynamic decoding
          outputs, final_context_state, _ = tf.contrib.seq2seq.dynamic_decode(
              my_decoder,
              maximum_iterations=maximum_iterations,
              output_time_major=self.time_major,
              swap_memory=True,
              scope=decoder_scope)
  
          if beam_width > 0:
            logits = tf.no_op()
            step_log_probs = outputs.step_log_probs
            sample_id = outputs.predicted_ids
          else:
            logits = outputs.rnn_output
            step_log_probs = tf.no_op()
            sample_id = outputs.sample_id
  
      return logits, sample_id, final_context_state, step_log_probs
  
    def get_max_time(self, tensor):
      time_axis = 0 if self.time_major else 1
      return tensor.shape[time_axis].value or tf.shape(tensor)[time_axis]
  
    @abc.abstractmethod
    def _build_decoder_cell(self, hparams, encoder_outputs, encoder_state,
                            source_sequence_length):
      """Subclass must implement this.
  
      Args:
        hparams: Hyperparameters configurations.
        encoder_outputs: The outputs of encoder for every time step.
        encoder_state: The final state of the encoder.
        source_sequence_length: sequence length of encoder_outputs.
  
      Returns:
        A tuple of a multi-layer RNN cell used by decoder
          and the intial state of the decoder RNN.
      """
      pass
  
  #   def _compute_loss(self, logits):
  #     """Compute optimization loss."""
  #     target_output = self.iterator.target_output
  #     if self.time_major:
  #       target_output = tf.transpose(target_output)
  #     max_time = self.get_max_time(target_output)
  #     crossent = tf.nn.sparse_softmax_cross_entropy_with_logits(
  #         labels=target_output, logits=logits)
  #     # crossent = tf.Print(crossent, [tf.shape(crossent), tf.shape(logits), tf.shape(target_output)])
  #     target_weights = tf.sequence_mask(
  #         self.iterator.target_sequence_length, max_time, dtype=logits.dtype)
  #     if self.time_major:
  #       target_weights = tf.transpose(target_weights)
  
  #     loss = tf.reduce_sum(
  #         crossent * target_weights) / tf.to_float(self.batch_size)
  #     return loss
    
    def _compute_loss(self, logits): # focal_loss(prediction_tensor, target_tensor, alpha=0.25, gamma=2):
      """Compute optimization loss.
      Focal loss
      Multi-labels Focal loss formula:
        FL = -alpha * (z-p)^gamma * log(p) -(1-alpha) * p^gamma * log(1-p)
          ,which alpha = 1.0, gamma = 1.0, p = sigmoid(x), z = target_tensor.
      """
      gamma = 1.0
      alpha = 0.25
      target_output = self.iterator.target_output
      if self.time_major:
        target_output = tf.transpose(target_output)
      max_time = self.get_max_time(target_output)
      
      if self.mode == tf.contrib.learn.ModeKeys.TRAIN:
        onehot_labels = tf.one_hot(target_output, self.tgt_vocab_size)
  
        ce = tf.nn.sigmoid_cross_entropy_with_logits(labels=onehot_labels, logits=logits)
        predictions = tf.sigmoid(logits)
        predictions_pt = tf.where(tf.equal(onehot_labels, 1), predictions, 1.-predictions)
        # add small value to avoid 0
        alpha_t = tf.scalar_mul(alpha, tf.ones_like(onehot_labels, dtype=tf.float32))
        alpha_t = tf.where(tf.equal(onehot_labels, 1.0), alpha_t, 1-alpha_t)
        crossent = ce * tf.pow(1-predictions_pt, gamma) * alpha_t
        crossent = tf.reduce_sum(crossent, -1)
      else:
        crossent = tf.nn.sparse_softmax_cross_entropy_with_logits(
            labels=target_output, logits=logits)
      
      target_weights = tf.sequence_mask(
          self.iterator.target_sequence_length, max_time, dtype=logits.dtype)
      if self.time_major:
        target_weights = tf.transpose(target_weights)
      
      loss = tf.reduce_sum(
        crossent * target_weights) / tf.to_float(self.batch_size)
      return loss
  
    def _get_infer_summary(self, hparams):
      return tf.no_op()
  
    def infer(self, sess, step):
      assert self.mode == tf.contrib.learn.ModeKeys.INFER
      out = sess.run([
          self.infer_logits, self.infer_summary, self.sample_id, self.sample_words, self.step_log_probs
      ])
      return out
  
    def decode(self, sess, step): # , summary_writer, step):
      """Decode a batch.
  
      Args:
        sess: tensorflow session to use.
  
      Returns:
        A tuple consiting of outputs, infer_summary.
          outputs: of size [batch_size, time]
      """
      _, infer_summary, _, sample_words, step_log_probs = self.infer(sess, step) #, summary_writer, step)
      
      # make sure outputs is of shape [batch_size, time] or [beam_width,
      # batch_size, time] when using beam search.
      if self.time_major:
        sample_words = sample_words.transpose()
        step_log_probs = step_log_probs.transpose()
      elif sample_words.ndim == 3:  # beam search output in [batch_size,
                                    # time, beam_width] shape.
        sample_words = sample_words.transpose([2, 0, 1])
        step_log_probs = step_log_probs.transpose([2, 0, 1])
      return sample_words, step_log_probs, infer_summary
  
  
  class Model(BaseModel):
    """Sequence-to-sequence dynamic model.
  
    This class implements a multi-layer recurrent neural network as encoder,
    and a multi-layer recurrent neural network decoder.
    """
  
    def _build_encoder(self, hparams):
      """Build an encoder."""
      iterator = self.iterator
  
      encoder_input = iterator.encoder_input
      encoder_input.set_shape([None, aa_conf.max_spec_length_cnn, len(aa_conf.mass_ID_np)+1 +1 ]) # +1 for original, +1 for position
  
      with tf.variable_scope("encoder") as scope:
        dtype = scope.dtype
        # MLP
        device_str = model_helper.get_device_str(0, self.num_gpus)
        
        with tf.device(device_str):
  
          
          # [-1, 50000, vocab_size] -> [-1, 50000, 32]
          conv1 = tf.layers.conv1d(
              inputs=encoder_input,
              filters=32,
              kernel_size=1,
              padding="same",
              activation=tf.nn.relu)
          conv2 = tf.layers.conv1d(
              inputs=conv1,
              filters=64,
              kernel_size=1,
              padding="same",
              activation=tf.nn.relu) # [-1, 50000, 64]
          conv3 = tf.layers.conv1d(
              inputs=conv2,
              filters=2,
              kernel_size=1,
              padding="same",
              activation=tf.nn.relu) # [-1, 50000, 2]
          
          flatten_enc = tf.contrib.layers.flatten(conv3) # [batch_size, 50000 * 5] -> 100000
          dense1 = tf.layers.dense(inputs=flatten_enc, units=512, activation=tf.nn.relu, name="dense1")
          dense1_drop = tf.layers.dropout(inputs=dense1, rate=0.4)
          dense2 = tf.layers.dense(inputs=dense1_drop, units=512, activation=tf.nn.relu, name="dense2")
          dense3 = tf.layers.dense(inputs=dense2, units=1024, activation=tf.nn.relu, name="dense3")
  
        h, c = tf.split(dense3, 2, axis=1)
        one_state = tf.nn.rnn_cell.LSTMStateTuple(c, h)
        
        init_state = []
        for i in range(self.num_decoder_layers):
          init_state.append(one_state)
                
        if len(init_state) == 1:  # Single layer.
          return (0), one_state
      
        return (0), tuple(init_state)
  
    def _build_bidirectional_rnn(self, inputs, sequence_length,
                                 dtype, hparams,
                                 num_bi_layers,
                                 num_bi_residual_layers,
                                 base_gpu=0):
      """Create and call biddirectional RNN cells.
  
      Args:
        num_residual_layers: Number of residual layers from top to bottom. For
          example, if `num_bi_layers=4` and `num_residual_layers=2`, the last 2 RNN
          layers in each RNN cell will be wrapped with `ResidualWrapper`.
        base_gpu: The gpu device id to use for the first forward RNN layer. The
          i-th forward RNN layer will use `(base_gpu + i) % num_gpus` as its
          device id. The `base_gpu` for backward RNN cell is `(base_gpu +
          num_bi_layers)`.
  
      Returns:
        The concatenated bidirectional output and the bidirectional RNN cell"s
        state.
      """
      # Construct forward and backward cells
      fw_cell = self._build_encoder_cell(hparams,
                                         num_bi_layers,
                                         num_bi_residual_layers,
                                         base_gpu=base_gpu)
      bw_cell = self._build_encoder_cell(hparams,
                                         num_bi_layers,
                                         num_bi_residual_layers,
                                         base_gpu=(base_gpu + num_bi_layers))
  
      bi_outputs, bi_state = tf.nn.bidirectional_dynamic_rnn(
          fw_cell,
          bw_cell,
          inputs,
          dtype=dtype,
          sequence_length=sequence_length,
          time_major=self.time_major,
          swap_memory=True)
  
      return tf.concat(bi_outputs, -1), bi_state
  
    def _build_decoder_cell(self, hparams, encoder_outputs, encoder_state,
                            source_sequence_length):
      """Build an RNN cell that can be used by decoder."""
      # We only make use of encoder_outputs in attention-based models
      if hparams.attention:
        raise ValueError("BasicModel doesn't support attention.")
  
      cell = model_helper.create_rnn_cell(
          unit_type=hparams.unit_type,
          num_units=hparams.num_units,
          num_layers=self.num_decoder_layers,
          num_residual_layers=self.num_decoder_residual_layers,
          forget_bias=hparams.forget_bias,
          dropout=hparams.dropout,
          num_gpus=self.num_gpus,
          mode=self.mode,
          single_cell_fn=self.single_cell_fn)
  
      # For beam search, we need to replicate encoder infos beam_width times
      if self.mode == tf.contrib.learn.ModeKeys.INFER and hparams.beam_width > 0:
        decoder_initial_state = tf.contrib.seq2seq.tile_batch(
            encoder_state, multiplier=hparams.beam_width)
      else:
        decoder_initial_state = encoder_state
      
      return cell, decoder_initial_state
  