import torch.nn as nn
class TransformerModel(nn.Module):

    def __init__(self, encoder, decoder):
        super().__init__()
        self._is_generation_fast = False
        self.encoder = encoder
        self.decoder = decoder
        
    def forward(self, src_tokens, src_lengths, prev_output_tokens):
        encoder_out, padding_mask = self.encoder(src_tokens, src_lengths)
        decoder_out = self.decoder(prev_output_tokens, encoder_out, padding_mask)
        return decoder_out

import numpy as np
import matplotlib.pyplot as plt

import torch
import torch.nn as nn
import torch.nn.functional as F
import math, copy, time
import seaborn as sns

from encoder_demos.demo_fairseq.models.fairseq_model import BaseFairseqModel, FairseqDecoder, FairseqEncoder
from encoder_demos.demo_fairseq.models.fairseq_incremental_decoder import FairseqIncrementalDecoder

import torch.nn as nn
class TransformerEncoder(nn.Module):
    """Transformer encoder."""

    def __init__(self, args, embed_tokens, left_pad=True):
        super().__init__()
        self.dropout = args.dropout
        self.fuse_dropout_add = args.fuse_dropout_add
        self.fuse_relu_dropout = args.fuse_relu_dropout

        embed_dim = embed_tokens.embedding_dim
        self.padding_idx = embed_tokens.padding_idx
        self.max_source_positions = args.max_source_positions

        self.embed_tokens = embed_tokens
        self.embed_scale = math.sqrt(embed_dim)
        self.embed_positions = PositionalEmbedding(
            args.max_source_positions, embed_dim, self.padding_idx,
            left_pad=left_pad,
            learned=args.encoder_learned_pos,
        ) if not args.no_token_positional_embeddings else None

        # 2️⃣.2 for문에 붙어있는 함수 실행
        self.layers = nn.ModuleList([])
        self.layers.extend([
            TransformerEncoderLayer(args)
            for i in range(args.encoder_layers)
        ])

        self.normalize = args.encoder_normalize_before
        if self.normalize:
            self.layer_norm = FusedLayerNorm(embed_dim) if args.fuse_layer_norm else nn.LayerNorm(embed_dim)

    # ```TransformerEncoder.forward()```가 호출됨
    def forward(self, src_tokens, src_lengths):
        # 1️⃣ embed tokens and positions
        x = self.embed_scale * self.embed_tokens(src_tokens)
        if self.embed_positions is not None:
            x += self.embed_positions(src_tokens)
        x = F.dropout(x, p=self.dropout, training=self.training)

        # B x T x C -> T x B x C
        # The tensor needs to copy transposed because
        # fused dropout is not capable of handing strided data
        if self.fuse_dropout_add :
            x = x.transpose(0, 1).contiguous()
        else :
            x = x.transpose(0, 1)

        # compute padding mask
        encoder_padding_mask = src_tokens.eq(self.padding_idx)
        if not encoder_padding_mask.any():
            _encoder_padding_mask = None
        else:
            _encoder_padding_mask = encoder_padding_mask

        # 2️⃣.1 encoder layers : self.layers를 시작하는 루프에서 6개의 TranformerEncoderLayers 각각을 통과
        for layer in self.layers:
            x = layer(x, _encoder_padding_mask)

        if self.normalize:
            x = self.layer_norm(x)

        return x, encoder_padding_mask # x.shape == T x B x C, encoder_padding_mask.shape == B x T

class TransformerEncoderLayer(nn.Module):
    """Encoder layer block.
    In the original paper each operation (multi-head attention or FFN) is
    postprocessed with: dropout -> add residual -> layernorm.
    In the tensor2tensor code they suggest that learning is more robust when
    preprocessing each layer with layernorm and postprocessing with:
    dropout -> add residual.
    We default to the approach in the paper, but the tensor2tensor approach can
    be enabled by setting `normalize_before=True`.
    """

    def __init__(self, args):
        super().__init__()
        self.embed_dim = args.encoder_embed_dim
        self.self_attn = MultiheadAttention(
            self.embed_dim, args.encoder_attention_heads,
            dropout=args.attention_dropout,
        )
        self.dropout = args.dropout
        self.relu_dropout = args.relu_dropout
        self.fuse_dropout_add = args.fuse_dropout_add
        self.fuse_relu_dropout = args.fuse_relu_dropout
        self.normalize_before = args.encoder_normalize_before
        self.fc1 = Linear(self.embed_dim, args.encoder_ffn_embed_dim)
        self.fc2 = Linear(args.encoder_ffn_embed_dim, self.embed_dim)
        self.layer_norms = nn.ModuleList([FusedLayerNorm(self.embed_dim) for i in range(2)])


    def forward(self, x, encoder_padding_mask):
        residual = x

        x = self.maybe_layer_norm(0, x, before=True)
        # 1️⃣ Self-Attention 통과
        x, _ = self.self_attn(query=x, key=x, value=x, key_padding_mask=encoder_padding_mask)
        # 2️⃣ 드롭아웃이 적용
        if self.fuse_dropout_add and self.training :
            x = jit_dropout_add(x, residual, self.dropout, self.training)
        else :
            x = F.dropout(x, p=self.dropout, training=self.training)
            x = residual + x
        x = self.maybe_layer_norm(0, x, after=True)

        residual = x
        x = self.maybe_layer_norm(1, x, before=True)

        # 3️⃣ 드롭아웃 X
        if self.fuse_relu_dropout :
            x = jit_relu_dropout(self.fc1(x), self.relu_dropout, self.training)
        else :
            x = F.threshold(self.fc1(x),0,0)
            x = F.dropout(x, p=self.relu_dropout, training=self.training)
        x = self.fc2(x)

        if self.fuse_dropout_add and self.training :
            x = jit_dropout_add(x, residual, self.dropout, self.training)
        else :
            x = F.dropout(x, p=self.dropout, training=self.training)
            x = residual + x
        x = self.maybe_layer_norm(1, x, after=True)
        return x

    def maybe_layer_norm(self, i, x, before=False, after=False):
        assert before ^ after
        if after ^ self.normalize_before:
            return self.layer_norms[i](x)
        else:
            return x

import encoder_demos.tokenize as tok

input_text = "I am looking for a place to eat."
tokens = tok.demo(input_text)
print("\nInput text:        %s\nTokenized output: " % input_text, tokens[0].numpy())

| [en] dictionary: 33712 types
| [de] dictionary: 33712 types

Input text:        I am looking for a place to eat.
Tokenized output:  [  40  105 1804   19   14  358   12 9637    5    2]

# Translate an input sentence in English to German

import encoder_demos.functional_translation as ft
input_sentence = "I am looking for a place to eat."

e, g, h = ft.demo(input_sentence)
print("En:", e)
print("German:", g)
print('')
print("H:", h)
print('H is the hypothesis along with an average log-likelihood')

| [en] dictionary: 33712 types
| [de] dictionary: 33712 types
| loading model(s) from /data/checkpoints/JoC_Transformer_FP32_PyT_20190304.pt
| Sentences are being padded to multiples of: 1
generated batches in  0.00028586387634277344 s
En: I am looking for a place to eat.
German: Ich bin auf der Suche nach einem Ort zu ße.

H: -0.3850874900817871
H is the hypothesis along with an average log-likelihood

print('Probability of H = {}'.format(round(np.exp(h),4)))

Probability of H = 0.6804

from torch.autograd import Variable
class PositionalEncoding(nn.Module):
    "Implement the PE function."
    def __init__(self, d_model, dropout, max_len=500):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        print("dropout:", dropout)
        # Compute the positional encodings once in log space.
        pe = torch.zeros(max_len, d_model)
        
        print("d_model:", d_model)
        position = torch.arange(0.0, max_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0.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)
        print("pe:", pe[:,0:2])
    
        self.register_buffer('pe', pe)
        
    def forward(self, x):
        
        x = x + Variable(self.pe[:, :x.size(1)], 
                         requires_grad=False)
        return self.dropout(x)

# visualization from harvardnlp/annotated-transformer GitHub (MIT license)
plt.figure(figsize=(15, 5))
pe = PositionalEncoding(20, 0)
y = pe.forward(Variable(torch.zeros(1, 100, 20)))
plt.plot(np.arange(100), y[0, :, 4:8].data.numpy())
_ = plt.legend(["dim %d"%p for p in [4,5,6,7]])

dropout: 0
d_model: 20
pe: tensor([[[0.0000e+00, 1.0000e+00, 0.0000e+00, 1.0000e+00, 0.0000e+00, 1.0000e+00, 0.0000e+00, 1.0000e+00, 0.0000e+00, 1.0000e+00, 0.0000e+00, 1.0000e+00, 0.0000e+00, 1.0000e+00, 0.0000e+00, 1.0000e+00, 0.0000e+00, 1.0000e+00, 0.0000e+00, 1.0000e+00],
         [8.4147e-01, 5.4030e-01, 3.8767e-01, 9.2180e-01, 1.5783e-01, 9.8747e-01, 6.3054e-02, 9.9801e-01, 2.5116e-02, 9.9968e-01, 9.9998e-03, 9.9995e-01, 3.9811e-03, 9.9999e-01, 1.5849e-03, 1.0000e+00, 6.3096e-04, 1.0000e+00, 2.5119e-04, 1.0000e+00]]])

import encoder_demos.self_attention as sa
def normalize(arr):
    out = np.zeros_like(arr)
    for rowno in range(arr.shape[0]):
        vals = arr[rowno, :]
        out[rowno, :] = (vals - np.min(vals)) / (np.max(vals) - np.min(vals))
    return out

sentence = "I am looking for a place to eat."
#sentence = "This is a much more complex sentence and, as a result, is much longer."
attn, attn_weights = sa.demo(sentence, return_early='self_attn')
a_w =attn_weights[0,:, :].cpu().numpy()  #you can print the attention weights.

sentence += " EOS"
sentence = sentence.replace(".", " .").replace(",", " ,").split(" ")

#a_w = normalize(a_w)

sns.set()
plt.figure(figsize=(15,15))
ax = sns.heatmap(a_w, vmin=0, vmax=1, cmap=sns.diverging_palette(200,10, n=200), xticklabels=sentence, square=True, yticklabels=sentence, cbar=True,cbar_kws={"shrink": .82})
ax.xaxis.set_ticks_position('top')
plt.yticks(rotation=0)
plt.show()

| [en] dictionary: 33712 types
| [de] dictionary: 33712 types
| loading model(s) from /data/checkpoints/JoC_Transformer_FP32_PyT_20190304.pt
| Sentences are being padded to multiples of: 1
generated batches in  0.00019598007202148438 s  

sentence = "I am looking for a place to eat."
#sentence = "This is a much more complex sentence and, as a result, is much longer."
attn, attn_weights = sa.demo(sentence, return_early='self_attn')
a_w =attn_weights[0,:, :].cpu().numpy()  #you can print the attention weights.

sentence += " EOS"
sentence = sentence.replace(".", " .").replace(",", " ,").split(" ")

a_w = normalize(a_w)

sns.set()
plt.figure(figsize=(15,15))
ax = sns.heatmap(a_w, vmin=0, vmax=1, cmap=sns.diverging_palette(200,10, n=200), xticklabels=sentence, square=True, yticklabels=sentence, cbar=True,cbar_kws={"shrink": .82})
ax.xaxis.set_ticks_position('top')
plt.yticks(rotation=0)
plt.show()

| [en] dictionary: 33712 types
| [de] dictionary: 33712 types
| loading model(s) from /data/checkpoints/JoC_Transformer_FP32_PyT_20190304.pt
| Sentences are being padded to multiples of: 1
generated batches in  0.00020933151245117188 s

# Import libraries
import os
import io
import sys
import PIL

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns 

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from typing import Optional, Dict
import math, copy, time

from encoder_demos.demo_fairseq.models.fairseq_model import BaseFairseqModel, FairseqDecoder, FairseqEncoder
from encoder_demos.demo_fairseq.models.fairseq_incremental_decoder import FairseqIncrementalDecoder

from torch import Tensor
from typing import Optional, Dict
class TransformerDecoder(FairseqIncrementalDecoder):
    """Transformer decoder."""

    def __init__(self, args, embed_tokens, no_encoder_attn=False, left_pad=False):
        super().__init__()
        self.dropout = args.dropout
        self.share_input_output_embed = args.share_decoder_input_output_embed
        self.fuse_dropout_add = args.fuse_dropout_add
        self.fuse_relu_dropout = args.fuse_relu_dropout

        embed_dim = embed_tokens.embedding_dim
        padding_idx = embed_tokens.padding_idx
        self.max_target_positions = args.max_target_positions

        self.embed_tokens = embed_tokens
        self.embed_scale = math.sqrt(embed_dim)
        self.embed_positions = PositionalEmbedding(
            args.max_target_positions, embed_dim, padding_idx,
            left_pad=left_pad,
            learned=args.decoder_learned_pos,
        ) if not args.no_token_positional_embeddings else None

        self.layers = nn.ModuleList([])
        self.layers.extend([
            TransformerDecoderLayer(args, no_encoder_attn)
            for _ in range(args.decoder_layers)
        ])


    def forward(self,
                prev_output_tokens: Tensor,
                encoder_out: Tensor,
                encoder_padding_mask: Tensor,
                incremental_state: Optional[Dict[str, Dict[str, Tensor]]]=None):
        # embed positions
        positions = self.embed_positions(
            prev_output_tokens,
            incremental_state=incremental_state,
        ) if self.embed_positions is not None else None

        if incremental_state is not None:
            prev_output_tokens = prev_output_tokens[:, -1:]
            if positions is not None:
                positions = positions[:, -1:]

        # embed tokens and positions
        x = self.embed_scale * self.embed_tokens(prev_output_tokens)
        if positions is not None:
            x += positions
        x = F.dropout(x, p=self.dropout, training=self.training)

        # B x T x C -> T x B x C
        # The tensor needs to copy transposed because
        # fused dropout is not capable of handing strided data
        if self.fuse_dropout_add :
            x = x.transpose(0, 1).contiguous()
        else :
            x = x.transpose(0, 1)
        attn = None

        # decoder layers
        for layer in self.layers:
            x, attn = layer(
                x,
                encoder_out,
                encoder_padding_mask if encoder_padding_mask.any() else None,
                incremental_state,
            )


        return x, attn

def subsequent_mask(size):
    "Mask out subsequent positions."
    attn_shape = (1, size, size)
    subsequent_mask = np.tril(np.ones((1,10,10))).astype('uint8')    
    return torch.as_tensor(subsequent_mask)

subsequent_mask(10)[0]

tensor([[1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
        [1, 1, 0, 0, 0, 0, 0, 0, 0, 0],
        [1, 1, 1, 0, 0, 0, 0, 0, 0, 0],
        [1, 1, 1, 1, 0, 0, 0, 0, 0, 0],
        [1, 1, 1, 1, 1, 0, 0, 0, 0, 0],
        [1, 1, 1, 1, 1, 1, 0, 0, 0, 0],
        [1, 1, 1, 1, 1, 1, 1, 0, 0, 0],
        [1, 1, 1, 1, 1, 1, 1, 1, 0, 0],
        [1, 1, 1, 1, 1, 1, 1, 1, 1, 0],
        [1, 1, 1, 1, 1, 1, 1, 1, 1, 1]], dtype=torch.uint8)

plt.figure(figsize=(5,5))
cmap = plt.cm.GnBu_r
plt.imshow(subsequent_mask(10)[0], cmap=cmap)
None