model.py

# 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