Annotated_Transformer_English_to_Chinese_Translator.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Annotated_Transformer_English_to_Chinese_Translator\n",
    "\n",
    "```python\n",
    "Created on   :2020/09/19 23:27:28\n",
    "@author      :Caihao (Chris) Cui\n",
    "@file        :Annotated_Transformer_English_to_Chinese_Translator.ipynb\n",
    "@content     :Transformer English-to-Chinese.\n",
    "@version     :0.1.0\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook, I will build a Transformer model (neural network model) to translate English sentence to Chinese sentences.\n",
    "\n",
    "The reference paper is [Attention Is All You Need](https://arxiv.org/pdf/1706.03762.pdf)\n",
    "\n",
    "![](document/images/google_translate.png )\n",
    "\n",
    "The full data set only contains around 10,000 sentence pairs. But this is very good practice to imporve your python programming skills and get deeper understanding of pytorch and Transformaer(neural networks).\n",
    "\n",
    "![Attention Is All You Need](document/images/google_paper.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Understand the Transformer Model\n",
    " \n",
    "The Whole Transformer  encoder-decoder model architecture  services for the following purposes. \n",
    "\n",
    "\n",
    "- Encoder(s): the encoding process transforms the input sentence (list of English words) into numeric matrix format (embedding ), consider this step is to extract useful and necessary information for the decoder. In Fig 06, the embedding is represented by the green matrix.\n",
    "\n",
    "- Decoder(s): then the decoding process mapping these embeddings back to another language sequence as Fig 06 shown, which helps us to solve all kinds of supervised NLP tasks, like machine translation (in this blog), sentiment classification, entity recognition,  summary generation, semantic relation extraction and so on. \n",
    " \n",
    "![Understand_How_Transformer_Work](./document/images/Understand_How_Transformer_Work.png)\n",
    "\n",
    " "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Encoder \n",
    "  \n",
    "We will focus on the structure of the encoder in this section, because after understanding the structure of the encoder, understanding the decoder will be very simple. Moreover we can just use the encoder to complete some of the mainstream tasks in NLP, such as sentiment classification, semantic relationship analysis, named entity recognition and so on.\n",
    "\n",
    "Recall that the Encoder denotes the process of mapping natural language sequences to mathematical expressions to hidden layers outputs."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Here is a Transformer Encoder Block structure**\n",
    "> Notification: the following sections will refer to the 1,2,3,4 blocks.\n",
    "\n",
    "![Transformer Encoder Stacks](./document/images/encoder.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.0 Data Preparation: English-to-Chinese Translator Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T00:27:00.977365Z",
     "start_time": "2020-09-20T00:27:00.835708Z"
    }
   },
   "outputs": [],
   "source": [
    "!tree data/nmt/en-cn "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T00:27:00.985865Z",
     "start_time": "2020-09-20T00:27:00.980036Z"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import math\n",
    "import copy\n",
    "import time\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from nltk import word_tokenize\n",
    "from collections import Counter\n",
    "from torch.autograd import Variable\n",
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T00:27:00.994553Z",
     "start_time": "2020-09-20T00:27:00.989214Z"
    }
   },
   "outputs": [],
   "source": [
    "# init parameters\n",
    "UNK = 0  # unknow word-id\n",
    "PAD = 1  # padding word-id\n",
    "BATCH_SIZE = 64   \n",
    "\n",
    "DEBUG = True    # Debug / Learning Purposes. \n",
    "# DEBUG = False # Build the model, better with GPU CUDA enabled.\n",
    "\n",
    "if DEBUG:        \n",
    "    EPOCHS  = 2   \n",
    "    LAYERS  = 3   \n",
    "    H_NUM   = 8    \n",
    "    D_MODEL = 128   \n",
    "    D_FF    = 256      \n",
    "    DROPOUT = 0.1   \n",
    "    MAX_LENGTH = 60   \n",
    "    TRAIN_FILE = 'data/nmt/en-cn/train_mini.txt'   \n",
    "    DEV_FILE   = 'data/nmt/en-cn/dev_mini.txt'   \n",
    "    SAVE_FILE  = 'save/models/model.pt'   \n",
    "else:\n",
    "    EPOCHS  = 20  \n",
    "    LAYERS  = 6    \n",
    "    H_NUM   = 8    \n",
    "    D_MODEL = 256    \n",
    "    D_FF    = 1024   \n",
    "    DROPOUT = 0.1    \n",
    "    MAX_LENGTH = 60  \n",
    "    TRAIN_FILE = 'data/nmt/en-cn/train.txt'   \n",
    "    DEV_FILE   = 'data/nmt/en-cn/dev.txt'   \n",
    "    SAVE_FILE  = 'save/models/large_model.pt'  \n",
    "\n",
    "DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T00:27:01.012800Z",
     "start_time": "2020-09-20T00:27:00.996896Z"
    }
   },
   "outputs": [],
   "source": [
    "def seq_padding(X, padding=0):\n",
    "    \"\"\"\n",
    "    add padding to a batch data \n",
    "    \"\"\"\n",
    "    L = [len(x) for x in X]\n",
    "    ML = max(L)\n",
    "    return np.array([\n",
    "        np.concatenate([x, [padding] * (ML - len(x))]) if len(x) < ML else x for x in X\n",
    "    ])    \n",
    "\n",
    "class PrepareData:\n",
    "    def __init__(self, train_file, dev_file):\n",
    "        # 01. Read the data and tokenize\n",
    "        self.train_en, self.train_cn = self.load_data(train_file)\n",
    "        self.dev_en, self.dev_cn     = self.load_data(dev_file)\n",
    "\n",
    "        # 02. build dictionary: English and Chinese\n",
    "        self.en_word_dict, self.en_total_words, self.en_index_dict = self.build_dict(self.train_en)\n",
    "        self.cn_word_dict, self.cn_total_words, self.cn_index_dict = self.build_dict(self.train_cn)\n",
    "\n",
    "        # 03. word to id by dictionary \n",
    "        self.train_en, self.train_cn = self.wordToID(self.train_en, self.train_cn, self.en_word_dict, self.cn_word_dict)\n",
    "        self.dev_en, self.dev_cn     = self.wordToID(self.dev_en, self.dev_cn, self.en_word_dict, self.cn_word_dict)\n",
    "\n",
    "        # 04. batch + padding + mask\n",
    "        self.train_data = self.splitBatch(self.train_en, self.train_cn, BATCH_SIZE)\n",
    "        self.dev_data   = self.splitBatch(self.dev_en, self.dev_cn, BATCH_SIZE)\n",
    "\n",
    "    def load_data(self, path):\n",
    "        \"\"\"\n",
    "        Read English and Chinese Data \n",
    "        tokenize the sentence and add start/end marks(Begin of Sentence; End of Sentence)\n",
    "        en = [['BOS', 'i', 'love', 'you', 'EOS'], \n",
    "              ['BOS', 'me', 'too', 'EOS'], ...]\n",
    "        cn = [['BOS', '我', '爱', '你', 'EOS'], \n",
    "              ['BOS', '我', '也', '是', 'EOS'], ...]\n",
    "        \"\"\"\n",
    "        en = []\n",
    "        cn = []\n",
    "        with open(path, 'r', encoding='utf-8') as f:\n",
    "            for line in f:\n",
    "                line = line.strip().split('\\t')\n",
    "                en.append([\"BOS\"] + word_tokenize(line[0].lower()) + [\"EOS\"])\n",
    "                cn.append([\"BOS\"] + word_tokenize(\" \".join([w for w in line[1]])) + [\"EOS\"])\n",
    "        return en, cn\n",
    "    \n",
    "    def build_dict(self, sentences, max_words = 50000):\n",
    "        \"\"\"\n",
    "        sentences: list of word list \n",
    "        build dictonary as {key(word): value(id)}\n",
    "        \"\"\"\n",
    "        word_count = Counter()\n",
    "        for sentence in sentences:\n",
    "            for s in sentence:\n",
    "                word_count[s] += 1\n",
    "\n",
    "        ls = word_count.most_common(max_words)\n",
    "        total_words = len(ls) + 2\n",
    "        word_dict = {w[0]: index + 2 for index, w in enumerate(ls)}\n",
    "        word_dict['UNK'] = UNK\n",
    "        word_dict['PAD'] = PAD\n",
    "        # inverted index: {key(id): value(word)}\n",
    "        index_dict = {v: k for k, v in word_dict.items()}\n",
    "        return word_dict, total_words, index_dict\n",
    "\n",
    "    def wordToID(self, en, cn, en_dict, cn_dict, sort=True):\n",
    "        \"\"\"\n",
    "        convert input/output word lists to id lists. \n",
    "        Use input word list length to sort, reduce padding.\n",
    "        \"\"\"\n",
    "        length = len(en)\n",
    "        out_en_ids = [[en_dict.get(w, 0) for w in sent] for sent in en]\n",
    "        out_cn_ids = [[cn_dict.get(w, 0) for w in sent] for sent in cn]\n",
    "\n",
    "        def len_argsort(seq):\n",
    "            \"\"\"\n",
    "            get sorted index w.r.t length.\n",
    "            \"\"\"\n",
    "            return sorted(range(len(seq)), key=lambda x: len(seq[x]))\n",
    "\n",
    "        if sort: # update index\n",
    "            sorted_index = len_argsort(out_en_ids) # English\n",
    "            out_en_ids = [out_en_ids[id] for id in sorted_index]\n",
    "            out_cn_ids = [out_cn_ids[id] for id in sorted_index]\n",
    "        return out_en_ids, out_cn_ids\n",
    "\n",
    "    def splitBatch(self, en, cn, batch_size, shuffle=True):\n",
    "        \"\"\"\n",
    "        get data into batches\n",
    "        \"\"\"\n",
    "        idx_list = np.arange(0, len(en), batch_size)\n",
    "        if shuffle:\n",
    "            np.random.shuffle(idx_list)\n",
    "\n",
    "        batch_indexs = []\n",
    "        for idx in idx_list:\n",
    "            batch_indexs.append(np.arange(idx, min(idx + batch_size, len(en))))\n",
    "        \n",
    "        batches = []\n",
    "        for batch_index in batch_indexs:\n",
    "            batch_en = [en[index] for index in batch_index]  \n",
    "            batch_cn = [cn[index] for index in batch_index]\n",
    "            # paddings: batch, batch_size, batch_MaxLength\n",
    "            batch_cn = seq_padding(batch_cn)\n",
    "            batch_en = seq_padding(batch_en)\n",
    "            batches.append(Batch(batch_en, batch_cn)) \n",
    "            #!!! 'Batch' Class is called here but defined in later section.\n",
    "        return batches"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> Notification, the Batch class relates to the Attention/Mask in Encoder class. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Input/Output Embeddings\n",
    "Similary to all sequential model, we used learned embedding to convert the input/output vectors' dimensionality to $d_{model}$.\n",
    "In our model, the two embedding layers and pre-softmax layer will share weight matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T00:27:01.020560Z",
     "start_time": "2020-09-20T00:27:01.016575Z"
    }
   },
   "outputs": [],
   "source": [
    "class Embeddings(nn.Module):\n",
    "    def __init__(self, d_model, vocab):\n",
    "        super(Embeddings, self).__init__() \n",
    "        self.lut = nn.Embedding(vocab, d_model)\n",
    "        self.d_model = d_model\n",
    "\n",
    "    def forward(self, x):\n",
    "        # return x's embedding vector（times math.sqrt(d_model)）\n",
    "        return self.lut(x) * math.sqrt(self.d_model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we have all the code for data preprocessing. Let's focus on the understand and build Transformer mode. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 Positional Encoding"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-19T02:20:41.866079Z",
     "start_time": "2020-09-19T02:20:41.858378Z"
    }
   },
   "source": [
    "The $Transformer$ does **not** contain iteration operation like RNN or LSTM in encoders, so we have to offer the position information of the words to the model, so the model learns the order in the input sequence.  \n",
    "\n",
    "Thus, we define the **positional encoding** as [max_sequence_length, embedding_dimension]\n",
    "\n",
    "In the paper, we use sine and cosine function to provide the position information.\n",
    "\n",
    "$$PE_{(pos,2i)} = sin(pos / 10000^{2i/d_{\\text{model}}}) \\quad \\quad PE_{(pos,2i+1)} = cos(pos / 10000^{2i/d_{\\text{model}}})\\tag{eq.1}$$  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T00:27:01.028057Z",
     "start_time": "2020-09-20T00:27:01.022506Z"
    }
   },
   "outputs": [],
   "source": [
    "class PositionalEncoding(nn.Module):\n",
    "    def __init__(self, d_model, dropout, max_len=5000):\n",
    "        super(PositionalEncoding, self).__init__()\n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "                                  \n",
    "        pe = torch.zeros(max_len, d_model, device=DEVICE)\n",
    "        position = torch.arange(0., max_len, device=DEVICE).unsqueeze(1)\n",
    "        div_term = torch.exp(torch.arange(0., d_model, 2, device=DEVICE) * -(math.log(10000.0) / d_model))\n",
    "        pe_pos   = torch.mul(position, div_term)\n",
    "        pe[:, 0::2] = torch.sin(pe_pos)\n",
    "        pe[:, 1::2] = torch.cos(pe_pos)\n",
    "        pe = pe.unsqueeze(0) \n",
    "                                  \n",
    "        self.register_buffer('pe', pe) # pe\n",
    "\n",
    "    def forward(self, x):\n",
    "        #  build pe w.r.t to the max_length\n",
    "        x = x + Variable(self.pe[:, :x.size(1)], requires_grad=False)\n",
    "        return self.dropout(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "See, we first build the position encoding based on x and then add the 'pe' to the x in the forward function.\n",
    "\n",
    "> Notification: Set 'requires_grad=False'，because we do not need to train pe."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here are the position embedding visualisations, you can find the pattern changes with the increasing embedding dimensions. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T00:27:01.923703Z",
     "start_time": "2020-09-20T00:27:01.029961Z"
    }
   },
   "outputs": [],
   "source": [
    "pe = PositionalEncoding(32, 0, 100)  # d_model, dropout-ratio, max_len\n",
    "positional_encoding = pe.forward(Variable(torch.zeros(1, 100, 32))) # sequence length, d_model\n",
    "plt.figure(figsize=(10,10))\n",
    "sns.heatmap(positional_encoding.squeeze()) # 100x32 matrix\n",
    "plt.title(\"Sinusoidal Function\")\n",
    "plt.xlabel(\"hidden dimension\")\n",
    "plt.ylabel(\"sequence length\")\n",
    "None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T00:27:02.109211Z",
     "start_time": "2020-09-20T00:27:01.927376Z"
    }
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(15, 5))\n",
    "pe = PositionalEncoding(24, 0)\n",
    "y = pe.forward(Variable(torch.zeros(1, 100, 24)))\n",
    "plt.plot(np.arange(100), y[0, :, 5:10].data.numpy())\n",
    "plt.legend([\"dim %d\"%p for p in [5,6,7,8,9]])\n",
    "None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2 Self Attention and Mask"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "An attention function can be described as mapping a query and a set of key-value pairs to an output, where the query, keys, values, and output are all vectors. The output is computed as a weighted sum of the values, where the weight assigned to each value is computed by a compatibility function of the query with the corresponding key.\n",
    "\n",
    "We call our particular attention “Scaled Dot-Product Attention”. The input consists of queries and keys of dimension $d_k$, and values of dimension $d_v$. \n",
    "\n",
    "We compute the dot products of the query with all keys, divide each by $\\sqrt{d_k}$, and apply a softmax function to obtain the weights on the values.\n",
    "\n",
    "![self_attention](document/images/self-attention.png)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " The two most commonly used attention functions are additive attention, and dot-product (multiplicative) attention. Dot-product attention is identical to our algorithm, except for the scaling factor of $\\frac{1}{\\sqrt{d_k}}$\n",
    ". Additive attention computes the compatibility function using a feed-forward network with a single hidden layer. While the two are similar in theoretical complexity, dot-product attention is much faster and more space-efficient in practice, since it can be implemented using highly optimized matrix multiplication code.\n",
    "\n",
    "To illustrate why the dot products get large, assume that the components of q and k are independent random variables with mean 0  and variance 1 . Then their dot product, $q\\cdot k $ has mean 0 and variance $d_k$,To counteract this effect, we scale the dot products by  $\\frac{1}{\\sqrt{d_k}}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T00:27:02.115993Z",
     "start_time": "2020-09-20T00:27:02.112265Z"
    }
   },
   "outputs": [],
   "source": [
    "def attention(query, key, value, mask=None, dropout=None):\n",
    "    \"Compute 'Scaled Dot Product Attention'\"\n",
    "    d_k = query.size(-1) \n",
    "    scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)\n",
    "    if mask is not None:\n",
    "        scores = scores.masked_fill(mask == 0, -1e9)\n",
    "    p_attn = F.softmax(scores, dim=-1)\n",
    "    if dropout is not None:\n",
    "        p_attn = dropout(p_attn)\n",
    "    return torch.matmul(p_attn, value), p_attn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Multi-head attention allows the model to jointly attend to information from different representation subspaces at different positions. With a single attention head, averaging inhibits this."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T00:27:02.122947Z",
     "start_time": "2020-09-20T00:27:02.117596Z"
    }
   },
   "outputs": [],
   "source": [
    "class MultiHeadedAttention(nn.Module):\n",
    "    def __init__(self, h, d_model, dropout=0.1):\n",
    "        \"Take in model size and number of heads.\"\n",
    "        super(MultiHeadedAttention, self).__init__()\n",
    "        assert d_model % h == 0 # check the h number\n",
    "        self.d_k = d_model // h\n",
    "        self.h = h\n",
    "        # 4 linear layers: WQ WK WV and final linear mapping WO  \n",
    "        self.linears = clones(nn.Linear(d_model, d_model), 4)\n",
    "        self.attn = None\n",
    "        self.dropout = nn.Dropout(p=dropout)\n",
    "\n",
    "    def forward(self, query, key, value, mask=None):\n",
    "        if mask is not None:\n",
    "            # Same mask applied to all h heads.\n",
    "            mask = mask.unsqueeze(1)\n",
    "        nbatches = query.size(0) # get batch size\n",
    "        # 1) Do all the linear projections in batch from d_model => h x d_k \n",
    "        # parttion into h sections，switch 2,3 axis for computation. \n",
    "        query, key, value = [l(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2) \n",
    "                             for l, x in zip(self.linears, (query, key, value))]\n",
    "        # 2) Apply attention on all the projected vectors in batch.\n",
    "        x, self.attn = attention(query, key, value, mask=mask, dropout=self.dropout)\n",
    "        # 3) \"Concat\" using a view and apply a final linear. \n",
    "        x = x.transpose(1, 2).contiguous().view(nbatches, -1, self.h * self.d_k)\n",
    "        return self.linears[-1](x) # final linear layer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Attention Mask**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The input $X$ is $[batch-size,  sequence-length]$, we use 'padding' to fill the matrix with 0 with respect to the longest sequence. \n",
    "\n",
    "But this will case issues for the softmax computation. \n",
    "$\\sigma(\\mathbf {z})_{i}={\\frac {e^{z_{i}}}{\\sum _{j=1}^{K}e^{z_{j}}}}$, where $e^0=1$.\n",
    "\n",
    "This means the padding sections join the computation, but they shouldn't. So we create this mask to ignore these area by assign a large negative bias.\n",
    " \n",
    "$$z_{illegal} = z_{illegal} + bias_{illegal}$$\n",
    "$$bias_{illegal} \\to -\\infty$$\n",
    "$$e^{z_{illegal}} \\to 0 $$  \n",
    "  \n",
    "Thus, the masked area will lead to 0 so we avoid them in computation.\n",
    "\n",
    "> Notification: in self-attention compution, we use mini-batch data as input, means we feed multiply lines of sentences into the model for training and computation. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](document/images/attention_mask.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In Transformer, both encoder and decoder attention computations need masking operation, but their functions are different.\n",
    "\n",
    "In the decoder, the self-attention layer is only allowed to attend to earlier positions in the output sequence. This is done by masking future positions (setting them to '-inf') before the softmax step in the self-attention calculation.\n",
    "\n",
    "The “Encoder-Decoder Attention” layer works just like multiheaded self-attention, except it creates its Queries matrix from the layer below it, and takes the Keys and Values matrix from the output of the encoder stack.\n",
    "  \n",
    "Here, we define a batch object that holds the src (English) and target sentences (Chinese) for training, as well as constructing the masks.\n",
    "\n",
    "> Notification: Mask(Opt.) is between Scale and Softmax\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T00:27:02.129099Z",
     "start_time": "2020-09-20T00:27:02.124593Z"
    }
   },
   "outputs": [],
   "source": [
    "class Batch:\n",
    "    \"Object for holding a batch of data with mask during training.\"\n",
    "    def __init__(self, src, trg=None, pad=0):\n",
    "        # convert words id to long format.  \n",
    "        src = torch.from_numpy(src).to(DEVICE).long()\n",
    "        trg = torch.from_numpy(trg).to(DEVICE).long()\n",
    "        self.src = src\n",
    "        # get the padding postion binary mask\n",
    "        # change the matrix shape to  1×seq.length\n",
    "        self.src_mask = (src != pad).unsqueeze(-2)\n",
    "        # 如果输出目标不为空，则需要对decoder要使用到的target句子进行mask\n",
    "        if trg is not None:\n",
    "            # decoder input from target \n",
    "            self.trg = trg[:, :-1]\n",
    "            # decoder target from trg \n",
    "            self.trg_y = trg[:, 1:]\n",
    "            # add attention mask to decoder input  \n",
    "            self.trg_mask = self.make_std_mask(self.trg, pad)\n",
    "            # check decoder output padding number\n",
    "            self.ntokens = (self.trg_y != pad).data.sum()\n",
    "    \n",
    "    # Mask \n",
    "    @staticmethod\n",
    "    def make_std_mask(tgt, pad):\n",
    "        \"Create a mask to hide padding and future words.\"\n",
    "        tgt_mask = (tgt != pad).unsqueeze(-2)\n",
    "        tgt_mask = tgt_mask & Variable(subsequent_mask(tgt.size(-1)).type_as(tgt_mask.data))\n",
    "        return tgt_mask # subsequent_mask is defined in 'decoder' section."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.3 Layer Normalization and Residual Connection"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1). **LayerNorm**:   \n",
    "  \n",
    "**Layer Normalization** normalize the hidden layer output into standard format, i.i.d, to boost the training efficency and model weight convergence (row wise). \n",
    "$$\\mu_{i}=\\frac{1}{m} \\sum^{m}_{i=1}x_{ij}$$  \n",
    "  \n",
    "$$\\sigma^{2}_{j}=\\frac{1}{m} \\sum^{m}_{i=1}\n",
    "(x_{ij}-\\mu_{j})^{2}$$  \n",
    "  \n",
    "$$LayerNorm(x)=\\alpha \\odot \\frac{x_{ij}-\\mu_{i}}\n",
    "{\\sqrt{\\sigma^{2}_{i}+\\epsilon}} + \\beta \\tag{eq.5}$$  \n",
    "  \n",
    "$\\epsilon$ is to avoid $0$ division; **$\\alpha, \\beta$** are parameter, $\\odot$ denotes element-wise product. Normally, we initialize $\\alpha$ as 1s and $\\beta$ as 0s."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2). **Residual Connection**:   \n",
    " \n",
    "We employ a residual connection  around each of the two sub-layers, followed by layer normalization.\n",
    "We get the Value matrix with the weights from attenations $Attention(Q,  K, V)$, and then we \n",
    "transpose it to make sure it shares the same shape of $X_{embedding}:[batch.size, sequence.length, embedding.dimension]$. And then add them together.\n",
    "\n",
    "$$X_{embedding} + Attention(Q, K, V)$$  \n",
    "  \n",
    "In the following compuations, after each module, we add the input with the output of the module to get residual connection, which allows the gradients be back-propogated to the start layers. \n",
    "$$X + SubLayer(X) \\tag{eq. 6}$$  \n",
    "  \n",
    "> **Notification**: to $SubLayer(X)$ we call the dropout function and then add x, $X + Dropout(SubLayer(X))$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T00:27:02.134759Z",
     "start_time": "2020-09-20T00:27:02.130657Z"
    }
   },
   "outputs": [],
   "source": [
    "class LayerNorm(nn.Module):\n",
    "    def __init__(self, features, eps=1e-6):\n",
    "        super(LayerNorm, self).__init__()\n",
    "        self.a_2 = nn.Parameter(torch.ones(features))\n",
    "        self.b_2 = nn.Parameter(torch.zeros(features))\n",
    "        self.eps = eps\n",
    "\n",
    "    def forward(self, x):\n",
    "        mean = x.mean(-1, keepdim=True) # rows\n",
    "        std = x.std(-1, keepdim=True)\n",
    "        x_zscore = (x - mean)/ torch.sqrt(std ** 2 + self.eps) \n",
    "        return self.a_2*x_zscore+self.b_2 "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "PyTorch has nn.LayerNorm，but we apply math equations here to learn. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T00:27:02.139180Z",
     "start_time": "2020-09-20T00:27:02.136279Z"
    }
   },
   "outputs": [],
   "source": [
    "class SublayerConnection(nn.Module):\n",
    "    \"\"\"\n",
    "    A residual connection followed by a layer norm.\n",
    "    Note for code simplicity the norm is first as opposed to last.\n",
    "    SublayerConnection: connect Multi-Head Attention and Feed Forward Layers \n",
    "    \"\"\"\n",
    "    def __init__(self, size, dropout):\n",
    "        super(SublayerConnection, self).__init__()\n",
    "        self.norm = LayerNorm(size)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x, sublayer):\n",
    "        \"Apply residual connection to any sublayer with the same size.\"\n",
    "        return x + self.dropout(sublayer(self.norm(x)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.4 Feedforwad Networks \n",
    "**Position-wise Feed-Forward Networks**\n",
    "\n",
    "In addition to attention sub-layers, each of the layers in our encoder and decoder contains a fully connected feed-forward network, which is applied to each position separately and identically. This consists of two linear transformations with a ReLU activation in between.\n",
    "\n",
    "$$ FFN(x) = max(0, xW_1 + b_1)W_2 + b_2$$\n",
    "\n",
    "While the linear transformations are the same across different positions, they use different parameters from layer to layer. Another way of describing this is as two convolutions with kernel size 1.\n",
    "\n",
    "The dimensionality of input and output is $d_{model}$, and the inner-layer has dimensionality $d_{ff}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T00:27:02.143660Z",
     "start_time": "2020-09-20T00:27:02.140463Z"
    }
   },
   "outputs": [],
   "source": [
    "class PositionwiseFeedForward(nn.Module):\n",
    "    def __init__(self, d_model, d_ff, dropout=0.1):\n",
    "        super(PositionwiseFeedForward, self).__init__()\n",
    "        self.w_1 = nn.Linear(d_model, d_ff)\n",
    "        self.w_2 = nn.Linear(d_ff, d_model)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x):\n",
    "        h1 = self.w_1(x)\n",
    "        h2 = self.dropout(h1)\n",
    "        return self.w_2(h2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.5 Transformer Encoder Overview"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we have programmed the four parts of the Transformer Encoder. Let us review how the data are transformed through all these layers.\n",
    "\n",
    "1. **Word Embedding and Positional Encoding**:\n",
    "\n",
    "$$X = EmbeddingLookup(X) + PositionalEncoding(X)$$\n",
    "\n",
    "$$X \\in \\mathbb{R}^{batch.size \\times  seq.len.\\times   embed.dim.} $$  \n",
    "  \n",
    "2. **Self-Attention and Mask**:\n",
    "\n",
    "$$Q = Linear(X) = XW_{Q}$$ \n",
    "$$K = Linear(X) = XW_{K}$$\n",
    "$$V = Linear(X) = XW_{V}$$\n",
    "$$X_{attention} = SelfAttention(Q, K, V)$$  \n",
    "\n",
    "3. **Residual Connection and Layer Normalization**:\n",
    "\n",
    "$$X_{attention} = LayerNorm(X_{attention})$$\n",
    "$$X_{attention} = X + X_{attention} $$  \n",
    "\n",
    "4. **Position-wise Feed-Forward Networks** two linear mappings with ReLU Avtivation function:\n",
    "$$X_{hidden} = Linear(Activate(Linear(X_{attention})))$$  \n",
    "  \n",
    "5. **Repeat 3** :\n",
    "$$X_{hidden} = LayerNorm(X_{hidden})$$\n",
    "$$X_{hidden} = X_{attention} + X_{hidden}$$\n",
    "\n",
    "$$X_{hidden} \\in \\mathbb{R}^{batch.size \\times  seq.len.\\times   embed.dim.}$$  \n",
    "  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T00:27:02.147675Z",
     "start_time": "2020-09-20T00:27:02.145260Z"
    }
   },
   "outputs": [],
   "source": [
    "def clones(module, N):\n",
    "    \"\"\"\n",
    "    \"Produce N identical layers.\"\n",
    "    Use deepcopy the weight are indenpendent.\n",
    "    \"\"\"\n",
    "    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the paper, $Encoder$ has $N=6$ blocks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T00:27:02.152154Z",
     "start_time": "2020-09-20T00:27:02.149114Z"
    }
   },
   "outputs": [],
   "source": [
    "class Encoder(nn.Module):\n",
    "    \"Core encoder is a stack of N layers (blocks)\"\n",
    "    def __init__(self, layer, N):\n",
    "        super(Encoder, self).__init__()\n",
    "        self.layers = clones(layer, N)\n",
    "        self.norm = LayerNorm(layer.size)\n",
    "\n",
    "    def forward(self, x, mask):\n",
    "        \"\"\"\n",
    "        Pass the input (and mask) through each layer in turn.\n",
    "        \"\"\"\n",
    "        for layer in self.layers:\n",
    "            x = layer(x, mask)\n",
    "        return self.norm(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Each **Encoder Block** contains two sub-layers(**Self-Attention**,**Position-wise**) and 2 sublayer-connetions:\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T00:27:02.157528Z",
     "start_time": "2020-09-20T00:27:02.153637Z"
    }
   },
   "outputs": [],
   "source": [
    "class EncoderLayer(nn.Module):\n",
    "    def __init__(self, size, self_attn, feed_forward, dropout):\n",
    "        super(EncoderLayer, self).__init__()\n",
    "        self.self_attn = self_attn\n",
    "        self.feed_forward = feed_forward\n",
    "        self.sublayer = clones(SublayerConnection(size, dropout), 2)\n",
    "        self.size = size # d_model\n",
    "\n",
    "    def forward(self, x, mask):\n",
    "        # X-embedding to Multi-head-Attention\n",
    "        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))\n",
    "        # X-embedding to feed-forwad nn\n",
    "        return self.sublayer[1](x, self.feed_forward)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3.  Decoder"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-19T11:07:22.025892Z",
     "start_time": "2020-09-19T11:07:22.008459Z"
    }
   },
   "source": [
    "After the introduction of the encoder structure, we can see the decoder shares a lot similarities of encoder.\n",
    "It also stacks N times. But there is a Encoder-Deconder-Contex-Attention layer (sublayer[1]) between the Masked MHA[0] and FFN[2]. It use the output of the decoder as query to search the output of encoder with MHA, which makes decoder see all the outputs from encoder.\n",
    "\n",
    "Decoding process:\n",
    "- Input: Encoding output(memory) and i-1 position decoder output/\n",
    "- Output: i position output work probabilities.\n",
    "- decoding process works like RNN."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![Decoder Structure](document/images/decoder.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T00:27:02.169437Z",
     "start_time": "2020-09-20T00:27:02.159137Z"
    }
   },
   "outputs": [],
   "source": [
    "class Decoder(nn.Module):\n",
    "    def __init__(self, layer, N):\n",
    "        \"Generic N layer decoder with masking.\"\n",
    "        super(Decoder, self).__init__()\n",
    "        self.layers = clones(layer, N)\n",
    "        self.norm = LayerNorm(layer.size)\n",
    "\n",
    "    def forward(self, x, memory, src_mask, tgt_mask):\n",
    "        \"\"\"\n",
    "        Repeat decoder N times\n",
    "        Decoderlayer get a input attention mask (src) \n",
    "        and a output attention mask (tgt) + subsequent mask \n",
    "        \"\"\"\n",
    "        for layer in self.layers:\n",
    "            x = layer(x, memory, src_mask, tgt_mask)\n",
    "        return self.norm(x)\n",
    "\n",
    "\n",
    "class DecoderLayer(nn.Module):\n",
    "    def __init__(self, size, self_attn, src_attn, feed_forward, dropout):\n",
    "        super(DecoderLayer, self).__init__()\n",
    "        self.size = size\n",
    "        self.self_attn = self_attn\n",
    "        self.src_attn = src_attn\n",
    "        self.feed_forward = feed_forward\n",
    "        self.sublayer = clones(SublayerConnection(size, dropout), 3)\n",
    "\n",
    "    def forward(self, x, memory, src_mask, tgt_mask):\n",
    "        m = memory # encoder output embedding\n",
    "        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask))\n",
    "        x = self.sublayer[1](x, lambda x: self.src_attn(x, m, m, src_mask))\n",
    "        # Context-Attention：q=decoder hidden，k,v from encoder hidden\n",
    "        return self.sublayer[2](x, self.feed_forward)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We also modify the self-attention sub-layer in the decoder stack to prevent positions from attending to subsequent positions. This masking (**subsequent_mask**), combined with fact that the output embeddings are offset by one position, ensures that the predictions for position i can depend only on the known outputs at positions less than i .\n",
    "\n",
    "\n",
    "For Encoder src-mask, just mask the padding cells. \n",
    "But for decoder trg-mask, we need mask the padding and add the subsequent-mask process.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T00:27:02.174560Z",
     "start_time": "2020-09-20T00:27:02.171267Z"
    }
   },
   "outputs": [],
   "source": [
    "def subsequent_mask(size):\n",
    "    \"Mask out subsequent positions.\"\n",
    "    attn_shape = (1, size, size)\n",
    "    subsequent_mask = np.triu(np.ones(attn_shape), k=1).astype('uint8')\n",
    "    return torch.from_numpy(subsequent_mask) == 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below the attention mask shows the position each tgt word (row) is allowed to look at (column). Words are blocked for attending to future words during training.\n",
    "\"Yellow\" color denote True."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T00:27:02.323043Z",
     "start_time": "2020-09-20T00:27:02.176195Z"
    }
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(5,5))\n",
    "plt.imshow(subsequent_mask(20)[0])\n",
    "None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Transformer Model\n",
    "  \n",
    "Finally, let us put encoder and decoder together with the 'generator'.\n",
    "\n",
    "![](document/images/English-to-Chinese.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Recall the decoder then generates an output sequence, of symbols one element at a time. At each step the model is auto-regressive (cite), consuming the previously generated symbols as additional input when generating the next. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T00:27:02.332072Z",
     "start_time": "2020-09-20T00:27:02.325677Z"
    }
   },
   "outputs": [],
   "source": [
    "class Transformer(nn.Module):\n",
    "    def __init__(self, encoder, decoder, src_embed, tgt_embed, generator):\n",
    "        super(Transformer, self).__init__()\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "        self.src_embed = src_embed\n",
    "        self.tgt_embed = tgt_embed\n",
    "        self.generator = generator \n",
    "\n",
    "    def encode(self, src, src_mask):\n",
    "        return self.encoder(self.src_embed(src), src_mask)\n",
    "\n",
    "    def decode(self, memory, src_mask, tgt, tgt_mask):\n",
    "        return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask)\n",
    "\n",
    "    def forward(self, src, tgt, src_mask, tgt_mask):\n",
    "        \"Take in and process masked src and target sequences.\"\n",
    "        # encoder output will be the decoder's memory for decoding\n",
    "        return self.decode(self.encode(src, src_mask), src_mask, tgt, tgt_mask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T00:27:02.341832Z",
     "start_time": "2020-09-20T00:27:02.338069Z"
    }
   },
   "outputs": [],
   "source": [
    "class Generator(nn.Module):\n",
    "    def __init__(self, d_model, vocab):\n",
    "        super(Generator, self).__init__()\n",
    "        # decode: d_model to vocab mapping\n",
    "        self.proj = nn.Linear(d_model, vocab)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return F.log_softmax(self.proj(x), dim=-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Set Parameters and Create the Full Transformer model Function**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T00:27:02.355463Z",
     "start_time": "2020-09-20T00:27:02.349624Z"
    }
   },
   "outputs": [],
   "source": [
    "def make_model(src_vocab, tgt_vocab, N=6, d_model=512, d_ff=2048, h = 8, dropout=0.1):\n",
    "    c = copy.deepcopy\n",
    "    #  Attention \n",
    "    attn = MultiHeadedAttention(h, d_model).to(DEVICE)\n",
    "    #  FeedForward \n",
    "    ff = PositionwiseFeedForward(d_model, d_ff, dropout).to(DEVICE)\n",
    "    #  Positional Encoding\n",
    "    position = PositionalEncoding(d_model, dropout).to(DEVICE)\n",
    "    #  Transformer \n",
    "    model = Transformer(\n",
    "        Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout).to(DEVICE), N).to(DEVICE),\n",
    "        Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), dropout).to(DEVICE), N).to(DEVICE),\n",
    "        nn.Sequential(Embeddings(d_model, src_vocab).to(DEVICE), c(position)),\n",
    "        nn.Sequential(Embeddings(d_model, tgt_vocab).to(DEVICE), c(position)),\n",
    "        Generator(d_model, tgt_vocab)).to(DEVICE)\n",
    "    \n",
    "    # This was important from their code. \n",
    "    # Initialize parameters with Glorot / fan_avg.\n",
    "    # Paper title: Understanding the difficulty of training deep feedforward neural networks Xavier\n",
    "    for p in model.parameters():\n",
    "        if p.dim() > 1:\n",
    "            nn.init.xavier_uniform_(p)\n",
    "    return model.to(DEVICE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Transformer Model Training: English-to-Chinese  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Regularization **Label Smoothing**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "During training, we employed label smoothing of value $\\epsilon_{ls}=0.1$ (https://arxiv.org/pdf/1512.00567.pdf). \n",
    "\n",
    "This hurts perplexity, as the model learns to be more unsure, but improves accuracy and BLEU score.\n",
    "  \n",
    ">We implement label smoothing using the KL div loss. Instead of using a one-hot target distribution, we create a distribution that has confidence of the correct word and the rest of the smoothing mass distributed throughout the vocabulary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T00:27:02.365730Z",
     "start_time": "2020-09-20T00:27:02.359819Z"
    }
   },
   "outputs": [],
   "source": [
    "class LabelSmoothing(nn.Module):\n",
    "    \"Implement label smoothing.\"\n",
    "    def __init__(self, size, padding_idx, smoothing=0.0):\n",
    "        super(LabelSmoothing, self).__init__()\n",
    "        self.criterion = nn.KLDivLoss(reduction='sum') # 2020 update\n",
    "        self.padding_idx = padding_idx\n",
    "        self.confidence = 1.0 - smoothing\n",
    "        self.smoothing = smoothing\n",
    "        self.size = size\n",
    "        self.true_dist = None\n",
    "        \n",
    "    def forward(self, x, target):\n",
    "        assert x.size(1) == self.size\n",
    "        true_dist = x.data.clone()\n",
    "        true_dist.fill_(self.smoothing / (self.size - 2))\n",
    "        true_dist.scatter_(1, target.data.unsqueeze(1), self.confidence)\n",
    "        true_dist[:, self.padding_idx] = 0\n",
    "        mask = torch.nonzero(target.data == self.padding_idx)\n",
    "        if mask.dim() > 0:\n",
    "            true_dist.index_fill_(0, mask.squeeze(), 0.0)\n",
    "        self.true_dist = true_dist\n",
    "        return self.criterion(x, Variable(true_dist, requires_grad=False))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here, 'size' is from vocab，'smoothing' is value to be distributed on non-groundtruth.\n",
    "We can see an example of how the mass is distributed to the words based on confidence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T00:27:02.500729Z",
     "start_time": "2020-09-20T00:27:02.367927Z"
    }
   },
   "outputs": [],
   "source": [
    "# Example of label smoothing.\n",
    "crit = LabelSmoothing(5, 0, 0.1)  #  ϵ=0.4\n",
    "predict = torch.FloatTensor([[0, 0.2, 0.7, 0.1, 0],\n",
    "                             [0, 0.2, 0.7, 0.1, 0], \n",
    "                             [0, 0.2, 0.7, 0.1, 0]])\n",
    "v = crit(Variable(predict.log()), Variable(torch.LongTensor([2, 1, 0])))\n",
    "\n",
    "# Show the target distributions expected by the system.\n",
    "plt.imshow(crit.true_dist)\n",
    "None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Label smoothing actually starts to penalize the model if it gets very confident about a given choice."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T00:27:02.604827Z",
     "start_time": "2020-09-20T00:27:02.502536Z"
    }
   },
   "outputs": [],
   "source": [
    "crit = LabelSmoothing(5, 0, 0.1)\n",
    "def loss(x):\n",
    "    d = x + 3 * 1\n",
    "    predict = torch.FloatTensor([[0, x / d, 1 / d, 1 / d, 1 / d]])\n",
    "    return crit(Variable(predict.log()), Variable(torch.LongTensor([1]))).item()\n",
    "\n",
    "plt.plot(np.arange(1, 100), [loss(x) for x in range(1, 100)])\n",
    "None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Loss Computation**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T00:27:02.611144Z",
     "start_time": "2020-09-20T00:27:02.606658Z"
    }
   },
   "outputs": [],
   "source": [
    "class SimpleLossCompute:\n",
    "    def __init__(self, generator, criterion, opt=None):\n",
    "        self.generator = generator\n",
    "        self.criterion = criterion\n",
    "        self.opt = opt\n",
    "        \n",
    "    def __call__(self, x, y, norm):\n",
    "        x = self.generator(x)\n",
    "        loss = self.criterion(x.contiguous().view(-1, x.size(-1)), \n",
    "                              y.contiguous().view(-1)) / norm\n",
    "        loss.backward()\n",
    "        if self.opt is not None:\n",
    "            self.opt.step()\n",
    "            self.opt.optimizer.zero_grad()\n",
    "        return loss.data.item() * norm.float()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Optimizer with Warmup Learning Rate**"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "According to the paper, they applied a warmup learning rate with Adam Optimizer with $\\beta_1=0.9、\\beta_2=0.98$ 和 $\\epsilon = 10^{−9}$.  \n",
    "\n",
    "This will update the learning rate over the course of training, according to the formula:\n",
    "\n",
    "$$lrate = d^{−0.5}_{model}⋅min(step\\_num^{−0.5},\\; step\\_num⋅warmup\\_steps^{−1.5})$$  \n",
    "\n",
    "This corresponds to increasing the learning rate linearly for the first \"warmup_steps\" training steps, and decreasing it thereafter proportionally to the inverse square root of the step number."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T00:27:02.618655Z",
     "start_time": "2020-09-20T00:27:02.612941Z"
    }
   },
   "outputs": [],
   "source": [
    "class NoamOpt:\n",
    "    \"Optim wrapper that implements rate.\"\n",
    "    def __init__(self, model_size, factor, warmup, optimizer):\n",
    "        self.optimizer = optimizer\n",
    "        self._step = 0\n",
    "        self.warmup = warmup\n",
    "        self.factor = factor\n",
    "        self.model_size = model_size\n",
    "        self._rate = 0\n",
    "        \n",
    "    def step(self):\n",
    "        \"Update parameters and rate\"\n",
    "        self._step += 1\n",
    "        rate = self.rate()\n",
    "        for p in self.optimizer.param_groups:\n",
    "            p['lr'] = rate\n",
    "        self._rate = rate\n",
    "        self.optimizer.step()\n",
    "        \n",
    "    def rate(self, step = None):\n",
    "        \"Implement `lrate` above\"\n",
    "        if step is None:\n",
    "            step = self._step\n",
    "        return self.factor * (self.model_size ** (-0.5) * min(step ** (-0.5), step * self.warmup ** (-1.5)))\n",
    "    \n",
    "# We used factor=2, warmup-step = 4000\n",
    "def get_std_opt(model):\n",
    "    return NoamOpt(model.src_embed[0].d_model, 2, 4000,\n",
    "            torch.optim.Adam(model.parameters(), lr=0, betas=(0.9, 0.98), eps=1e-9))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "all the updates are for the learning rate, \n",
    "- model-size denotes $d_{model}$. \n",
    "- warmup denotes  warmup-steps.\n",
    "- factor is a scalar."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Example of the curves of this model for different model sizes and for optimization hyperparameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-20T00:27:08.648356Z",
     "start_time": "2020-09-20T00:27:08.258785Z"
    }
   },
   "outputs": [],
   "source": [
    "# Three settings of the lrate hyperparameters.\n",
    "opts = [NoamOpt(512, 1, 4000, None), \n",
    "        NoamOpt(512, 1, 8000, None),\n",
    "        NoamOpt(256, 1, 4000, None)]\n",
    "plt.plot(np.arange(1, 20000), [[opt.rate(i) for opt in opts] for i in range(1, 20000)])\n",
    "plt.legend([\"512:4000\", \"512:8000\", \"256:4000\"])\n",
    "None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Training Iterators**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-19T13:33:21.978199Z",
     "start_time": "2020-09-19T13:33:21.974190Z"
    }
   },
   "outputs": [],
   "source": [
    "def run_epoch(data, model, loss_compute, epoch):\n",
    "    start = time.time()\n",
    "    total_tokens = 0.\n",
    "    total_loss = 0.\n",
    "    tokens = 0.\n",
    "    for i , batch in enumerate(data):\n",
    "        out = model(batch.src, batch.trg, batch.src_mask, batch.trg_mask)\n",
    "        loss = loss_compute(out, batch.trg_y, batch.ntokens)\n",
    "        total_loss += loss\n",
    "        total_tokens += batch.ntokens\n",
    "        tokens += batch.ntokens\n",
    "        if i % 50 == 1:\n",
    "            elapsed = time.time() - start\n",
    "            print(\"Epoch {:d} Batch: {:d} Loss: {:.4f} Tokens per Sec: {:.2f}s\".format(epoch, i - 1, loss / batch.ntokens, (tokens.float() / elapsed / 1000.)))\n",
    "            start = time.time()\n",
    "            tokens = 0\n",
    "    return total_loss / total_tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-19T13:33:21.983286Z",
     "start_time": "2020-09-19T13:33:21.979822Z"
    }
   },
   "outputs": [],
   "source": [
    "def train(data, model, criterion, optimizer):\n",
    "    \"\"\"\n",
    "    Train and Save the model.\n",
    "    \"\"\"\n",
    "    # init loss as a large value\n",
    "    best_dev_loss = 1e5\n",
    "    \n",
    "    for epoch in range(EPOCHS):\n",
    "        # Train model \n",
    "        model.train()\n",
    "        run_epoch(data.train_data, model, SimpleLossCompute(model.generator, criterion, optimizer), epoch)\n",
    "        model.eval()\n",
    "\n",
    "        # validate model on dev dataset\n",
    "        print('>>>>> Evaluate')\n",
    "        dev_loss = run_epoch(data.dev_data, model, SimpleLossCompute(model.generator, criterion, None), epoch)\n",
    "        print('<<<<< Evaluate loss: {:.2f}'.format(dev_loss))\n",
    "        \n",
    "        # save the model with best-dev-loss\n",
    "        if dev_loss < best_dev_loss:\n",
    "            best_dev_loss = dev_loss\n",
    "            torch.save(model.state_dict(), SAVE_FILE) # SAVE_FILE = 'save/model.pt'\n",
    "            \n",
    "        print(f\">>>>> current best loss: {best_dev_loss}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-19T13:33:30.293504Z",
     "start_time": "2020-09-19T13:33:21.984823Z"
    }
   },
   "outputs": [],
   "source": [
    "# Step 1: Data Preprocessing\n",
    "data = PrepareData(TRAIN_FILE, DEV_FILE)\n",
    "src_vocab = len(data.en_word_dict)\n",
    "tgt_vocab = len(data.cn_word_dict)\n",
    "print(f\"src_vocab {src_vocab}\")\n",
    "print(f\"tgt_vocab {tgt_vocab}\")\n",
    "\n",
    "# Step 2: Init model\n",
    "model = make_model(\n",
    "                    src_vocab, \n",
    "                    tgt_vocab, \n",
    "                    LAYERS, \n",
    "                    D_MODEL, \n",
    "                    D_FF,\n",
    "                    H_NUM,\n",
    "                    DROPOUT\n",
    "                )\n",
    "\n",
    "# Step 3: Training model\n",
    "print(\">>>>>>> start train\")\n",
    "train_start = time.time()\n",
    "criterion = LabelSmoothing(tgt_vocab, padding_idx = 0, smoothing= 0.0)\n",
    "optimizer = NoamOpt(D_MODEL, 1, 2000, torch.optim.Adam(model.parameters(), lr=0, betas=(0.9,0.98), eps=1e-9))\n",
    "\n",
    "train(data, model, criterion, optimizer)\n",
    "print(f\"<<<<<<< finished train, cost {time.time()-train_start:.4f} seconds\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Prediction with English-to-Chinese Translator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-19T13:33:30.299613Z",
     "start_time": "2020-09-19T13:33:30.295197Z"
    }
   },
   "outputs": [],
   "source": [
    "def greedy_decode(model, src, src_mask, max_len, start_symbol):\n",
    "    \"\"\"\n",
    "    Translate src with model\n",
    "    \"\"\"\n",
    "    # decode the src \n",
    "    memory = model.encode(src, src_mask)\n",
    "    # init 1×1 tensor as prediction，fill in ('BOS')id, type: (LongTensor)\n",
    "    ys = torch.ones(1, 1).fill_(start_symbol).type_as(src.data)\n",
    "    #  run 遍历输出的长度下标\n",
    "    for i in range(max_len-1):\n",
    "        # decode one by one\n",
    "        out = model.decode(memory, \n",
    "                           src_mask, \n",
    "                           Variable(ys), \n",
    "                           Variable(subsequent_mask(ys.size(1)).type_as(src.data)))\n",
    "        #  out to log_softmax \n",
    "        prob = model.generator(out[:, -1])\n",
    "        #  get the max-prob id\n",
    "        _, next_word = torch.max(prob, dim = 1)\n",
    "        next_word = next_word.data[0]\n",
    "        #  concatnate with early predictions\n",
    "        ys = torch.cat([ys,torch.ones(1, 1).type_as(src.data).fill_(next_word)], dim=1)\n",
    "    return ys"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "English to Chinese Translations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-19T13:33:30.306447Z",
     "start_time": "2020-09-19T13:33:30.301271Z"
    }
   },
   "outputs": [],
   "source": [
    "def evaluate(data, model):\n",
    "    \"\"\"\n",
    "    Make prediction with trained model, and print results.\n",
    "    \"\"\"\n",
    "    with torch.no_grad():\n",
    "        #  pick some random sentences from dev data.         \n",
    "        for i in np.random.randint(len(data.dev_en), size=10):\n",
    "            # Print English sentence\n",
    "            en_sent = \" \".join([data.en_index_dict[w] for w in data.dev_en[i]])\n",
    "            print(\"\\n\" + en_sent)\n",
    "            \n",
    "            # Print Target Chinese sentence\n",
    "            cn_sent =  \" \".join([data.cn_index_dict[w] for w in data.dev_cn[i]])\n",
    "            print(\"\".join(cn_sent))\n",
    "            \n",
    "            # conver English to tensor  \n",
    "            src = torch.from_numpy(np.array(data.dev_en[i])).long().to(DEVICE)\n",
    "            src = src.unsqueeze(0)\n",
    "            # set attention mask\n",
    "            src_mask = (src != 0).unsqueeze(-2)\n",
    "            # apply model to decode, make prediction\n",
    "            out = greedy_decode(model, src, src_mask, max_len=MAX_LENGTH, start_symbol=data.cn_word_dict[\"BOS\"])\n",
    "            # save all in the translation list \n",
    "            translation = []\n",
    "            # convert id to Chinese, skip 'BOS' 0.\n",
    "            # 遍历翻译输出字符的下标（注意：跳过开始符\"BOS\"的索引 0）\n",
    "            for j in range(1, out.size(1)):\n",
    "                sym = data.cn_index_dict[out[0, j].item()]\n",
    "                if sym != 'EOS':\n",
    "                    translation.append(sym)\n",
    "                else:\n",
    "                    break\n",
    "            print(\"translation: {}\".format(\" \".join(translation)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**English to Chinese Translator** "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-19T13:33:32.915034Z",
     "start_time": "2020-09-19T13:33:30.308167Z"
    }
   },
   "outputs": [],
   "source": [
    "# Predition\n",
    "model.load_state_dict(torch.load(SAVE_FILE))\n",
    "print(\">>>>>>> start evaluate\")\n",
    "evaluate_start  = time.time()\n",
    "evaluate(data, model)         \n",
    "print(f\"<<<<<<< finished evaluate, cost {time.time()-evaluate_start:.4f} seconds\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If your translation does not look good, that is totally OK. \n",
    "You can change the “Debug“ to False and retrain the model with CUDA support."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-19T13:33:32.919407Z",
     "start_time": "2020-09-19T13:33:32.916832Z"
    }
   },
   "outputs": [],
   "source": [
    "# Here is my pretrained model with CUDA support \n",
    "if DEVICE==\"cuda\":\n",
    "    SAVE_FILE_EXTRA  = 'save/models/large_model.pt'  \n",
    "    model.load_state_dict(torch.load(SAVE_FILE_EXTRA))\n",
    "    print(\">>>>>>> start evaluate\")\n",
    "    evaluate_start  = time.time()\n",
    "    evaluate(data, model)         \n",
    "    print(f\"<<<<<<< finished evaluate, cost {time.time()-evaluate_start:.4f} seconds\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2020-09-19T12:43:58.615447Z",
     "start_time": "2020-09-19T12:43:58.319968Z"
    }
   },
   "source": [
    "## Reference\n",
    "\n",
    "- The Annotated Transformer http://nlp.seas.harvard.edu/2018/04/03/attention.html\n",
    "- The Illustrated Transformer http://jalammar.github.io/illustrated-transformer/\n",
    "- GreedAI NLP Training Camp: https://www.greedyai.com/ \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "THE END"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Base",
   "language": "python",
   "name": "base"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.8"
  },
  "nteract": {
   "version": "0.25.0"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": false,
   "sideBar": true,
   "skip_h1_title": true,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": true,
   "toc_position": {
    "height": "calc(100% - 180px)",
    "left": "10px",
    "top": "150px",
    "width": "256px"
   },
   "toc_section_display": true,
   "toc_window_display": false
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "position": {
    "height": "557px",
    "left": "447px",
    "right": "20px",
    "top": "255px",
    "width": "683px"
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}