attention_model.py

# -*- coding: utf-8 -*-
#   Copyright (c) 2019 PaddlePaddle Authors. All Rights Reserve.
#
# 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 __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import paddle.fluid as fluid
import numpy as np
from paddle.fluid import ParamAttr
from paddle.fluid.dygraph.base import to_variable
from paddle.fluid.dygraph.nn import Embedding
from rnn import BasicLSTMUnit
import numpy as np

INF = 1. * 1e5
alpha = 0.6
uniform_initializer = lambda x: fluid.initializer.UniformInitializer(low=-x, high=x)
zero_constant = fluid.initializer.Constant(0.0)

class AttentionModel(fluid.dygraph.Layer):
    def __init__(self,
                 hidden_size,
                 src_vocab_size,
                 tar_vocab_size,
                 batch_size,
                 num_layers=1,
                 init_scale=0.1,
                 dropout=None,
                 beam_size=1,
                 beam_start_token=1,
                 beam_end_token=2,
                 beam_max_step_num=100,
                 mode='train'):
        super(AttentionModel, self).__init__()
        self.hidden_size = hidden_size
        self.src_vocab_size = src_vocab_size
        self.tar_vocab_size = tar_vocab_size
        self.batch_size = batch_size
        self.num_layers = num_layers
        self.init_scale = init_scale
        self.dropout = dropout
        self.beam_size = beam_size
        self.beam_start_token = beam_start_token
        self.beam_end_token = beam_end_token
        self.beam_max_step_num = beam_max_step_num
        self.mode = mode
        self.kinf = 1e9

        param_attr = ParamAttr(initializer=uniform_initializer(self.init_scale))
        bias_attr = ParamAttr(initializer=zero_constant)
        forget_bias = 1.0

        self.src_embeder = Embedding(
            size=[self.src_vocab_size, self.hidden_size],
            param_attr=fluid.ParamAttr(
                name='source_embedding',
                initializer=uniform_initializer(init_scale)))

        self.tar_embeder = Embedding(
            size=[self.tar_vocab_size, self.hidden_size],
            is_sparse=False,
            param_attr=fluid.ParamAttr(
                name='target_embedding',
                initializer=uniform_initializer(init_scale)))

        self.enc_units = []
        for i in range(num_layers):
            self.enc_units.append(
                self.add_sublayer("enc_units_%d" % i,
                    BasicLSTMUnit(
                    hidden_size=self.hidden_size, 
                    input_size=self.hidden_size,
                    param_attr=param_attr, 
                    bias_attr=bias_attr, 
                    forget_bias=forget_bias)))

        self.dec_units = []
        for i in range(num_layers):
            if i == 0:
                self.dec_units.append(
                    self.add_sublayer("dec_units_%d" % i,
                        BasicLSTMUnit(
                        hidden_size=self.hidden_size, 
                        input_size=self.hidden_size * 2,
                        param_attr=param_attr, 
                        bias_attr=bias_attr, 
                        forget_bias=forget_bias)))
            else:
                self.dec_units.append(
                    self.add_sublayer("dec_units_%d" % i,
                    BasicLSTMUnit(
                    hidden_size=self.hidden_size, 
                    input_size=self.hidden_size,
                    param_attr=param_attr, 
                    bias_attr=bias_attr, 
                    forget_bias=forget_bias)))
        
        self.fc = fluid.dygraph.nn.Linear(self.hidden_size,
                self.tar_vocab_size, 
                param_attr=param_attr,
                bias_attr=False)
        
        self.attn_fc = fluid.dygraph.nn.Linear(self.hidden_size, 
                self.hidden_size, 
                param_attr=param_attr,
                bias_attr=False)
        
        self.concat_fc = fluid.dygraph.nn.Linear(2 * self.hidden_size, 
                self.hidden_size, 
                param_attr=param_attr,
                bias_attr=False)

    def _transpose_batch_time(self, x):
        return fluid.layers.transpose(x, [1, 0] + list(range(2, len(x.shape))))

    def _merge_batch_beams(self, x):
        return fluid.layers.reshape(x, shape=(-1,x.shape[2]))

    def tile_beam_merge_with_batch(self, x):
        x = fluid.layers.unsqueeze(x, [1])  # [batch_size, 1, ...]
        expand_times = [1] * len(x.shape)
        expand_times[1] = self.beam_size
        x = fluid.layers.expand(x, expand_times)  # [batch_size, beam_size, ...]
        x = fluid.layers.transpose(x, list(range(2, len(x.shape))) +
                         [0, 1])  # [..., batch_size, beam_size]
        # use 0 to copy to avoid wrong shape
        x = fluid.layers.reshape(
            x, shape=[0] *
            (len(x.shape) - 2) + [-1])  # [..., batch_size * beam_size]
        x = fluid.layers.transpose(
            x, [len(x.shape) - 1] +
            list(range(0, len(x.shape) - 1)))  # [batch_size * beam_size, ...]
        return x

    def _split_batch_beams(self, x):
        return fluid.layers.reshape(x, shape=(-1, self.beam_size, x.shape[1]))

    def _expand_to_beam_size(self, x):
        x = fluid.layers.unsqueeze(x, [1])
        expand_times = [1] * len(x.shape)
        expand_times[1] = self.beam_size
        x = fluid.layers.expand(x, expand_times)
        return x

    def _real_state(self, state, new_state, step_mask):
        new_state = fluid.layers.elementwise_mul(new_state, step_mask, axis=0) - \
                    fluid.layers.elementwise_mul(state, (step_mask - 1), axis=0)
        return new_state
        
    def _gather(self, x, indices, batch_pos):
        topk_coordinates = fluid.layers.stack([batch_pos, indices], axis=2)
        return fluid.layers.gather_nd(x, topk_coordinates)

    def attention(self, query, enc_output, mask=None):
        query = fluid.layers.unsqueeze(query, [1])
        memory = self.attn_fc(enc_output)
        attn = fluid.layers.matmul(query, memory, transpose_y=True)

        if mask:
            attn = fluid.layers.transpose(attn, [1, 0, 2])
            attn = fluid.layers.elementwise_add(attn, mask * 1000000000, -1)
            attn = fluid.layers.transpose(attn, [1, 0, 2])
        weight = fluid.layers.softmax(attn)
        weight_memory = fluid.layers.matmul(weight, memory)

        return weight_memory

    def forward(self, inputs):
        inputs = [fluid.dygraph.to_variable(np_inp) for np_inp in inputs]
        src, tar, label, src_sequence_length, tar_sequence_length = inputs
        if src.shape[0] < self.batch_size:
            self.batch_size = src.shape[0]
        src_emb = self.src_embeder(self._transpose_batch_time(src))
        
        enc_hidden = to_variable(np.zeros((self.num_layers, self.batch_size, self.hidden_size), dtype='float32'))
        enc_cell = to_variable(np.zeros((self.num_layers, self.batch_size, self.hidden_size), dtype='float32'))

        max_seq_len = src_emb.shape[0]
        enc_len_mask = fluid.layers.sequence_mask(src_sequence_length, maxlen=max_seq_len, dtype="float32")
        enc_padding_mask = (enc_len_mask - 1.0)
        enc_len_mask = fluid.layers.transpose(enc_len_mask, [1, 0])
        enc_states = [[enc_hidden, enc_cell]]
        enc_outputs = []
        for l in range(max_seq_len):
            step_input = src_emb[l]
            step_mask = enc_len_mask[l]
            enc_hidden, enc_cell = enc_states[l]
            new_enc_hidden, new_enc_cell = [], []
            for i in range(self.num_layers):
                new_hidden, new_cell = self.enc_units[i](step_input, enc_hidden[i], enc_cell[i])
                new_enc_hidden.append(new_hidden)
                new_enc_cell.append(new_cell)
                if self.dropout != None and self.dropout > 0.0:
                    step_input = fluid.layers.dropout(
                        new_hidden,
                        dropout_prob=self.dropout,
                        dropout_implementation='upscale_in_train')
                else:
                    step_input = new_hidden
            new_enc_hidden = [self._real_state(enc_hidden[i], new_enc_hidden[i], step_mask) for i in range(self.num_layers)]
            new_enc_cell = [self._real_state(enc_cell[i], new_enc_cell[i], step_mask) for i in range(self.num_layers)]
            enc_states.append([new_enc_hidden, new_enc_cell])
            enc_outputs.append(step_input)
        enc_outputs = fluid.layers.stack(enc_outputs)
        enc_outputs = self._transpose_batch_time(enc_outputs)

        if self.mode in ['train', 'eval']: 
            # calculation with input_feed derives from paper: https://arxiv.org/pdf/1508.04025.pdf
            input_feed = to_variable(np.zeros((self.batch_size, self.hidden_size), dtype='float32'))
            
            dec_hidden, dec_cell = enc_states[-1]
            tar_emb = self.tar_embeder(self._transpose_batch_time(tar))
            max_seq_len = tar_emb.shape[0]
            dec_output = []

            for step_idx in range(max_seq_len):
                step_input = tar_emb[step_idx]
                step_input = fluid.layers.concat([step_input, input_feed], 1)
                new_dec_hidden, new_dec_cell = [], []
                for i in range(self.num_layers):
                    new_hidden, new_cell = self.dec_units[i](step_input, dec_hidden[i], dec_cell[i])

                    new_dec_hidden.append(new_hidden)
                    new_dec_cell.append(new_cell)
                    if self.dropout != None and self.dropout > 0.0:
                        step_input = fluid.layers.dropout(
                            new_hidden,
                            dropout_prob=self.dropout,
                            dropout_implementation='upscale_in_train')
                    else:
                        step_input = new_hidden

                dec_att = self.attention(step_input, enc_outputs, enc_padding_mask)
                dec_att = fluid.layers.squeeze(dec_att, [1])
                concat_att_out = fluid.layers.concat([dec_att, step_input], 1)
                out = self.concat_fc(concat_att_out)
                input_feed = out
                dec_output.append(out)
                dec_hidden, dec_cell = new_dec_hidden, new_dec_cell

            dec_output = fluid.layers.stack(dec_output)
            dec_output = self.fc(self._transpose_batch_time(dec_output))
        
            loss = fluid.layers.softmax_with_cross_entropy(
            logits=dec_output, label=label, soft_label=False)
            loss = fluid.layers.squeeze(loss, axes=[2])
            max_tar_seq_len = fluid.layers.shape(tar)[1]
            tar_mask = fluid.layers.sequence_mask(
                tar_sequence_length, maxlen=max_tar_seq_len, dtype='float32')
            loss = loss * tar_mask
            loss = fluid.layers.reduce_mean(loss, dim=[0])
            loss = fluid.layers.reduce_sum(loss)
            return loss
            
        elif self.mode in ['beam_search']:
            enc_outputs = self.tile_beam_merge_with_batch(enc_outputs)
            enc_padding_mask = self.tile_beam_merge_with_batch(enc_padding_mask)
            batch_beam_shape = (self.batch_size, self.beam_size)
            vocab_size_tensor = to_variable(np.full((1), self.tar_vocab_size))
            start_token_tensor = to_variable(np.full(batch_beam_shape, self.beam_start_token, dtype='int64'))
            end_token_tensor = to_variable(np.full(batch_beam_shape, self.beam_end_token, dtype='int64'))
            step_input = self.tar_embeder(start_token_tensor)
            input_feed = to_variable(np.zeros((self.batch_size, self.hidden_size), dtype='float32'))
            input_feed = self._expand_to_beam_size(input_feed)
            input_feed = self._merge_batch_beams(input_feed)
            beam_finished = to_variable(np.full(batch_beam_shape, 0, dtype='float32'))
            beam_state_log_probs = to_variable(np.array([[0.] + [-self.kinf] * (self.beam_size - 1)], dtype="float32"))
            beam_state_log_probs = fluid.layers.expand(beam_state_log_probs, [self.batch_size, 1])
            
            dec_hidden, dec_cell = enc_states[-1]
            dec_hidden = [self._expand_to_beam_size(state) for state in dec_hidden]
            dec_cell = [self._expand_to_beam_size(state) for state in dec_cell]
            
            batch_pos = fluid.layers.expand(
                fluid.layers.unsqueeze(to_variable(np.arange(0, self.batch_size, 1, dtype="int64")), [1]),
                [1, self.beam_size])
            predicted_ids = []
            parent_ids = []

            for step_idx in range(self.beam_max_step_num):
                if fluid.layers.reduce_sum(1 - beam_finished).numpy()[0] == 0:
                    break
                step_input = self._merge_batch_beams(step_input)
                step_input = fluid.layers.concat([step_input, input_feed], 1)
                new_dec_hidden, new_dec_cell = [], []
                dec_hidden = [self._merge_batch_beams(state) for state in dec_hidden]
                dec_cell = [self._merge_batch_beams(state) for state in dec_cell]

                for i in range(self.num_layers):
                    new_hidden, new_cell = self.dec_units[i](step_input, dec_hidden[i], dec_cell[i])
                    new_dec_hidden.append(new_hidden)
                    new_dec_cell.append(new_cell)
                    if self.dropout != None and self.dropout > 0.0:
                        step_input = fluid.layers.dropout(
                            new_hidden,
                            dropout_prob=self.dropout,
                            dropout_implementation='upscale_in_train')
                    else:
                        step_input = new_hidden
                dec_att = self.attention(step_input, enc_outputs, enc_padding_mask)
                dec_att = fluid.layers.squeeze(dec_att, [1])
                concat_att_out = fluid.layers.concat([dec_att, step_input], 1)
                out = self.concat_fc(concat_att_out)
                input_feed = out
                cell_outputs = self._split_batch_beams(out)
                cell_outputs = self.fc(cell_outputs)
                step_log_probs = fluid.layers.log(fluid.layers.softmax(cell_outputs))
                noend_array = [-self.kinf] * self.tar_vocab_size
                noend_array[self.beam_end_token] = 0 # [-kinf, -kinf, ..., 0, -kinf, ...]
                noend_mask_tensor = to_variable(np.array(noend_array,dtype='float32'))
                # set finished position to one-hot probability of <eos>
                step_log_probs = fluid.layers.elementwise_mul(
                        fluid.layers.expand(fluid.layers.unsqueeze(beam_finished, [2]), [1, 1, self.tar_vocab_size]),
                    noend_mask_tensor, axis=-1) - \
                    fluid.layers.elementwise_mul(step_log_probs, (beam_finished - 1), axis=0)
                log_probs = fluid.layers.elementwise_add(
                    x=step_log_probs, y=beam_state_log_probs, axis=0)
                scores = fluid.layers.reshape(log_probs, [-1, self.beam_size * self.tar_vocab_size])
                topk_scores, topk_indices = fluid.layers.topk(input=scores, k=self.beam_size)
                beam_indices = fluid.layers.elementwise_floordiv(topk_indices, vocab_size_tensor) # in which beam
                token_indices = fluid.layers.elementwise_mod(topk_indices, vocab_size_tensor) # position in beam
                next_log_probs = self._gather(scores, topk_indices, batch_pos) # 

                new_dec_hidden = [self._split_batch_beams(state) for state in new_dec_hidden]
                new_dec_cell = [self._split_batch_beams(state) for state in new_dec_cell]
                new_dec_hidden = [self._gather(x, beam_indices, batch_pos) for x in new_dec_hidden]
                new_dec_cell = [self._gather(x, beam_indices, batch_pos) for x in new_dec_cell]
                
                next_finished = self._gather(beam_finished, beam_indices, batch_pos)              
                next_finished = fluid.layers.cast(next_finished, "bool")
                next_finished = fluid.layers.logical_or(next_finished, fluid.layers.equal(token_indices, end_token_tensor))
                next_finished = fluid.layers.cast(next_finished, "float32")
                # prepare for next step
                dec_hidden, dec_cell = new_dec_hidden, new_dec_cell
                beam_finished = next_finished
                beam_state_log_probs = next_log_probs
                step_input = self.tar_embeder(token_indices)
                predicted_ids.append(token_indices)
                parent_ids.append(beam_indices)

            predicted_ids = fluid.layers.stack(predicted_ids)
            parent_ids = fluid.layers.stack(parent_ids)
            predicted_ids = fluid.layers.gather_tree(predicted_ids, parent_ids)
            predicted_ids = self._transpose_batch_time(predicted_ids)
            return predicted_ids
        else:
            print("not support mode ", self.mode)
            raise Exception("not support mode: " + self.mode)