attention_decoder.py

# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
# Modifications Copyright 2017 Abigail See
#
# 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.
# ==============================================================================

"""This file defines the decoder"""

import tensorflow as tf
from tensorflow.python.ops import variable_scope
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import nn_ops
from tensorflow.python.ops import math_ops

FLAGS = tf.flags.FLAGS


def attention_decoder(decoder_inputs, initial_state, encoder_states, enc_padding_mask, local_states, global_states,
                      event_padding_mask, event_len, time_key, cell,
                      initial_state_attention=False, pointer_gen=True, use_coverage=False, prev_coverage=None):
    """
    Args:
      decoder_inputs: A list of 2D Tensors [batch_size x input_size].
      initial_state: 2D Tensor [batch_size x cell.state_size].
      encoder_states: 3D Tensor [batch_size x attn_length x attn_size].
      enc_padding_mask: 2D Tensor [batch_size x attn_length] containing 1s and 0s; indicates which of the encoder locations are padding (0) or a real token (1).
      cell: rnn_cell.RNNCell defining the cell function and size.
      initial_state_attention:
        Note that this attention decoder passes each decoder input through a linear layer with the previous step's context vector to get a modified version of the input. If initial_state_attention is False, on the first decoder step the "previous context vector" is just a zero vector. If initial_state_attention is True, we use initial_state to (re)calculate the previous step's context vector. We set this to False for train/eval mode (because we call attention_decoder once for all decoder steps) and True for decode mode (because we call attention_decoder once for each decoder step).
      pointer_gen: boolean. If True, calculate the generation probability p_gen for each decoder step.
      use_coverage: boolean. If True, use coverage mechanism.
      prev_coverage:
        If not None, a tensor with shape (batch_size, attn_length). The previous step's coverage vector. This is only not None in decode mode when using coverage.

    Returns:
      outputs: A list of the same length as decoder_inputs of 2D Tensors of
        shape [batch_size x cell.output_size]. The output vectors.
      state: The final state of the decoder. A tensor shape [batch_size x cell.state_size].
      attn_dists: A list containing tensors of shape (batch_size,attn_length).
        The attention distributions for each decoder step.
      p_gens: List of scalars. The values of p_gen for each decoder step. Empty list if pointer_gen=False.
      coverage: Coverage vector on the last step computed. None if use_coverage=False.
    """
    with variable_scope.variable_scope("attention_decoder") as scope:
        batch_size = encoder_states.get_shape()[0].value
        attn_size = encoder_states.get_shape()[2].value

        time_size = time_key.get_shape()[-1].value

        # Reshape encoder_states (need to insert a dim)
        encoder_states = tf.expand_dims(encoder_states, axis=2)  # now is shape (batch_size, attn_len, 1, attn_size)
        local_states = tf.expand_dims(local_states, axis=2)
        global_states = tf.expand_dims(global_states, axis=2)
        time_states = tf.expand_dims(tf.tile(tf.expand_dims(time_key, axis=0), [batch_size, 1, 1]), axis=2)

        # To calculate attention, we calculate
        #   v^T tanh(W_h h_i + W_s s_t + b_attn)
        # where h_i is an encoder state, and s_t a decoder state.
        # attn_vec_size is the length of the vectors v, b_attn, (W_h h_i) and (W_s s_t).
        # We set it to be equal to the size of the encoder states.
        attention_vec_size = attn_size

        # Get the weight matrix W_h and apply it to each encoder state to get (W_h h_i), the encoder features
        W_h = variable_scope.get_variable("W_h", [1, 1, attn_size, attention_vec_size])
        W_t = variable_scope.get_variable("W_t", [1, 1, time_size, attention_vec_size])
        encoder_features = nn_ops.conv2d(encoder_states, W_h, [1, 1, 1, 1], "SAME")
        local_features = nn_ops.conv2d(local_states, W_h, [1, 1, 1, 1], "SAME")
        time_features = nn_ops.conv2d(time_states, W_t, [1, 1, 1, 1], "SAME")

        # Get the weight vectors v and w_c (w_c is for coverage)
        v = variable_scope.get_variable("v", [attention_vec_size])
        v_e = variable_scope.get_variable("v_e", [attention_vec_size])
        v_t = variable_scope.get_variable("v_t", [attention_vec_size])

        if use_coverage:
            with variable_scope.variable_scope("coverage"):
                w_c = variable_scope.get_variable("w_c", [1, 1, 1, attention_vec_size])

        if prev_coverage is not None:  # for beam search mode with coverage
            # reshape from (batch_size, attn_length) to (batch_size, attn_len, 1, 1)
            prev_coverage = tf.expand_dims(tf.expand_dims(prev_coverage, 2), 3)

        def attention(decoder_state, coverage=None):
            """Calculate the context vector and attention distribution from the decoder state.

            Args:
              decoder_state: state of the decoder
              coverage: Optional. Previous timestep's coverage vector, shape (batch_size, attn_len, 1, 1).

            Returns:
              context_vector: weighted sum of encoder_states
              attn_dist: attention distribution
              coverage: new coverage vector. shape (batch_size, attn_len, 1, 1)
            """
            with variable_scope.variable_scope("Attention"):
                # Pass the decoder state through a linear layer (this is W_s s_t + b_attn in the paper)
                decoder_features = linear(decoder_state, attention_vec_size,
                                          True)  # shape (batch_size, attention_vec_size)
                decoder_features = tf.expand_dims(tf.expand_dims(decoder_features, 1),
                                                  1)  # reshape to (batch_size, 1, 1, attention_vec_size)

                def masked_attention(e, mask):
                    """Take softmax of e then apply enc_padding_mask and re-normalize"""
                    attn_dist = nn_ops.softmax(e)  # take softmax. shape (batch_size, attn_length)
                    attn_dist *= mask  # apply mask
                    masked_sums = tf.reduce_sum(attn_dist, axis=1)  # shape (batch_size)
                    return attn_dist / tf.reshape(masked_sums, [-1, 1])  # re-normalize

                if use_coverage and coverage is not None:  # non-first step of coverage
                    # Multiply coverage vector by w_c to get coverage_features.
                    coverage_features = nn_ops.conv2d(coverage, w_c, [1, 1, 1, 1],
                                                      "SAME")  # c has shape (batch_size, attn_length, 1, attention_vec_size)

                    # Calculate v^T tanh(W_h h_i + W_s s_t + w_c c_i^t + b_attn)
                    e = math_ops.reduce_sum(v * math_ops.tanh(encoder_features + decoder_features + coverage_features),
                                            [2, 3])  # shape (batch_size,attn_length)
                    l_e = math_ops.reduce_sum(v_e * math_ops.tanh(local_features + decoder_features),
                                              [2, 3])  # calculate e

                    l_t = math_ops.reduce_sum(v_t * math_ops.tanh(time_features + decoder_features),
                                              [2, 3])  # calculate e

                    # Calculate attention distribution
                    attn_dist = masked_attention(e, enc_padding_mask)

                    local_attn_dist = masked_attention(l_e, event_padding_mask)

                    time_attn_dist = masked_attention(l_t, event_padding_mask)

                    # Update coverage vector
                    coverage += array_ops.reshape(attn_dist, [batch_size, -1, 1, 1])

                    event_rep = math_ops.reduce_sum(
                        array_ops.reshape(local_attn_dist, [batch_size, -1, 1, 1]) * local_states,
                        [1, 2])  # shape (batch_size, attn_size).

                    local_rep = math_ops.reduce_sum(
                        array_ops.reshape(time_attn_dist, [batch_size, -1, 1, 1]) * local_states,
                        [1, 2])  # shape (batch_size, attn_size).

                    global_rep = math_ops.reduce_sum(
                        array_ops.reshape(time_attn_dist, [batch_size, -1, 1, 1]) * global_states,
                        [1, 2])  # shape (batch_size, attn_size).

                else:
                    # Calculate v^T tanh(W_h h_i + W_s s_t + b_attn)
                    e = math_ops.reduce_sum(v * math_ops.tanh(encoder_features + decoder_features),
                                            [2, 3])  # calculate e

                    l_e = math_ops.reduce_sum(v_e * math_ops.tanh(local_features + decoder_features),
                                              [2, 3])  # calculate e

                    l_t = math_ops.reduce_sum(v_t * math_ops.tanh(time_features + decoder_features),
                                              [2, 3])  # calculate e

                    # Calculate attention distribution
                    attn_dist = masked_attention(e, enc_padding_mask)

                    local_attn_dist = masked_attention(l_e, event_padding_mask)

                    time_attn_dist = masked_attention(l_t, event_padding_mask)

                    # local_attn_dist_tmp = tf.reshape(tf.tile(tf.expand_dims(local_attn_dist, 2), [1, 1, event_len]), [batch_size, -1])
                    # attn_dist = attn_dist * local_attn_dist_tmp

                    if use_coverage:  # first step of training
                        coverage = tf.expand_dims(tf.expand_dims(attn_dist, 2), 2)  # initialize coverage

                    event_rep = math_ops.reduce_sum(
                        array_ops.reshape(local_attn_dist, [batch_size, -1, 1, 1]) * local_states,
                        [1, 2])  # shape (batch_size, attn_size).

                    local_rep = math_ops.reduce_sum(
                        array_ops.reshape(time_attn_dist, [batch_size, -1, 1, 1]) * local_states,
                        [1, 2])  # shape (batch_size, attn_size).

                    global_rep = math_ops.reduce_sum(
                        array_ops.reshape(time_attn_dist, [batch_size, -1, 1, 1]) * global_states,
                        [1, 2])  # shape (batch_size, attn_size).

                # Calculate the context vector from attn_dist and encoder_states
                context_vector = math_ops.reduce_sum(
                    array_ops.reshape(attn_dist, [batch_size, -1, 1, 1]) * encoder_states,
                    [1, 2])  # shape (batch_size, attn_size).
                context_vector = array_ops.reshape(context_vector, [-1, attn_size])

            return context_vector, attn_dist, coverage, event_rep, local_rep, global_rep, time_attn_dist

        outputs = []
        attn_dists = []
        event_dists = []
        p_gens = []
        state = initial_state
        coverage = prev_coverage  # initialize coverage to None or whatever was passed in
        context_vector = array_ops.zeros([batch_size, attn_size])
        context_vector.set_shape([None, attn_size])  # Ensure the second shape of attention vectors is set.
        if initial_state_attention:  # true in decode mode
            # Re-calculate the context vector from the previous step so that we can pass it through a linear layer with this step's input to get a modified version of the input
            context_vector, _, coverage, event_rep, local_rep, global_rep, time_attn_dist = attention(initial_state,
                                                                                                      coverage)  # in decode mode, this is what updates the coverage vector
        for i, inp in enumerate(decoder_inputs):
            if i > 0:
                variable_scope.get_variable_scope().reuse_variables()

            # Merge input and previous attentions into one vector x of the same size as inp
            input_size = inp.get_shape().with_rank(2)[1]
            if input_size.value is None:
                raise ValueError("Could not infer input size from input: %s" % inp.name)
            x = linear([inp] + [context_vector], input_size, True)

            # Run the decoder RNN cell. cell_output = decoder state
            cell_output, state = cell(x, state)

            # Run the attention mechanism.
            if i == 0 and initial_state_attention:  # always true in decode mode
                with variable_scope.variable_scope(variable_scope.get_variable_scope(),
                                                   reuse=True):  # you need this because you've already run the initial attention(...) call
                    context_vector, attn_dist, _, event_rep, local_rep, global_rep, time_attn_dist = attention(state,
                                                                                                               coverage)  # don't allow coverage to update
            else:
                context_vector, attn_dist, coverage, event_rep, local_rep, global_rep, time_attn_dist = attention(state,
                                                                                                                  coverage)
            attn_dists.append(attn_dist)
            event_dists.append(time_attn_dist)

            with tf.variable_scope('memory', reuse=tf.AUTO_REUSE):
                # add memory representations
                g_1 = tf.layers.dense(tf.concat([cell_output, context_vector, local_rep], -1), 1, tf.nn.sigmoid)
                g_2 = tf.layers.dense(tf.concat([cell_output, context_vector, global_rep], -1), 1, tf.nn.sigmoid)

                local_rep = local_rep * g_1

                state_h = (1 - g_2) * state.h + g_2 * tf.layers.dense(global_rep, FLAGS.hidden_dim)
                state = tf.contrib.rnn.LSTMStateTuple(state.c, state_h)

            # Calculate p_gen
            if pointer_gen:
                with tf.variable_scope('calculate_pgen'):
                    p_gen = linear([context_vector, state.c, state.h, x], 1, True)  # a scalar
                    p_gen = tf.sigmoid(p_gen)
                    # if i > 20:
                    #     p_gen = tf.ones_like(p_gen)
                    p_gens.append(p_gen)

            # Concatenate the cell_output (= decoder state) and the context vector, and pass them through a linear layer
            # This is V[s_t, h*_t] + b in the paper
            with variable_scope.variable_scope("AttnOutputProjection"):
                output = linear([cell_output] + [context_vector] + [local_rep], cell.output_size, True)
                # output = linear([cell_output] + [context_vector], cell.output_size, True)

            outputs.append(output)
            if i % 10 == 0:
                tf.logging.info('Building %d decoder steps...' % i)

        # If using coverage, reshape it
        if coverage is not None:
            coverage = array_ops.reshape(coverage, [batch_size, -1])

        return outputs, state, attn_dists, event_dists, p_gens, coverage


def linear(args, output_size, bias, bias_start=0.0, scope=None):
    """Linear map: sum_i(args[i] * W[i]), where W[i] is a variable.

    Args:
      args: a 2D Tensor or a list of 2D, batch x n, Tensors.
      output_size: int, second dimension of W[i].
      bias: boolean, whether to add a bias term or not.
      bias_start: starting value to initialize the bias; 0 by default.
      scope: VariableScope for the created subgraph; defaults to "Linear".

    Returns:
      A 2D Tensor with shape [batch x output_size] equal to
      sum_i(args[i] * W[i]), where W[i]s are newly created matrices.

    Raises:
      ValueError: if some of the arguments has unspecified or wrong shape.
    """
    if args is None or (isinstance(args, (list, tuple)) and not args):
        raise ValueError("`args` must be specified")
    if not isinstance(args, (list, tuple)):
        args = [args]

    # Calculate the total size of arguments on dimension 1.
    total_arg_size = 0
    shapes = [a.get_shape().as_list() for a in args]
    for shape in shapes:
        if len(shape) != 2:
            raise ValueError("Linear is expecting 2D arguments: %s" % str(shapes))
        if not shape[1]:
            raise ValueError("Linear expects shape[1] of arguments: %s" % str(shapes))
        else:
            total_arg_size += shape[1]

    # Now the computation.
    with tf.variable_scope(scope or "Linear"):
        matrix = tf.get_variable("Matrix", [total_arg_size, output_size])
        if len(args) == 1:
            res = tf.matmul(args[0], matrix)
        else:
            res = tf.matmul(tf.concat(axis=1, values=args), matrix)
        if not bias:
            return res
        bias_term = tf.get_variable(
            "Bias", [output_size], initializer=tf.constant_initializer(bias_start))
    return res + bias_term