text: # Copyright 2016 The TensorFlow Authors. All Rights Reserved.
  # Modifications Copyright 2017 Abigail See
  # Modifications Copyright 2018 Arman Cohan
  #
  # 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.
  # ==============================================================================
  from tensorflow.contrib.slim.python.slim import learning
  import sys
  
  """This file contains code to build and run the tensorflow graph for the sequence-to-sequence model"""
  
  import os
  import time
  import numpy as np
  import tensorflow as tf
  from util import load_embeddings
  from tensorflow.contrib.tensorboard.plugins import projector
  from tensorflow.python.ops import array_ops
  from six.moves import xrange
  from tensorflow.python import debug as tf_debug
  FLAGS = tf.app.flags.FLAGS
  
  class SummarizationModel(object):
    """A class to represent a sequence-to-sequence model for text summarization. Supports both baseline mode, pointer-generator mode, and coverage"""
  
    def __init__(self, hps, vocab, num_gpus):
      self._hps = hps
      self._cur_gpu = 0
      self._vocab = vocab
      self._num_gpus = num_gpus
      if FLAGS.new_attention and FLAGS.hier:
        print('using linear attention mechanism for considering sections')
        from attention_decoder_new import attention_decoder
      else:
        print('using hierarchical attention mechanism for considering sections')
        from attention_decoder import attention_decoder
      self.attn_decoder = attention_decoder
  
    def _add_placeholders(self):
      """Add placeholders to the graph. These are entry points for any input data."""
      hps = self._hps
  
      # encoder part
      self._enc_batch = tf.placeholder(tf.int32, [hps.batch_size, None], name='enc_batch')
      self._enc_padding_mask = tf.placeholder(tf.float32, [hps.batch_size, None], name='enc_padding_mask')
      if self._hps.hier:
        self._enc_batch_sections = tf.placeholder(tf.int32, [hps.batch_size, hps.num_sections, None], name='enc_batch_sections')
        self._doc_sec_lens = tf.placeholder(tf.int32, [hps.batch_size]) # length of doc in num sections
        self._batch_sections_len = tf.placeholder(tf.int32, [hps.batch_size, hps.num_sections])
        self._enc_section_padding_mask = tf.placeholder(tf.int32, [hps.batch_size, hps.num_sections, None], name='enc_section_padding_mask')
      self._enc_lens = tf.placeholder(tf.int32, [hps.batch_size], name='enc_lens')
      if FLAGS.pointer_gen:
        self._enc_batch_extend_vocab = tf.placeholder(tf.int32, [hps.batch_size, None], name='enc_batch_extend_vocab')
        self._max_art_oovs = tf.placeholder(tf.int32, [], name='max_art_oovs')
  
      # decoder part
      self._dec_batch = tf.placeholder(tf.int32, [hps.batch_size, hps.max_dec_steps], name='dec_batch')
      self._target_batch = tf.placeholder(tf.int32, [hps.batch_size, hps.max_dec_steps], name='target_batch')
      self._dec_padding_mask = tf.placeholder(tf.float32, [hps.batch_size, hps.max_dec_steps], name='padding_mask')
  
      if hps.mode=="decode" and hps.coverage:
        self.prev_coverage = tf.placeholder(tf.float32, [hps.batch_size, None], name='prev_coverage')
  
  
    def _make_feed_dict(self, batch, just_enc=False):
      """Make a feed dictionary mapping parts of the batch to the appropriate placeholders.
  
      Args:
        batch: Batch object
        just_enc: Boolean. If True, only feed the parts needed for the encoder.
      """
      feed_dict = {}
      if self._hps.hier:
        feed_dict[self._enc_batch_sections] = batch.batch_sections  # shape=[batch-size, num-sections, enc-seq-len]
        feed_dict[self._batch_sections_len] = batch.batch_sections_len # length of sections in the entire batch (num sections). [[400, 400, 400, 400],[...]]
        feed_dict[self._doc_sec_lens] = batch.batch_doc_sec_lens
        feed_dict[self._enc_section_padding_mask] = batch.enc_section_padding_mask
      else:
        feed_dict[self._enc_batch] = batch.enc_batch
        feed_dict[self._enc_padding_mask] = batch.enc_padding_mask
      feed_dict[self._enc_lens] = batch.enc_lens
      if FLAGS.pointer_gen:
        feed_dict[self._enc_batch_extend_vocab] = batch.enc_batch_extend_vocab
        feed_dict[self._max_art_oovs] = batch.max_art_oovs
      if not just_enc:
        feed_dict[self._dec_batch] = batch.dec_batch
        feed_dict[self._target_batch] = batch.target_batch
        feed_dict[self._dec_padding_mask] = batch.dec_padding_mask
      return feed_dict
  
    def _add_encoder(self, encoder_inputs, seq_len):
      """Add a single-layer bidirectional LSTM encoder to the graph.
  
      Args:
        encoder_inputs: A tensor of shape [batch_size, <=max_enc_steps, emb_size].
        seq_len: Lengths of encoder_inputs (before padding). A tensor of shape [batch_size].
  
      Returns:
        encoder_outputs:
          A tensor of shape [batch_size, <=max_enc_steps, 2*hidden_dim]. It's 2*hidden_dim because it's the concatenation of the forwards and backwards states.
        fw_state, bw_state:
          Each are LSTMStateTuples of shape ([batch_size,hidden_dim],[batch_size,hidden_dim])
      """
      with tf.variable_scope('encoder'):
        cell_fw = tf.contrib.rnn.LSTMCell(self._hps.hidden_dim, initializer=self.rand_unif_init, state_is_tuple=True)
        cell_bw = tf.contrib.rnn.LSTMCell(self._hps.hidden_dim, initializer=self.rand_unif_init, state_is_tuple=True)
        (encoder_outputs, (fw_st, bw_st)) = tf.nn.bidirectional_dynamic_rnn(cell_fw, cell_bw, encoder_inputs, dtype=tf.float32, sequence_length=seq_len, swap_memory=True)
        encoder_outputs = tf.concat(axis=2, values=encoder_outputs) # concatenate the forwards and backwards states
      return encoder_outputs, fw_st, bw_st
  
  
    def _reduce_states(self, fw_st, bw_st):
      """Add to the graph a linear layer to reduce the encoder's final FW and BW state
      into a single initial state for the decoder. This is needed
      because the encoder is bidirectional but the decoder is not.
  
      Args:
        fw_st: LSTMStateTuple with hidden_dim units.
        bw_st: LSTMStateTuple with hidden_dim units.
  
      Returns:
        state: LSTMStateTuple with hidden_dim units.
      """
      hidden_dim = self._hps.hidden_dim
      with tf.variable_scope('reduce_final_st'):
  
        # Define weights and biases to reduce the cell and reduce the state
        w_reduce_c = tf.get_variable('w_reduce_c', [hidden_dim * 2, hidden_dim], dtype=tf.float32, initializer=self.trunc_norm_init)
        w_reduce_h = tf.get_variable('w_reduce_h', [hidden_dim * 2, hidden_dim], dtype=tf.float32, initializer=self.trunc_norm_init)
        bias_reduce_c = tf.get_variable('bias_reduce_c', [hidden_dim], dtype=tf.float32, initializer=self.trunc_norm_init)
        bias_reduce_h = tf.get_variable('bias_reduce_h', [hidden_dim], dtype=tf.float32, initializer=self.trunc_norm_init)
  
        # Apply linear layer
        old_c = tf.concat(axis=1, values=[fw_st.c, bw_st.c]) # Concatenation of fw and bw cell
        old_h = tf.concat(axis=1, values=[fw_st.h, bw_st.h]) # Concatenation of fw and bw state
        new_c = tf.nn.relu(tf.matmul(old_c, w_reduce_c) + bias_reduce_c) # Get new cell from old cell
        new_h = tf.nn.relu(tf.matmul(old_h, w_reduce_h) + bias_reduce_h) # Get new state from old state
        return tf.contrib.rnn.LSTMStateTuple(new_c, new_h) # Return new cell and state
  
  
    def _calc_final_dist(self, vocab_dists, attn_dists):
      """Calculate the final distribution, for the pointer-generator model
  
      Args:
        vocab_dists: The vocabulary distributions. List length max_dec_steps of
          (batch_size, vsize) arrays. The words are in the order they appear in the
          vocabulary file.
        attn_dists: The attention distributions. List length max_dec_steps of
          (batch_size, attn_len) arrays
  
      Returns:
        final_dists: The final distributions. List length max_dec_steps of
          (batch_size, extended_vsize) arrays. extended_vsize is the vocab + article OOV
      """
  
      with tf.variable_scope('final_distribution'):
        vocab_dists = [p_gen * dist for (p_gen,dist) in zip(self.p_gens, vocab_dists)]
        attn_dists = [(1-p_gen) * dist for (p_gen,dist) in zip(self.p_gens, attn_dists)]
  
        # Extend the vocabulary dist with zeros (for OOV words
        extended_vsize = self._vocab.size() + self._max_art_oovs 
        extra_zeros = tf.zeros((self._hps.batch_size, self._max_art_oovs))
        # list length max_dec_steps of shape (batch_size, extended_vsize)
        vocab_dists_extended = [tf.concat(axis=1, values=[dist, extra_zeros])
                                for dist in vocab_dists]
  
  
        # Project the values in the attention distributions onto the appropriate
        # entries in the final distributions
        batch_nums = tf.range(0, limit=self._hps.batch_size) # shape (batch_size)
        batch_nums = tf.expand_dims(batch_nums, 1) # shape (batch_size, 1)
  
        attn_len = tf.shape(self._enc_batch_extend_vocab)[1] 
        batch_nums = tf.tile(batch_nums, [1, attn_len]) # shape (batch_size, attn_len)
        indices = tf.stack( (batch_nums, self._enc_batch_extend_vocab), axis=2) 
        shape = [self._hps.batch_size, extended_vsize]
        # indices has shape [batch_size, extended_vsize, 2]
        # sample slice: [[[0, 701], ... ], [[1, 529], ...], [[2, 728], ...]]
        # scatter the distribution among corresponding batches and vocabulary index
        
  
        # list length max_dec_steps (batch_size, extended_vsize)
        attn_dists_projected = [tf.scatter_nd(indices, copy_dist, shape)
                                for copy_dist in attn_dists]
  
        # Add a very small epsilon
        def _add_epsilon(dist, epsilon=sys.float_info.epsilon):
          epsilon_mask = tf.ones_like(dist) * epsilon
          return dist + epsilon_mask
        # Add the vocab distributions and the copy distributions together to get
        # the final distributions
        final_dists = [vocab_dist + copy_dist
                       for (vocab_dist, copy_dist)
                       in zip(vocab_dists_extended, attn_dists_projected)]
        final_dists = [_add_epsilon(dist) for dist in final_dists]
  
        return final_dists
  
    def _add_emb_vis(self, embedding_var):
      """Do setup so that we can view word embedding visualization in Tensorboard, as described here:
      https://www.tensorflow.org/get_started/embedding_viz
      Make the vocab metadata file, then make the projector config file pointing to it."""
      train_dir = os.path.join(FLAGS.log_root, "train")
      vocab_metadata_path = os.path.join(train_dir, "vocab_metadata.tsv")
      self._vocab.write_metadata(vocab_metadata_path) # write metadata file
      summary_writer = tf.summary.FileWriter(train_dir)
      config = projector.ProjectorConfig()
      embedding = config.embeddings.add()
      embedding.tensor_name = embedding_var.name
      embedding.metadata_path = vocab_metadata_path
      chkpt_dir = tf.train.latest_checkpoint(train_dir)
      print('chkpt_dir for embeddings: ', chkpt_dir)
      if chkpt_dir:
        config.model_checkpoint_path = chkpt_dir
      else:
        chkpt_dir = train_dir
      projector.visualize_embeddings(summary_writer, config)
      
    def _next_device(self):
      """Round robin the gpu device. (Reserve last gpu for expensive op)."""
      if self._num_gpus == 0:
        return ''
      dev = '/gpu:%d' % self._cur_gpu
      if self._num_gpus > 1:
        self._cur_gpu = (self._cur_gpu + 1) % (self._num_gpus-1)
      return dev
    
    def _get_gpu(self, gpu_id):
      if self._num_gpus <= 0 or gpu_id >= self._num_gpus:
        return ''
      return '/gpu:%d' % gpu_id
  
    def _add_seq2seq(self):
      """Add the whole sequence-to-sequence model to the graph."""
      hps = self._hps
      vsize = self._vocab.size() # size of the vocabulary
      
      with tf.variable_scope('seq2seq'):
        # Some initializers
        self.rand_unif_init = tf.random_uniform_initializer(-hps.rand_unif_init_mag, hps.rand_unif_init_mag, seed=123)
        self.trunc_norm_init = tf.truncated_normal_initializer(stddev=hps.trunc_norm_init_std)
  
  
        with tf.variable_scope('embedding'):
          if hps.pretrained_embeddings:
            word2vec = load_embeddings(hps.embeddings_path, self._vocab.word2id, hps.rand_unif_init_mag)
            self.embedding = tf.get_variable('embedding', [vsize, hps.emb_dim],
                                      dtype=tf.float32, initializer=tf.constant_initializer(word2vec))
            # self.assign_embedding = tf.assign(self.embedding, word2vec)
          else:
            self.embedding = tf.get_variable('embedding', [vsize, hps.emb_dim],
                                        dtype=tf.float32, initializer=self.trunc_norm_init)
          if hps.mode=="train": self._add_emb_vis(self.embedding) # add to tensorboard
  
          # tensor with shape (batch_size, max_enc_steps, emb_size)
          emb_enc_inputs = tf.nn.embedding_lookup(self.embedding, self._enc_batch)
          if self._hps.hier:
            enc_batch_sections = tf.unstack(self._enc_batch_sections, axis=1)
            sec_emb_enc_inputs = [tf.nn.embedding_lookup(self.embedding, section)
                                  for section in enc_batch_sections]
          # list length max_dec_steps containing shape (batch_size, emb_size)
          emb_dec_inputs = [tf.nn.embedding_lookup(self.embedding, x)
                            for x in tf.unstack(self._dec_batch, axis=1)]
  
  
        # Hierarchical attention model
        if self._hps.hier:
          with tf.variable_scope('encoder'), tf.device(self._next_device()):
            sec_enc_outs = []
            states_fw = []
            states_bw = []
            states = []
  
            # level 1, encode words to sections
            with tf.variable_scope("word_level_encoder", reuse=tf.AUTO_REUSE) as scope:
              encoder_outputs_words = []
              cell_fw = tf.contrib.rnn.LSTMCell(self._hps.hidden_dim, initializer=self.rand_unif_init, state_is_tuple=True)
              cell_bw = tf.contrib.rnn.LSTMCell(self._hps.hidden_dim, initializer=self.rand_unif_init, state_is_tuple=True)
              fw_st, bw_st = None, None
              if self._hps.use_do:  # DropOut
                cell_fw = tf.contrib.rnn.DropoutWrapper(cell_fw, output_keep_prob=1.0 - self._hps.do_prob)
                cell_bw = tf.contrib.rnn.DropoutWrapper(cell_bw, output_keep_prob=1.0 - self._hps.do_prob)
              for i in range(self._hps.num_sections):
                encoder_tmp_output, (fw_st, bw_st) = tf.nn.bidirectional_dynamic_rnn(
                  cell_fw, cell_bw, inputs=sec_emb_enc_inputs[i], dtype=tf.float32,
                  sequence_length=self._batch_sections_len[:,i], swap_memory=True, initial_state_bw=bw_st, initial_state_fw=fw_st)
                # concatenate the forwards and backwards states
                encoder_tmp_output = tf.concat(axis=2, values=encoder_tmp_output) #shape=[batch x seq_len x hidden_size]
                   
                encoder_outputs_words.append(encoder_tmp_output)
                # instead of concating the fw and bw states, we use a ff network
                combined_state = self._reduce_states(fw_st, bw_st)
                states.append(combined_state)
                scope.reuse_variables()
                
            # level 2, encode sections to doc
            encoder_outputs_words = tf.stack(encoder_outputs_words, axis=1)  # shape [batch x num_sections x seq_len x hidden_size]
            shapes = encoder_outputs_words.shape
            encoder_outputs_words = tf.reshape(encoder_outputs_words, (shapes[0].value, -1, shapes[-1].value))  #shape=[batch x (seq_len * num_sections) x hidden_size]
  
            doc_sections_h = tf.stack([s.h for s in states], axis=1)   # [batch x num_sections x hidden_size]
            doc_sections_c = tf.stack([s.c for s in states], axis=1)   # [batch x num_sections x hidden_size]
  
            with tf.variable_scope("section_level_encoder"):
              if FLAGS.section_level_encoder == 'RNN':
                cell_fw_1 = tf.contrib.rnn.LSTMCell(self._hps.hidden_dim, initializer=self.rand_unif_init, state_is_tuple=True)
                cell_bw_1 = tf.contrib.rnn.LSTMCell(self._hps.hidden_dim, initializer=self.rand_unif_init, state_is_tuple=True)
                if self._hps.use_do:
                    cell_fw_1 = tf.contrib.rnn.DropoutWrapper(cell_fw_1, output_keep_prob=1.0 - self._hps.do_prob)
                    cell_bw_1 = tf.contrib.rnn.DropoutWrapper(cell_bw_1, output_keep_prob=1.0 - self._hps.do_prob)
                encoder_output_sections, (fw_st_2, bw_st_2) =\
                    tf.nn.bidirectional_dynamic_rnn(cell_fw_1, cell_bw_1, inputs=doc_sections_h, sequence_length=self._doc_sec_lens, dtype=tf.float32, swap_memory=True)
                encoder_output_sections = tf.concat(axis=2, values=encoder_output_sections)
                doc_sections_state = self._reduce_states(fw_st_2, bw_st_2)
              else:
                if FLAGS.section_level_encoder == 'AVG': # average section cells
                  doc_sections_state_h = tf.reduce_mean(doc_sections_h, axis=1)
                  doc_sections_state_c = tf.reduce_mean(doc_sections_c, axis=1)
                elif FLAGS.section_level_encoder == 'FF':  # use a feedforward network to combine section cells
                  doc_sections_state_h = tf.reshape([doc_sections_h.shape[0].eval(), -1])
                  doc_sections_state_h = tf.layers.dense(
                    inputs=doc_sections_state_h,
                    units=self._hps.hidden,
                    activation=tf.nn.relu)              
                  doc_sections_state_c = tf.reshape([doc_sections_c.shape[0].eval(), -1])
                  doc_sections_state_c = tf.layers.dense(
                    inputs=doc_sections_state_c,
                    units=self._hps.hidden,
                    activation=tf.nn.relu)
                else:
                  raise AttributeError('FLAGS.section_level_encoder={} is not a valid option'.format(FLAGS.section_level_encoder))
                doc_sections_state = tf.contrib.rnn.LSTMStateTuple(doc_sections_state_c, doc_sections_state_h)
                encoder_output_sections = doc_sections_h         
        
        elif not self._hps.multi_layer_encoder:
          with tf.variable_scope('encoder'):
            with tf.variable_scope('word_level_encoder'):
              cell_fw = tf.contrib.rnn.LSTMCell(self._hps.hidden_dim, initializer=self.rand_unif_init, state_is_tuple=True)
              cell_bw = tf.contrib.rnn.LSTMCell(self._hps.hidden_dim, initializer=self.rand_unif_init, state_is_tuple=True)
              (encoder_outputs, (fw_st, bw_st)) =\
                tf.nn.bidirectional_dynamic_rnn(cell_fw, cell_bw, inputs=emb_enc_inputs, dtype=tf.float32, sequence_length=self._enc_lens, swap_memory=True)
              # concatenate the forwards and backwards states
              encoder_outputs = tf.concat(axis=2, values=encoder_outputs)
      
        # stack n layers of lstms for encoder
        elif self._hps.multi_layer_encoder:
          # TODO: check
          for layer_i in xrange(self._hps.enc_layers):
            with tf.variable_scope('encoder%d'%layer_i), tf.device(
                self._next_device()):
              cell_fw = tf.contrib.rnn.LSTMCell(self._hps.hidden_dim, initializer=self.rand_unif_init, state_is_tuple=True)
              cell_bw = tf.contrib.rnn.LSTMCell(self._hps.hidden_dim, initializer=self.rand_unif_init, state_is_tuple=True)
              if self._hps.use_do: # add dropout
                  cell_fw = tf.contrib.rnn.DropoutWrapper(cell_fw, output_keep_prob=1.0 - self._hps.do_prob)
                  cell_bw = tf.contrib.rnn.DropoutWrapper(cell_bw, output_keep_prob=1.0 - self._hps.do_prob)
              emb_enc_inputs, (fw_st, bw_st) =\
                tf.nn.bidirectional_dynamic_rnn(cell_fw, cell_bw, inputs=emb_enc_inputs, dtype=tf.float32, sequence_length=self._enc_lens, swap_memory=True)
              emb_enc_inputs = tf.concat(axis=2, values=emb_enc_inputs)
          encoder_outputs = emb_enc_inputs
        
        if self._hps.hier:
          self._enc_sec_states = encoder_output_sections
          self._enc_states = encoder_outputs_words 
        else:
          self._enc_states = encoder_outputs
          self._enc_sec_states = None
         
        # convert the encoder bidirectional hidden state to the decoder state
        # (unidirectional) by an MLP
        if self._hps.hier:
          self._dec_in_state = doc_sections_state
        else:
          with tf.variable_scope('encoder'):
            with tf.variable_scope('word_level_encoder'):
              self._dec_in_state = self._reduce_states(fw_st, bw_st) 
          
        # Add the decoder
  
        with tf.variable_scope('decoder'), tf.device(self._next_device()):
          cell = tf.contrib.rnn.LSTMCell(
            self._hps.hidden_dim,
            state_is_tuple=True,
          initializer=self.rand_unif_init)
      
          # We need to pass in the previous step's coverage vector each time
          prev_coverage = self.prev_coverage\
                           if hps.mode=="decode" and self._hps.coverage \
                           else None 
      
          
          if self._hps.hier:
            decoder_outputs, self._dec_out_state, self.attn_dists, self.p_gens, self.coverage, self.attn_dists_sec =\
              self.attn_decoder(emb_dec_inputs,
                                self._dec_in_state,
                                self._enc_states,
                                cell,
                                self._enc_sec_states,
                                num_words_section=self._batch_sections_len,
                                enc_padding_mask=self._enc_padding_mask,
                                enc_section_padding_mask=self._enc_section_padding_mask,
                                initial_state_attention=(self._hps.mode=="decode"),
                                pointer_gen=self._hps.pointer_gen,
                                use_coverage=self._hps.coverage,
                                prev_coverage=prev_coverage,
                                temperature=self._hps.temperature
                                )
                        
          else:
            decoder_outputs, self._dec_out_state, self.attn_dists, self.p_gens, self.coverage, _ =\
              self.attn_decoder(emb_dec_inputs,
                                self._dec_in_state,
                                self._enc_states,
                                cell,
                                encoder_section_states=None,
                                num_words_section=None,
                                enc_padding_mask=self._enc_padding_mask,
                                initial_state_attention=(self._hps.mode=="decode"),
                                pointer_gen=self._hps.pointer_gen,
                                use_coverage=self._hps.coverage,
                                prev_coverage=prev_coverage,
                                )      
          
  
        # Project decoder output to vocabulary
        with tf.variable_scope('output_projection'), tf.device(self._next_device()):
          if self._hps.output_weight_sharing:
            # share weights of embedding layer with projection
            # self.embedding is in shape [vsize, hps.emb_dim]
            w_proj = tf.get_variable('w_proj', [self._hps.emb_dim, self._hps.hidden_dim],
                                dtype=tf.float32, initializer=self.trunc_norm_init)
            w = tf.tanh(tf.transpose(tf.matmul(self.embedding, w_proj))) # shape = [vsize, hps.hidden_dim]
            
    #         w_t = tf.transpose(w)
            b = tf.get_variable('b', [vsize],
                                dtype=tf.float32, initializer=self.trunc_norm_init)
          else: 
            w = tf.get_variable('w', [self._hps.hidden_dim, vsize],
                                dtype=tf.float32, initializer=self.trunc_norm_init)
    #         w_t = tf.transpose(w)
            b = tf.get_variable('b', [vsize],
                                dtype=tf.float32, initializer=self.trunc_norm_init)
          # vocabulary score at each decoder step
          vocab_scores = []
          for i,output in enumerate(decoder_outputs):
            if i > 0:
              tf.get_variable_scope().reuse_variables()
            vocab_scores.append(tf.nn.xw_plus_b(output, w, b)) # apply the linear layer
  
          # the final vocab distribution for each decoder time step
          # shape of each element is [batch_size, vsize]
          vocab_dists = [tf.nn.softmax(s) for s in vocab_scores] 
  
        
        # pointing / generating
        if FLAGS.pointer_gen:
          final_dists = self._calc_final_dist(vocab_dists, self.attn_dists)
  #         log_dists = [tf.log(dist) for dist in final_dists]
        else:
  #         log_dists = [tf.log(dist) for dist in vocab_dists]
          final_dists = vocab_dists
          
  
        # Calculate Losses:
        
        if self._hps.mode in ['train', 'eval']:
          # Calculate the loss
          with tf.variable_scope('loss'), tf.device(self._next_device()):
            if FLAGS.pointer_gen:
              # Calculate the loss per step
              # This is fiddly; we use tf.gather_nd to pick out the gold target words
              # will be list length max_dec_steps containing shape (batch_size)
              loss_per_step = [] 
              batch_nums = tf.range(0, limit=hps.batch_size) # shape (batch_size)
              for dec_step, dist in enumerate(final_dists):
                # The indices of the target words. shape (batch_size)
                targets = self._target_batch[:,dec_step] 
                indices = tf.stack( (batch_nums, targets), axis=1) # shape (batch_size, 2)
                # shape (batch_size). loss on this step for each batch
                gold_probs = tf.gather_nd(dist, indices)
                losses = -tf.log(gold_probs)
                loss_per_step.append(losses)
  
              # Apply dec_padding_mask mask and get loss
              self._loss = _mask_and_avg(loss_per_step, self._dec_padding_mask)
              
  
            else: # baseline model
              # this applies softmax internally
              self._loss = tf.contrib.seq2seq.sequence_loss(
                tf.stack(vocab_scores, axis=1), self._target_batch, self._dec_padding_mask) # this applies softmax internally
  
            tf.summary.scalar('loss', self._loss)
  
            # Calculate coverage loss from the attention distributions
            if self._hps.coverage:
              with tf.variable_scope('coverage_loss'):
                self._coverage_loss = _coverage_loss(self.attn_dists, self._dec_padding_mask)
                tf.summary.scalar('coverage_loss', self._coverage_loss)
              self._total_loss = self._loss + self._hps.cov_loss_wt * self._coverage_loss
              tf.summary.scalar('total_loss', self._total_loss)
              
          # ---------------------------/
  
  
      if self._hps.mode == "decode":
          assert len(final_dists) == 1 # final_dists is a singleton list containing shape (batch_size, extended_vsize)
          final_dists = final_dists[0]
          topk_probs, self._topk_ids = tf.nn.top_k(final_dists, hps.batch_size*2) # take the k largest probs. note batch_size=beam_size in decode mode
          self._topk_log_probs = tf.log(topk_probs)
  
    def _add_train_op(self):
      """Sets self._train_op, the op to run for training."""
      self._lr_rate = tf.maximum(
        self._hps.min_lr,  # min_lr_rate.
        tf.train.exponential_decay(self._hps.lr, self.global_step, 30000, 0.98))
      
      
      # Take gradients of the trainable variables w.r.t. the loss function to minimize
      loss_to_minimize = self._total_loss if self._hps.coverage else self._loss
      tvars = tf.trainable_variables()
      gradients = tf.gradients(loss_to_minimize, tvars, aggregation_method=tf.AggregationMethod.EXPERIMENTAL_ACCUMULATE_N)
  
      # Clip the gradients
      with tf.device(self._get_gpu(self._num_gpus-1)):
        grads, global_norm = tf.clip_by_global_norm(gradients, self._hps.max_grad_norm)
  
      # Add a summary
      tf.summary.scalar('global_norm', global_norm)
  
      # Apply adagrad optimizer
      if self._hps.optimizer == 'adagrad':
        optimizer = tf.train.AdagradOptimizer(
          self._hps.lr, initial_accumulator_value=self._hps.adagrad_init_acc)
  
      elif self._hps.optimizer == 'adam':    
        # Adam
        optimizer = tf.train.AdamOptimizer()
      
      elif self._hps.optimizer == 'sgd':
        # SGD
        optimizer = tf.train.GradientDescentOptimizer(self._lr_rate)
        tf.summary.scalar('learning rate', self._lr_rate)
      
      else:
        raise Exception('Invalid optimizer: ', self._hps.optimizer)
  
      with tf.device(self._get_gpu(self._num_gpus-1)):
        self._train_op = optimizer.apply_gradients(
          zip(grads, tvars), global_step=self.global_step, name='train_step')
  
  
    def build_graph(self):
      """Add the placeholders, model, global step, train_op and summaries to the graph"""
      tf.logging.info('Building graph...')
      t0 = time.time()
      self._add_placeholders()
      with tf.device("/gpu:0"):
        self._add_seq2seq()
      self.global_step = tf.Variable(0, name='global_step', trainable=False)
      if self._hps.mode == 'train':
        self._add_train_op()
      self._summaries = tf.summary.merge_all()
      t1 = time.time()
      tf.logging.info('Time to build graph: %i seconds', t1 - t0)
      
      print('#'*78,'\nprinting model variables:')
      total_parameters = 0
      for variable in tf.trainable_variables():
          shape = variable.get_shape().as_list()
          variable_parameters = 1
          for dim in shape:
              variable_parameters *= dim
          print('{:}: shape={:}, variable_parameters={:}'.format(
            variable.name, shape, variable_parameters))
          total_parameters += variable_parameters
      print('total model parameters: {:}'.format(total_parameters))
      print('#'*78)
  
    def run_train_step(self, sess, batch):
      feed_dict = self._make_feed_dict(batch)
      to_return = {
          'train_op': self._train_op,
          'summaries': self._summaries,
          'loss': self._loss,
          'global_step': self.global_step,
      }
      if self._hps.coverage:
        to_return['coverage_loss'] = self._coverage_loss
      if FLAGS.debug:
        print('entering debug mode\n\n\n\n\n\n\n\n\n')
        sess = tf_debug.LocalCLIDebugWrapperSession(sess, dump_root=FLAGS.dump_root, ui_type=FLAGS.ui_type)
        sess.add_tensor_filter("has_inf_or_nan", tf_debug.has_inf_or_nan)
  #     try:    
      res = sess.run(to_return, feed_dict)
  
      if not np.isfinite(res['loss']):
        print('loss is nan!!!!!')
        raise Exception("Loss is not finite. Stopping.")
      
  #     except tf.errors.InvalidArgumentError:
  #       import pdb; pdb.set_trace()
      return res
  
    def run_eval_step(self, sess, batch):
      """Runs one evaluation iteration. Returns a dictionary containing summaries, loss, global_step and (optionally) coverage loss."""
      feed_dict = self._make_feed_dict(batch)
      to_return = {
          'summaries': self._summaries,
          'loss': self._loss,
          'global_step': self.global_step,
      }
      if self._hps.coverage:
        to_return['coverage_loss'] = self._coverage_loss
      res = sess.run(to_return, feed_dict)
      return res
  
    def run_encoder(self, sess, batch):
      """For beam search decoding. Run the encoder on the batch and return the encoder states and decoder initial state.
  
      Args:
        sess: Tensorflow session.
        batch: Batch object that is the same example repeated across the batch (for beam search)
  
      Returns:
        enc_states: The encoder states. A tensor of shape [batch_size, <=max_enc_steps, 2*hidden_dim].
        dec_in_state: A LSTMStateTuple of shape ([1,hidden_dim],[1,hidden_dim])
      """
      feed_dict = self._make_feed_dict(batch, just_enc=True) 
      (enc_states, dec_in_state, global_step) = sess.run(
        [self._enc_states, self._dec_in_state, self.global_step], feed_dict) # run the encoder
  
      # dec_in_state is LSTMStateTuple shape ([batch_size,hidden_dim],[batch_size,hidden_dim])
      # Given that the batch is a single example repeated, dec_in_state is identical across the batch so we just take the top row.
      dec_in_state = tf.contrib.rnn.LSTMStateTuple(dec_in_state.c[0], dec_in_state.h[0])
      return enc_states, dec_in_state
  
  
    def decode_onestep(self, sess, batch, latest_tokens, enc_states, dec_init_states, prev_coverage):
      """For beam search decoding. Run the decoder for one step.
  
      Args:
        sess: Tensorflow session.
        batch: Batch object containing single example repeated across the batch
        latest_tokens: Tokens to be fed as input into the decoder for this timestep
        enc_states: The encoder states.
        dec_init_states: List of beam_size LSTMStateTuples; the decoder states from the previous timestep
        prev_coverage: List of np arrays. The coverage vectors from the previous timestep. List of None if not using coverage.
  
      Returns:
        ids: top 2k ids. shape [beam_size, 2*beam_size]
        probs: top 2k log probabilities. shape [beam_size, 2*beam_size]
        new_states: new states of the decoder. a list length beam_size containing
          LSTMStateTuples each of shape ([hidden_dim,],[hidden_dim,])
        attn_dists: List length beam_size containing lists length attn_length.
        p_gens: Generation probabilities for this step. A list length beam_size. List of None if in baseline mode.
        new_coverage: Coverage vectors for this step. A list of arrays. List of None if coverage is not turned on.
      """
  
      beam_size = len(dec_init_states)
  
      # Turn dec_init_states (a list of LSTMStateTuples) into a single LSTMStateTuple for the batch
      cells = [np.expand_dims(state.c, axis=0) for state in dec_init_states]
      hiddens = [np.expand_dims(state.h, axis=0) for state in dec_init_states]
      new_c = np.concatenate(cells, axis=0)  # shape [batch_size,hidden_dim]
      new_h = np.concatenate(hiddens, axis=0)  # shape [batch_size,hidden_dim]
      new_dec_in_state = tf.contrib.rnn.LSTMStateTuple(new_c, new_h)
  
      feed = {
          self._enc_states: enc_states,
          self._enc_padding_mask: batch.enc_padding_mask,
          self._dec_in_state: new_dec_in_state,
          self._dec_batch: np.transpose(np.array([latest_tokens])),
      }
  
      to_return = {
        "ids": self._topk_ids,
        "probs": self._topk_log_probs,
        "states": self._dec_out_state,
        "attn_dists": self.attn_dists
      }
      
      if self._hps.hier:
        feed[self._enc_batch_sections] = batch.batch_sections  # shape=[batch-size, num-sections, enc-seq-len]
        feed[self._batch_sections_len] = batch.batch_sections_len
        feed[self._doc_sec_lens] = batch.batch_doc_sec_lens
        feed[self._enc_section_padding_mask] = batch.enc_section_padding_mask
        feed[self._enc_lens] = batch.enc_lens
        to_return['attn_dists_sec'] = self.attn_dists_sec
  
      if FLAGS.pointer_gen:
        feed[self._enc_batch_extend_vocab] = batch.enc_batch_extend_vocab
        feed[self._max_art_oovs] = batch.max_art_oovs
        to_return['p_gens'] = self.p_gens
  
      if self._hps.coverage:
        feed[self.prev_coverage] = np.stack(prev_coverage, axis=0)
        to_return['coverage'] = self.coverage
  
      results = sess.run(to_return, feed_dict=feed) # run the decoder step
  
      # Convert results['states'] (a single LSTMStateTuple) into a list of LSTMStateTuple -- one for each hypothesis
      new_states = [tf.contrib.rnn.LSTMStateTuple(results['states'].c[i, :], results['states'].h[i, :]) for i in xrange(beam_size)]
  
      # Convert singleton list containing a tensor to a list of k arrays
      assert len(results['attn_dists'])==1
      attn_dists = results['attn_dists'][0].tolist()
      
      if 'attn_dists_sec' in results:
        if len(results['attn_dists_sec']) > 0:
          attn_dists_sec = results['attn_dists_sec'][0].tolist()
        else: attn_dists_sec = None
      else:
        attn_dists_sec = None
  
      if FLAGS.pointer_gen:
        # Convert singleton list containing a tensor to a list of k arrays
        assert len(results['p_gens'])==1
        p_gens = results['p_gens'][0].tolist()
      else:
        p_gens = [None for _ in xrange(beam_size)]
  
      # Convert the coverage tensor to a list length k containing the coverage vector for each hypothesis
      if FLAGS.coverage:
        new_coverage = results['coverage'].tolist()
        assert len(new_coverage) == beam_size
      else:
        new_coverage = [None for _ in xrange(beam_size)]
  
      return results['ids'], results['probs'], new_states, attn_dists, p_gens, new_coverage, attn_dists_sec
  
  
  def _mask_and_avg(values, padding_mask):
    """Applies mask to values then returns overall average (a scalar)
  
    Args:
      values: a list length max_dec_steps containing arrays shape (batch_size).
      padding_mask: tensor shape (batch_size, max_dec_steps) containing 1s and 0s.
  
    Returns:
      a scalar
    """
  
    dec_lens = tf.reduce_sum(padding_mask, axis=1) # shape batch_size. float32
    values_per_step = [v * padding_mask[:,dec_step] for dec_step,v in enumerate(values)]
    values_per_ex = sum(values_per_step)/dec_lens # shape (batch_size); normalized value for each batch member
    return tf.reduce_mean(values_per_ex) # overall average
  
  
  def _coverage_loss(attn_dists, padding_mask):
    """Calculates the coverage loss from the attention distributions.
  
    Args:
      attn_dists: The attention distributions for each decoder timestep. A list length max_dec_steps containing shape (batch_size, attn_length)
      padding_mask: shape (batch_size, max_dec_steps).
  
    Returns:
      coverage_loss: scalar
    """
    coverage = tf.zeros_like(attn_dists[0]) # shape (batch_size, attn_length). Initial coverage is zero.
    covlosses = [] # Coverage loss per decoder timestep. Will be list length max_dec_steps containing shape (batch_size).
    for a in attn_dists:
      covloss = tf.reduce_sum(tf.minimum(a, coverage), [1]) # calculate the coverage loss for this step
      covlosses.append(covloss)
      coverage += a # update the coverage vector
    coverage_loss = _mask_and_avg(covlosses, padding_mask)
    return coverage_loss
  