model.py

import numpy as np
import math
import torch.nn as nn
from torch.autograd import Variable
from torch.nn.utils.rnn import pad_packed_sequence, pack_padded_sequence
import torch
import torch.nn.functional as F
from registry import register
from functools import partial
registry = {}
register = partial(register, registry=registry)


@register('rnn')
class RNN(nn.Module):
    def __init__(self, args):
        super().__init__()
        self.dim = args.emb_dim
        self.encoder_layer = args.encoder_layer
        self.embedding = nn.Embedding(args.vocab_size, self.dim, padding_idx=0)
        self.embedding.weight.requires_grad = True
        self.rnn = nn.LSTM(input_size=self.dim, hidden_size=int(self.dim // 2), batch_first=True, num_layers=self.encoder_layer, bidirectional=True)
        self.linear1 = nn.Linear(self.dim, self.dim)
        self.linear2 = nn.Linear(self.dim, self.dim)
        self.activiation = nn.Tanh()
        self.dropout = nn.Dropout(p=args.drop_rate)

    def forward(self, x, mask=None):
        x_embed = self.embedding(x)
        shape = x_embed.size()

        mask = mask.squeeze().cpu().detach().numpy()
        mask = [np.sum(e != 0) for e in mask]

        # rnn pack
        packed = pack_padded_sequence(x_embed, mask, batch_first=True, enforce_sorted=False)
        encoder_outputs_packed, (h_last, c_last) = self.rnn(packed)
        rnn_output, _ = pad_packed_sequence(encoder_outputs_packed, batch_first=True)

        output = list()
        for index in range(len(mask)):
            temp = rnn_output[index, mask[index]-1, :]
            output.append(temp)
        output = torch.reshape(torch.cat(output, dim=0), (shape[0], self.dim))

        output = self.dropout(output)
        output = self.linear1(output)
        output = self.activiation(output)
        output = self.dropout(output)
        output = self.linear2(output)
        return output


@register('cnn')
class CNN(nn.Module):
    def __init__(self, args):
        super().__init__()
        self.dim = args.emb_dim
        self.embedding = nn.Embedding(args.vocab_size, self.dim, padding_idx=0)
        self.embedding.weight.data.uniform_(-1, 1)
        self.windows = [3, 4, 5]
        self.convs = list()
        for window in self.windows:
            self.convs.append(nn.Conv2d(1, self.dim, (window, self.dim)).cuda())
        self.dropout = torch.nn.Dropout(p=args.drop_rate)
        self.proj = nn.Linear(len(self.windows)*self.dim, self.dim)

    def forward(self, x, x_lens):
        embeddings = self.embedding(x)
        embeddings = embeddings.unsqueeze(1)

        poolings = list()
        for conv in self.convs:
            conv_f = conv(embeddings)
            conv_f = F.relu(conv_f.squeeze(3))
            pooling = conv_f.max(dim=-1)[0]
            poolings.append(pooling)
        poolings = self.dropout(torch.cat(poolings, dim=-1))
        new_embed = self.proj(poolings)
        return new_embed


@register('attention')
class Attention(nn.Module):
    def __init__(self, args):
        super(Attention, self).__init__()
        self.dim = args.emb_dim
        self.embedding = nn.Embedding(args.vocab_size, self.dim, padding_idx=0)
        self.embedding.weight.requires_grad = True
        self.encoders = nn.ModuleList([OriAttention(args) for _ in range(1)])
        self.sublayer = SublayerConnection(args.drop_rate, self.dim)

    def forward(self, x, mask):
        x = self.embedding(x)
        for i, encoder in enumerate(self.encoders):
            x = self.sublayer(x, lambda x: encoder(x, mask))
        return x.masked_fill_(~mask, -float('inf')).max(dim=1)[0]


class OriAttention(nn.Module):
    def __init__(self, args):
        super().__init__()
        self.head = 1
        self.dropout = nn.Dropout(p=args.drop_rate)

    def forward(self, x, mask):
        q = x
        k = x
        v = x
        q, k, v = (split_last(a, (self.head, -1)).transpose(1, 2)
                   for a in [q, k, v])

        scores = torch.matmul(q, k.transpose(2, 3)) / (k.size(-1) ** 0.25)
        mask = torch.matmul(mask.float(), mask.transpose(1, 2).float()).bool()
        mask = mask.unsqueeze(1)
        mask = mask.repeat([1, self.head, 1, 1])
        scores.masked_fill_(~mask, -1e7)
        scores = F.softmax(scores, dim=2)
        scores = scores.transpose(2, 3)
        v_ = torch.matmul(scores, v)
        v_ = v_.transpose(1, 2).contiguous()
        v_ = merge_last(v_, 2)
        return v_


def l2norm(x):
    return x / x.norm(p=2, dim=1, keepdim=True)


@register('pam')
class Pamela(nn.Module):
    def __init__(self, args):
        super(Pamela, self).__init__()
        self.dim = args.emb_dim
        self.encoder_layer = args.encoder_layer
        self.embedding = nn.Embedding(args.vocab_size, self.dim, padding_idx=0)
        self.embedding.weight.requires_grad = True
        self.head = args.att_head_num
        self.encoders = nn.ModuleList([Pamelaformer(args) for _ in range(self.encoder_layer)])
        self.sublayer = SublayerConnection(args.drop_rate, self.dim)

    def forward(self, x, mask):
        x = self.embedding(x)
        shape = list(x.size())
        position = PositionalEncoding(shape[-1], shape[-2])
        pos_att = position(x)

        for i, encoder in enumerate(self.encoders):
            x = self.sublayer(x, lambda x: encoder(x, mask, pos_att))

        x = x.masked_fill_(~mask, 0).sum(dim=1)
        return l2norm(x)


class Pamelaformer(nn.Module):
    def __init__(self, args):
        super().__init__()
        self.self_attention = SAM(args)
        self.pos_attention = PAM(args)
        dim = args.emb_dim
        proj_dim = args.emb_dim
        self.merge = args.merge
        if self.merge:
            proj_dim = 2 * args.emb_dim
        self.projection = nn.Sequential(
            nn.Linear(proj_dim, dim),
            nn.ReLU()
        )
        self.dropout = nn.Dropout(p=args.drop_rate)

    def forward(self, x, mask, position, merge=True):
        att = self.self_attention(x, mask)
        pos = self.pos_attention(x, mask, position)
        if self.merge:
            c = self.projection(torch.cat([att, pos], dim=-1))
        else:
            c = self.projection(pos)
        return c


class PAM(nn.Module):
    def __init__(self, args):
        super().__init__()
        dim = args.emb_dim
        self.head = args.att_head_num
        self.projection = nn.Sequential(
            nn.Linear(dim, dim),
            nn.ReLU()
        )
        self.dropout = nn.Dropout(p=args.drop_rate)

    def forward(self, x, mask, pos):
        q = pos
        k = pos
        v = x
        q, k, v = (split_last(a, (self.head, -1)).transpose(1, 2) for a in [q, k, v])
        scores = torch.matmul(q, k.transpose(2, 3)) / (k.size(-1) ** 0.25)
        mask = torch.matmul(mask.float(), mask.transpose(1, 2).float()).bool()
        mask = mask.unsqueeze(1)
        mask = mask.repeat([1, self.head, 1, 1])
        scores.masked_fill_(~mask, -1e7)

        scores = F.softmax(scores, dim=2)
        scores = scores.transpose(2, 3)
        v_ = torch.matmul(scores, v)
        v_ = v_.transpose(1, 2).contiguous()
        v_ = merge_last(v_, 2)
        v_ = self.projection(v_)
        return v_


class SAM(nn.Module):
    def __init__(self, args):
        super().__init__()
        attention_dim = 32
        self.head = args.att_head_num
        hidden_size = args.emb_dim
        self.projectionq = nn.Sequential(
            nn.Linear(hidden_size, attention_dim),
            nn.ReLU()
        )
        self.projectionk = nn.Sequential(
            nn.Linear(hidden_size, attention_dim),
            nn.ReLU()
        )
        self.dropout = nn.Dropout(p=args.drop_rate)

    def forward(self, x, mask):
        q = self.projectionq(x)
        k = self.projectionk(x)
        v = x
        q, k, v = (split_last(a, (self.head, -1)).transpose(1, 2)
                   for a in [q, k, v])

        scores = torch.matmul(q, k.transpose(2, 3)) / (k.size(-1) ** 0.25)
        mask = torch.matmul(mask.float(), mask.transpose(1, 2).float()).bool()
        mask = mask.unsqueeze(1)
        mask = mask.repeat([1, self.head, 1, 1])
        scores.masked_fill_(~mask, -1e7)
        scores = F.softmax(scores, dim=2)
        scores = scores.transpose(2, 3)
        v_ = torch.matmul(scores, v)
        v_ = v_.transpose(1, 2).contiguous()
        v_ = merge_last(v_, 2)

        return v_


def split_last(x, shape):
    "split the last dimension to given shape"
    shape = list(shape)
    assert shape.count(-1) <= 1
    if -1 in shape:
        shape[shape.index(-1)] = int(x.size(-1) / -np.prod(shape))
    return x.view(*x.size()[:-1], *shape)


def merge_last(x, n_dims):
    "merge the last n_dims to a dimension"
    s = x.size()
    assert n_dims > 1 and n_dims < len(s)
    return x.view(*s[:-n_dims], -1)


class SublayerConnection(nn.Module):
    """
    A residual connection followed by a layer norm.
    Note for code simplicity the norm is first as opposed to last.
    """
    def __init__(self, dropout, dim):
        super(SublayerConnection, self).__init__()
        self.norm = LayerNorm(dim)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, sublayer):
        "Apply residual connection to any sublayer with the same size."
        return x + self.dropout(self.norm(sublayer(x)))


class LayerNorm(nn.Module):
    "Construct a layernorm module (See citation for details)."
    def __init__(self, features, eps=1e-6):
        super(LayerNorm, self).__init__()
        self.a_2 = nn.Parameter(torch.ones(features))
        self.b_2 = nn.Parameter(torch.zeros(features))
        self.eps = eps

    def forward(self, x):
        mean = x.mean(-1, keepdim=True)
        std = x.std(-1, keepdim=True)
        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2


def cal_fixed_pos_att(max_len, window_size):
    win = (window_size - 1) // 2
    weight = float(1 / window_size)
    attn_dict = dict()
    for sen_len in range(1, max_len+1):
        attn = np.eye(sen_len)
        if sen_len < window_size:
            attn_dict[sen_len] = attn
            continue
        for i in range(sen_len):
            attn[i, i-win:i+win+1] = weight
        attn[0, 0:win+1] = weight
        attn_dict[sen_len] = torch.FloatTensor(attn)

    return attn_dict


class PositionalAttCached(nn.Module):
    def __init__(self, d_model, pos_attns, max_len=5000):
        super(PositionalAttCached, self).__init__()
        # Compute the positional encodings once in log space.
        self.d_model = d_model
        self.pos_attns = pos_attns
        self.max_len = max_len

    def forward(self, x):
        shape = list(x.size())
        pos_attn = self.pos_attns[shape[1]]
        p_e = Variable(pos_attn, requires_grad=False).cuda()
        p_e = p_e.repeat([shape[0], 1, 1])
        return p_e

class PositionalEncoding(nn.Module):

    def __init__(self, d_model, max_len=5000):
        super(PositionalEncoding, self).__init__()
        # Compute the positional encodings once in log space.
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1)
        position = position * 1
        div_term = torch.exp(torch.arange(0, d_model, 2) *
                             -(math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0)
        self.register_buffer('pe', pe)

    def forward(self, x):
        shape = list(x.size())
        p_e = Variable(self.pe[:, :x.size(1)], requires_grad=False).cuda()
        p_e = p_e.repeat([shape[0], 1, 1])
        return p_e