<filename>nmtpytorch/layers/decoders/multisource_transformer.py
from copy import deepcopy
import torch
from torch import nn

from ..transformer import (
    PositionalEncoding, 
    MultiHeadAttention, 
    MultiSourceMultiHeadAttention, 
    PositionwiseFeedForward
)
from ..transformer import LabelSmoothingLoss
from ..transformer import get_embs, get_pad_attn_mask

from .transformer import DecoderLayer as TextDecoderLayer

class DecoderLayer(nn.Module):
    def __init__(self, model_dim, feat_dim, feat_ratio, key_size, n_heads, inner_dim, dropout, dropnet):
        super(DecoderLayer, self).__init__()
        
        self.self_attn = MultiHeadAttention(model_dim, n_heads, dropout=dropout)
        self.ctx_attn = MultiSourceMultiHeadAttention(model_dim, feat_dim, feat_ratio, n_heads, 
            dropout=dropout, dropnet=dropnet)
        self.ff = PositionwiseFeedForward(model_dim, inner_dim, dropout=dropout)

        self.layer_norm = nn.LayerNorm(model_dim)

    def forward(self, hs, enc_hs, feat, self_pad_attn_mask, ctx_pad_attn_mask, feat_pad_attn_mask=None):
        h1, self_attn = self.self_attn(hs, hs, hs, self_pad_attn_mask)
        h1 = self.layer_norm(h1 + hs)

        h2, ctx_attn = self.ctx_attn(h1, enc_hs, enc_hs, ctx_pad_attn_mask, 
            feat, feat, feat_pad_attn_mask)
        h2 = self.layer_norm(h2 + h1)

        h3 = self.ff(h2)
        h3 = self.layer_norm(h3 + h2)

        return h3, self_attn, ctx_attn

class MultiSourceTransformerDecoder(nn.Module):
    def __init__(self, n_vocab, model_dim, max_len, layer_types, key_size, n_heads, 
        inner_dim, dropout, dropnet, tied_emb_proj, eps, ctx_name, 
        feat_name, feat_dim, feat_ratio, **kwargs):
        super().__init__()

        self.embs = get_embs(n_vocab, max_len, model_dim, dropout)

        def get_text_layer():
            return TextDecoderLayer(
                model_dim = model_dim, 
                key_size = key_size, 
                n_heads = n_heads, 
                inner_dim = inner_dim,
                dropout = dropout)
        
        def get_fused_layer():
            return DecoderLayer(
                model_dim = model_dim,
                feat_dim = feat_dim,
                feat_ratio = feat_ratio,
                key_size = key_size, 
                n_heads = n_heads, 
                inner_dim = inner_dim,
                dropout = dropout,
                dropnet = dropnet,)

        get_layer = {'t': get_text_layer, 'f': get_fused_layer}

        self.layer_types = layer_types.lower()
        self.layers = nn.ModuleList([get_layer[t]() for t in self.layer_types])

        self.hs2prob = nn.Linear(model_dim, n_vocab)

        # shares embedding and projection weight
        if tied_emb_proj:
            self.hs2prob.weight = self.embs[0].weight
            self.logit_scale = model_dim ** -0.5
        else:
            self.logit_scale = 1.0

        self.log_softmax = nn.LogSoftmax(dim=-1)
        self.loss_fn = LabelSmoothingLoss(eps, n_vocab, padding_index=0)

        self.ctx_name = ctx_name
        self.feat_name = feat_name

    def reset_parameters(self, **kwargs):
        nn.init.normal_(self.embs[0].weight)
        for layer in self.layers:
            nn.init.xavier_normal_(layer.ff.w_1.weight)
            nn.init.xavier_normal_(layer.ff.w_2.weight)
            nn.init.zeros_(layer.ff.w_1.bias)
            nn.init.zeros_(layer.ff.w_2.bias)

    def get_emb(self, idxs, tstep):
        # Only used by search.py
        # idxs: eval_batch_size

        # ret: 1 x batch_size
        return idxs.unsqueeze(0)

    def f_init(self, ctx_dict):
        # Actually initial hidden state is no longer used in Transformer
        # Because transformer decoder needs target outputs so far that search.py does not provide,
        # we pass it by parameter 'h' in f_next function, which is designed to hold rnn state.
        enc_hs, x = ctx_dict[self.ctx_name]
        batch_size = x.size(1)
        return torch.LongTensor(batch_size, 0).to(x.device)

    def f_next(self, ctx_dict, y, h, training=False):
        # ctx_dict: enc_hs (encoder states), x (source sentence)
        # y: batch_size, len_y (default in nmtpytorch)
        # h: None (training), target sentence so far (test)
        y = y.transpose(0, 1)
        if type(h) != type(None): # test
            y = torch.cat([h, y], dim=-1)
        hs = self.embs(y)

        enc_hs, x = ctx_dict[self.ctx_name]
        x = x.transpose(0, 1) # len_x x batch_size
        enc_hs = enc_hs.transpose(0, 1) # len_x x batch_size x model_dim

        #
        feat, feat_mask = ctx_dict[self.feat_name]
        feat = feat.transpose(0, 1)
        feat_mask = feat_mask.transpose(0, 1)
        
        self_pad_attn_mask = get_pad_attn_mask(y, y, masking_seq=True)
        ctx_pad_attn_mask  = get_pad_attn_mask(y, x)
        feat_pad_attn_mask = get_pad_attn_mask(y, feat_mask)

        self_attns, ctx_attns = [], []
        for t, layer in zip(self.layer_types, self.layers):
            hs, self_attn, ctx_attn = layer(
                hs=hs, 
                enc_hs=enc_hs, 
                feat=feat, 
                self_pad_attn_mask=self_pad_attn_mask, 
                ctx_pad_attn_mask=ctx_pad_attn_mask, 
                feat_pad_attn_mask=feat_pad_attn_mask
            )

            self_attns.append(self_attn)
            ctx_attns.append(ctx_attn)

        logit = self.hs2prob(hs)
        if not training:
            logit = logit[:, -1, :]

        logit = logit * self.logit_scale
        log_p = self.log_softmax(logit)

        return log_p, y

    def forward(self, ctx_dict, y):
        # y: len_y x batch_size (default in nmtpytorch)
        # len_y, batch_size= y.size()

        log_p, _ = self.f_next(ctx_dict, y[:-1, :], None, training=True)
        loss = self.loss_fn(log_p, y[1:, :].transpose(0, 1).contiguous().view(-1))

        return {'loss': loss}
<filename>nmtpytorch/layers/decoders/multisource_transformer.py
from copy import deepcopy
import torch
from torch import nn

from ..transformer import (
    PositionalEncoding, 
    MultiHeadAttention, 
    MultiSourceMultiHeadAttention, 
    PositionwiseFeedForward
)
from ..transformer import LabelSmoothingLoss
from ..transformer import get_embs, get_pad_attn_mask

from .transformer import DecoderLayer as TextDecoderLayer

class DecoderLayer(nn.Module):
    def __init__(self, model_dim, feat_dim, feat_ratio, key_size, n_heads, inner_dim, dropout, dropnet):
        super(DecoderLayer, self).__init__()
        
        self.self_attn = MultiHeadAttention(model_dim, n_heads, dropout=dropout)
        self.ctx_attn = MultiSourceMultiHeadAttention(model_dim, feat_dim, feat_ratio, n_heads, 
            dropout=dropout, dropnet=dropnet)
        self.ff = PositionwiseFeedForward(model_dim, inner_dim, dropout=dropout)

        self.layer_norm = nn.LayerNorm(model_dim)

    def forward(self, hs, enc_hs, feat, self_pad_attn_mask, ctx_pad_attn_mask, feat_pad_attn_mask=None):
        h1, self_attn = self.self_attn(hs, hs, hs, self_pad_attn_mask)
        h1 = self.layer_norm(h1 + hs)

        h2, ctx_attn = self.ctx_attn(h1, enc_hs, enc_hs, ctx_pad_attn_mask, 
            feat, feat, feat_pad_attn_mask)
        h2 = self.layer_norm(h2 + h1)

        h3 = self.ff(h2)
        h3 = self.layer_norm(h3 + h2)

        return h3, self_attn, ctx_attn

class MultiSourceTransformerDecoder(nn.Module):
    def __init__(self, n_vocab, model_dim, max_len, layer_types, key_size, n_heads, 
        inner_dim, dropout, dropnet, tied_emb_proj, eps, ctx_name, 
        feat_name, feat_dim, feat_ratio, **kwargs):
        super().__init__()

        self.embs = get_embs(n_vocab, max_len, model_dim, dropout)

        def get_text_layer():
            return TextDecoderLayer(
                model_dim = model_dim, 
                key_size = key_size, 
                n_heads = n_heads, 
                inner_dim = inner_dim,
                dropout = dropout)
        
        def get_fused_layer():
            return DecoderLayer(
                model_dim = model_dim,
                feat_dim = feat_dim,
                feat_ratio = feat_ratio,
                key_size = key_size, 
                n_heads = n_heads, 
                inner_dim = inner_dim,
                dropout = dropout,
                dropnet = dropnet,)

        get_layer = {'t': get_text_layer, 'f': get_fused_layer}

        self.layer_types = layer_types.lower()
        self.layers = nn.ModuleList([get_layer[t]() for t in self.layer_types])

        self.hs2prob = nn.Linear(model_dim, n_vocab)

        # shares embedding and projection weight
        if tied_emb_proj:
            self.hs2prob.weight = self.embs[0].weight
            self.logit_scale = model_dim ** -0.5
        else:
            self.logit_scale = 1.0

        self.log_softmax = nn.LogSoftmax(dim=-1)
        self.loss_fn = LabelSmoothingLoss(eps, n_vocab, padding_index=0)

        self.ctx_name = ctx_name
        self.feat_name = feat_name

    def reset_parameters(self, **kwargs):
        nn.init.normal_(self.embs[0].weight)
        for layer in self.layers:
            nn.init.xavier_normal_(layer.ff.w_1.weight)
            nn.init.xavier_normal_(layer.ff.w_2.weight)
            nn.init.zeros_(layer.ff.w_1.bias)
            nn.init.zeros_(layer.ff.w_2.bias)

    def get_emb(self, idxs, tstep):
        # Only used by search.py
        # idxs: eval_batch_size

        # ret: 1 x batch_size
        return idxs.unsqueeze(0)

    def f_init(self, ctx_dict):
        # Actually initial hidden state is no longer used in Transformer
        # Because transformer decoder needs target outputs so far that search.py does not provide,
        # we pass it by parameter 'h' in f_next function, which is designed to hold rnn state.
        enc_hs, x = ctx_dict[self.ctx_name]
        batch_size = x.size(1)
        return torch.LongTensor(batch_size, 0).to(x.device)

    def f_next(self, ctx_dict, y, h, training=False):
        # ctx_dict: enc_hs (encoder states), x (source sentence)
        # y: batch_size, len_y (default in nmtpytorch)
        # h: None (training), target sentence so far (test)
        y = y.transpose(0, 1)
        if type(h) != type(None): # test
            y = torch.cat([h, y], dim=-1)
        hs = self.embs(y)

        enc_hs, x = ctx_dict[self.ctx_name]
        x = x.transpose(0, 1) # len_x x batch_size
        enc_hs = enc_hs.transpose(0, 1) # len_x x batch_size x model_dim

        #
        feat, feat_mask = ctx_dict[self.feat_name]
        feat = feat.transpose(0, 1)
        feat_mask = feat_mask.transpose(0, 1)
        
        self_pad_attn_mask = get_pad_attn_mask(y, y, masking_seq=True)
        ctx_pad_attn_mask  = get_pad_attn_mask(y, x)
        feat_pad_attn_mask = get_pad_attn_mask(y, feat_mask)

        self_attns, ctx_attns = [], []
        for t, layer in zip(self.layer_types, self.layers):
            hs, self_attn, ctx_attn = layer(
                hs=hs, 
                enc_hs=enc_hs, 
                feat=feat, 
                self_pad_attn_mask=self_pad_attn_mask, 
                ctx_pad_attn_mask=ctx_pad_attn_mask, 
                feat_pad_attn_mask=feat_pad_attn_mask
            )

            self_attns.append(self_attn)
            ctx_attns.append(ctx_attn)

        logit = self.hs2prob(hs)
        if not training:
            logit = logit[:, -1, :]

        logit = logit * self.logit_scale
        log_p = self.log_softmax(logit)

        return log_p, y

    def forward(self, ctx_dict, y):
        # y: len_y x batch_size (default in nmtpytorch)
        # len_y, batch_size= y.size()

        log_p, _ = self.f_next(ctx_dict, y[:-1, :], None, training=True)
        loss = self.loss_fn(log_p, y[1:, :].transpose(0, 1).contiguous().view(-1))

        return {'loss': loss}
# shares embedding and projection weight # Only used by search.py # idxs: eval_batch_size # ret: 1 x batch_size # Actually initial hidden state is no longer used in Transformer # Because transformer decoder needs target outputs so far that search.py does not provide, # we pass it by parameter 'h' in f_next function, which is designed to hold rnn state. # ctx_dict: enc_hs (encoder states), x (source sentence) # y: batch_size, len_y (default in nmtpytorch) # h: None (training), target sentence so far (test) # test # len_x x batch_size # len_x x batch_size x model_dim # # y: len_y x batch_size (default in nmtpytorch) # len_y, batch_size= y.size()