tutorials/attention_is_all_you_need/practical.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial: Attention is all you need\n",
    "\n",
    "In this tutorial we will implement [Vaswani et al.](https://arxiv.org/abs/1706.03762)'s transductive model that is entirely based on attention mechanism instead of previously used recurrence models like RNNs. \n",
    "For example, in the case of English to German translation, an RNN would need to be trained sequentially such that the hidden state at time `t-1` along with the input at time `t` need to be known to compute the German translation upto that point. \n",
    "This not only introduces difficulty in training variable sized sequences but also draws upon an assumption that the translation changes from one part of the sequence to the other. \n",
    "\n",
    "Transformer architecture proposed by [Vaswani et al.](https://arxiv.org/abs/1706.03762) overcomes these problems by making the model entirely based on self-attention mechanism, thereby allowing for faster training due to efficient batching as well as doing away with the assumption on sequential nature of translation task. \n",
    "We will build such a transformer from scratch and train it on [Multimodal Machine Translation](http://www.statmt.org/wmt16/multimodal-task.html#task1) task in this tutorial. \n",
    "This task is of appropriate size for us to get a full sense of how Transformer works. \n",
    "It has 29,000 English to German translations text for training, 1,014 for validation and 1,000 samples for test dataset. \n",
    "\n",
    "\n",
    "The final evaluation will be done using METEOR metric which computes weighted harmonic mean of precision and recall (recall is given a weight of 90%) of the candidate translation with respect to the reference translation. \n",
    "Read more about this metric [here](https://en.wikipedia.org/wiki/METEOR).\n",
    "\n",
    "\n",
    "Some usefule resources to understand Transformers (also helped in building this tutorial):\n",
    "1. [Attention Is All You Need [Original paper] ](https://arxiv.org/abs/1706.03762)\n",
    "2. [Transformers from Scratch](http://peterbloem.nl/blog/transformers)\n",
    "3. [The Illustrated Transformer](https://jalammar.github.io/illustrated-transformer/)\n",
    "4. [The Annotated Transformer](https://nlp.seas.harvard.edu/2018/04/03/attention.html)\n",
    "5. [How to code the transformer in PyTorch](https://towardsdatascience.com/how-to-code-the-transformer-in-pytorch-24db27c8f9ec)\n",
    "\n",
    "[[Vaswani et al. 2017] Attention Is All You Need](https://arxiv.org/abs/1706.03762)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "We will need the following libraries \n",
    "\n",
    "- [`torch`](https://pytorch.org/)\n",
    "- [`torchtext`](https://github.com/pytorch/text)\n",
    "- [`spacy`](https://spacy.io/usage) \n",
    "- [`nltk`](https://www.nltk.org/index.html) \n",
    "\n",
    "Follow the link to the libraries above to see the instructions to install the libraries. We will be using `spacy`'s specialized tokenizer (see below), so run the following commands to install the necessary modules. \n",
    "\n",
    "```python\n",
    "pip install -U spacy\n",
    "python -m spacy download en_core_web_sm\n",
    "python -m spacy download de_core_news_sm\n",
    "```\n",
    "\n",
    "We will use `nltk`'s [`nltk.translate.meteor_score`](https://www.nltk.org/_modules/nltk/translate/meteor_score.html) to evaluate the performance of our translator. \n",
    "This library requires `wordnet` corpus which can be installed using [`nltk.download()`](https://www.nltk.org/data.html). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.nn.utils.rnn import pad_sequence \n",
    "\n",
    "from torchtext.data.utils import get_tokenizer\n",
    "from torchtext.vocab import build_vocab_from_iterator\n",
    "from torchtext.datasets import Multi30k\n",
    "\n",
    "import nltk\n",
    "import math\n",
    "import numpy as np\n",
    "from collections import defaultdict, namedtuple\n",
    "import matplotlib.pyplot as plt \n",
    "import pathlib\n",
    "\n",
    "# fix seed for reproducibility \n",
    "rng = np.random.RandomState(1)\n",
    "torch.manual_seed(rng.randint(np.iinfo(int).max))\n",
    "\n",
    "# it is a good practice to define `device` globally\n",
    "if torch.cuda.is_available():\n",
    "    device = torch.device(\"cuda\")\n",
    "    print(\"Using GPU:\", device)\n",
    "else:\n",
    "    device = torch.device(\"cpu\")\n",
    "    print(\"No GPU -> using CPU:\", device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Machine Translation task\n",
    "\n",
    "\n",
    "In this tutorial, we will focus on building a model that translates English sentences to German.\n",
    "To access such dataset , we will use `torchtext.datasets.Multi30k` that provies us these sentence pairs.\n",
    "The function also separates out `train`, `valid`, and `test` datasets for us. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "LANGUAGE_PAIRS = ['de', 'en']\n",
    "Multi30k(root=\"./data\", split=(\"train\", \"valid\", \"test\"), language_pair=LANGUAGE_PAIRS)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load training data\n",
    "\n",
    "We will load English-German pairs of sentences in the form of an iterator. \n",
    "Since the iterators need to be loaded again once they reach the end we will write a function `get_data_iter` to do this. \n",
    "We will call this function whenever we need to refresh the data iterator. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_data_iter(itertype):\n",
    "    return Multi30k(root=\"./data\", split=itertype, language_pair=LANGUAGE_PAIRS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_iter = get_data_iter(\"train\")\n",
    "\n",
    "# example call\n",
    "print(\"An example of English-German pair of sentences\\n\")\n",
    "de, en= next(train_iter)\n",
    "print(en, de)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Preprocess text data\n",
    "\n",
    "Text data, in its raw form, can't be digested directly by the machine learning models.\n",
    "Therefore, we need to transform this data before machine learning models can learn any patterns in it. \n",
    "We will be using the following transformations (in order):\n",
    "\n",
    "- **tokenization**: It refers to breaking a sentence into the list of unique words and punctuations. We will use `spacy`'s specialized tokenizer for this purpose. \n",
    "\n",
    "\n",
    "- **token to index**: The tokenized words in a sentence need to be transformed to their unique indices. To do this, we will need to iterate through the entire tokenized corpus to build a mapping from tokens to index. With the help of this mapping, we can convert any list of token into its corresponding unique indices. We will need to specify the index with which we will denote special cases (e.g., unknown words) so that the vocabulary builder does not use those numbers to denote any tokens. \n",
    "    - *Unknown words*: Words can be classified as unknown either (a) when only the words above certain frequency in the corpus are recognized by the vocabulary builder or (b) when a word is not present in the training corpus. Therefore, we need to reserve a number or index to denote unknown words\n",
    "    \n",
    "    - *beginning-of-sentence (BOS)*: this index will denote the beginning of a sentence\n",
    "    \n",
    "    - *end-of-sentence (EOS)*: this index will denote the end of a sentence\n",
    "    \n",
    "    - *padding index*: While training, different sentences need to be batched together. Therefore, we need to define a padding index to mark the irrelevant part of the input tensors. \n",
    "    \n",
    "\n",
    "- **marking beginning and end of a sentence**: The beginning and ending of a sentence might be clear in the input, but the model need to predict the end-of-sentence token in its output. Therefore, we will pad the indexed tokens with BOS and EOS indices. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SRC_LANGUAGE = LANGUAGE_PAIRS[0]\n",
    "TGT_LANGUAGE = LANGUAGE_PAIRS[1]\n",
    "\n",
    "SPACY_LANGUAGE_MAP = {\n",
    "    'en': 'en_core_web_sm',\n",
    "    'de': 'de_core_news_sm'\n",
    "}\n",
    "\n",
    "## tokenization\n",
    "token_transform = {}\n",
    "token_transform[SRC_LANGUAGE] = get_tokenizer('spacy', language=SPACY_LANGUAGE_MAP[SRC_LANGUAGE])\n",
    "token_transform[TGT_LANGUAGE] = ## Define tokenizer for the target language\n",
    "\n",
    "## token to index\n",
    "\n",
    "# Define special symbols and indices\n",
    "UNK_IDX, PAD_IDX, BOS_IDX, EOS_IDX = 0, 1, 2, 3\n",
    "\n",
    "# Make sure the tokens are in order of their indices to properly insert them in vocab\n",
    "special_symbols = ['<unk>', '<pad>', '<bos>', '<eos>']\n",
    "\n",
    "\n",
    "def yield_tokens(data_iter, language):\n",
    "    idx = LANGUAGE_PAIRS.index(language)\n",
    "    \n",
    "    for sentence in data_iter:\n",
    "        yield ## call token_transform for language on appropriate sentence\n",
    "\n",
    "vocab_transform = {}\n",
    "vocab_transform[SRC_LANGUAGE] = build_vocab_from_iterator(yield_tokens(get_data_iter(\"train\"), SRC_LANGUAGE), min_freq=1, specials=special_symbols, special_first=True)\n",
    "vocab_transform[TGT_LANGUAGE] = ## build vocab for the target language\n",
    "\n",
    "for v in vocab_transform.values():\n",
    "    v.set_default_index(UNK_IDX)\n",
    "\n",
    "## marking beginning and end of a sentence\n",
    "def bos_eos_pad(tokens):\n",
    "    return torch.cat((\n",
    "            torch.tensor([BOS_IDX]), \n",
    "            torch.tensor(tokens), \n",
    "            torch.tensor([EOS_IDX])\n",
    "    ))\n",
    "\n",
    "## Putting it all together\n",
    "def get_text_transform(language):\n",
    "    def func(sentence):\n",
    "        x = token_transform[language](sentence)\n",
    "        x = vocab_transform[language](x)\n",
    "        x = bos_eos_pad(x)\n",
    "        return x\n",
    "    return func\n",
    "\n",
    "text_transform = {}\n",
    "text_transform[SRC_LANGUAGE] = get_text_transform(SRC_LANGUAGE)\n",
    "text_transform[TGT_LANGUAGE] = get_text_transform(TGT_LANGUAGE)\n",
    "\n",
    "# n tokens\n",
    "SRC_N_TOKENS = len(vocab_transform[SRC_LANGUAGE])\n",
    "TGT_N_TOKENS = len(vocab_transform[TGT_LANGUAGE])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# example call\n",
    "print(\"An example of English-German pair of sentences\\n\", \"-\"*50)\n",
    "train_iter = get_data_iter('train')\n",
    "de, en = next(train_iter)\n",
    "print(en, de)\n",
    "print(\"\\n After token transformation\\n\", \"-\"*50)\n",
    "en = token_transform['en'](en)\n",
    "de = token_transform['de'](de)\n",
    "print(en,\"\\n\", de)\n",
    "\n",
    "print(\"\\n After token to index transformation\\n\", \"-\"*50)\n",
    "en = vocab_transform['en'](en)\n",
    "de = vocab_transform['de'](de)\n",
    "print(en, \"\\n\", de)\n",
    "\n",
    "print(\"\\n After marking beginning and end of a sentence \\n\",\"-\"*50)\n",
    "en = bos_eos_pad(en)\n",
    "de = bos_eos_pad(de)\n",
    "print(en, \"\\n\", de)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Batching input\n",
    "\n",
    "As mentioned above, we will pad the sentences with `PAD_IDX` while batching several `text_transform`ed sentences together. \n",
    "In Pytorch, several cpu workers load observations (sentence pairs in our case) from the dataset. \n",
    "These observations are handed over [`torch.utils.data.DataLoader`](https://pytorch.org/docs/stable/data.html) which uses `collate_fn` to batch these observations together. \n",
    "In our case, we will overwrite the default `collate_fn` to transform and pad the input sentences with `PAD_IDX`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from https://pytorch.org/tutorials/beginner/translation_transformer.html#collation\n",
    "\n",
    "# function to collate data samples into batch tesors\n",
    "def collate_fn(batch):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        batch (iterator): each element is a source sentence and target sentence \n",
    "    \n",
    "    Returns:\n",
    "        src_batch (torch.tensor): [n_sentences x max_sentence_lenght]\n",
    "        tgt_batch (torch.tensor): [n_sentences x max_sentence_lenght]\n",
    "    \"\"\"\n",
    "    src_batch, tgt_batch = [], []\n",
    "    for src_sample, tgt_sample in batch:\n",
    "        src_batch.append(## USE APPROPRIATE TEXT TRANSFORM HERE )\n",
    "        tgt_batch.append(## USE APPROPRIATE TEXT TRANSFORM HERE)\n",
    "\n",
    "    src_batch = pad_sequence(src_batch, padding_value=PAD_IDX)\n",
    "    tgt_batch = pad_sequence(tgt_batch, padding_value=PAD_IDX)\n",
    "    return src_batch.permute(1,0), tgt_batch.permute(1,0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_iter = get_data_iter('train')\n",
    "batch = [next(train_iter) for _ in range(10)]\n",
    "src, tgt = collate_fn(batch)\n",
    "\n",
    "print(\"input shape:\", src.shape)\n",
    "print(\"output shape:\", tgt.shape)\n",
    "print(\"Note: each row is a single sentence. It begins with BOS_IDX and ends with an EOS_IDX. Each of them are padded with PAD_IDX towards the end.\")\n",
    "print(\"-\"*50)\n",
    "print(\"source:\")\n",
    "print(src)\n",
    "print('target:')\n",
    "print(tgt)\n",
    "print(\"-\"*50)\n",
    "print(\"Max length in the source batch:\", src.shape[1])\n",
    "print(\"Max length in the target batch:\", tgt.shape[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### A handy function to tests modules\n",
    "def get_sample_batch(type, n):\n",
    "    data_iter = get_data_iter(type)\n",
    "    batch = [next(data_iter) for _ in range(n)]\n",
    "    src, tgt = collate_fn(batch)\n",
    "    return src, tgt\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Masking Input\n",
    "\n",
    "A mask is needed to prevent the attention mechanism attending to unwanted parts of the input.\n",
    "We will come to specifics of masking later in this tutorial. \n",
    "Denoting the query matrix by `Q`, the key matrix by `K`, the attention matrix is given by $A' = softmax(\\frac{1}{\\sqrt{d}}QK^T)$.\n",
    "Given a mask $Mask$ of the same shape as the attention matrix $A'$ containing `-inf` at positions where the query should not attend to the key (and $0$ otherwise), the modified attention looks like $A'' = softmax(\\frac{1}{\\sqrt{d}}QK^T + Mask) $.\n",
    "Alternatively, one can use a boolean matrix as a mask and multiply $A'$ by it to achieve the same effect. \n",
    "Another alterative is to use mask  along with [`masked_fill`](https://pytorch.org/docs/stable/generated/torch.Tensor.masked_fill_.html) on $A'$.\n",
    "We will use `masked_fill`, so our masks will need to carry $0$ wherever we want attention scores to be zero. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Transformer Model\n",
    "\n",
    "Finally, we will focus on building the model itself.\n",
    "As a refresher, the Transformer consists of an  Encoder and a Decoder arranged as shown in the figure below \n",
    "\n",
    "<img src=\"img/transformer_architecture.jpeg\" width=500>\n",
    "\n",
    "Following modules will be required as we build a Transformer:\n",
    "- Attention \n",
    "- Multi-Head Attention (MHA)\n",
    "- Encoder layer (includes MHA, Feed-forward network and normalization layers)\n",
    "- Decoder layer (includes MHA and normalization layers)\n",
    "- Positional Encoder\n",
    "- Embedding \n",
    "- Transformer (includes all of the above)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Attention Mechanism\n",
    "\n",
    "Here we will assume $Mask$ is of the same shape as the final attention scores. We will come to building such masks later in the tutorial. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def attention(Q, K, V, mask=None, dropout=None):\n",
    "    \"\"\"\n",
    "    Implements attention mechanism.\n",
    "    \n",
    "    Args:\n",
    "        Q (torch.tensor): [batch_size x heads x max_sentence_length_query x dimension] Query matrix \n",
    "        K (torch.tensor): [batch_size x heads x max_sentence_length_key x dimension] Key matrix \n",
    "        V (torch.tensor): [batch_size x heads x max_sentence_length_value x dimension] Value matrix \n",
    "        mask (torch.tensor): [max_sentence_length x max_sentence_length] mask to prevent certain queries attending to certain keys\n",
    "        dropout (F.dropout): \n",
    "    \n",
    "    Returns:\n",
    "        (torch.tensor): Convex combination of V where weights are decided by the attention mechanism\n",
    "    \"\"\"\n",
    "    d_k = K.shape[-1]\n",
    "    \n",
    "    # final size: bs x n_heads x max_sentence_length_query x max_sentence_length_key\n",
    "    scores = ## REFER TO THE Equation of Attention above and write the appropriate code\n",
    "    \n",
    "    if mask is not None:\n",
    "        scores = scores.masked_fill(mask==0, -1e9)\n",
    "    \n",
    "    scores = F.softmax(scores, dim=-1)\n",
    "    \n",
    "    if dropout is not None:\n",
    "        scores = dropout(scores)\n",
    "\n",
    "    output = ## finally combine values (V) according to the attention scores \n",
    "    \n",
    "    return output\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# example call\n",
    "Q = torch.randn((2, 5, 6)) #  batches,  max lenght, dimensions\n",
    "K = torch.randn((2, 5, 6)) \n",
    "V = torch.randn((2, 5, 6)) \n",
    "mask = torch.triu(torch.ones((5, 5)))\n",
    "A = attention(Q, K, V, mask)\n",
    "print(A)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Multi-head attention\n",
    "\n",
    "We need to break down the computation of attention scores across several heads.\n",
    "Intuitively, each head captures a different context or meaning with which predcitions can be made. \n",
    "\n",
    "**NOTE**: There are no non-linearities in this entire module. There are only learnable parameters associated with affine transformations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadAttention(nn.Module):\n",
    "    def __init__(self, n_heads, d_model, dropout=0.1): \n",
    "        \n",
    "        super().__init__()\n",
    "        \n",
    "        self.n_heads = n_heads\n",
    "        self.d_model = d_model\n",
    "        self.d_head = d_model // n_heads\n",
    "        \n",
    "        self.Q = nn.Linear(d_model, d_model)\n",
    "        self.K = nn.Linear(d_model, d_model)\n",
    "        self.V = nn.Linear(d_model, d_model)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "    def forward(self, q, k, v, mask=None):\n",
    "        \n",
    "        batch_size = q.shape[0]\n",
    "        \n",
    "        # break the output such that the last two dimensions are heads x d_head\n",
    "        Q = ## Do a forward using using self.Q and reshape the tensor appropriately \n",
    "        K = ## Do a forward using using self.K and reshape the tensor appropriately \n",
    "        V = ## Do a forward using using self.V and reshape the tensor appropriately \n",
    "        \n",
    "        # take transpose so that last two dimensions are max_sentence_length x d_head (as required for the attention module)\n",
    "        Q = ## Take transpose\n",
    "        K = ## Take transpose\n",
    "        V = ## Take transpose\n",
    "        \n",
    "        # batch_size x n_heads x max_sentence_length x d_head\n",
    "        attn = attention(Q, K, V, mask, self.dropout)\n",
    "        \n",
    "        # transpose back to get batch_size x sentence_length x d_model\n",
    "        # use contiguous to reset the ordering of elements (i.e. stride and offset): https://stackoverflow.com/a/52229694/3413239\n",
    "        concat = attn.transpose(1, 2).contiguous().view(batch_size, -1, self.d_model)\n",
    "        \n",
    "        return concat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# example call\n",
    "\n",
    "mha = MultiHeadAttention(2, 6)\n",
    "\n",
    "Q = torch.randn((2, 5, 6)) #  batches,  max lenght, dimensions\n",
    "K = torch.randn((2, 5, 6)) \n",
    "V = torch.randn((2, 5, 6)) \n",
    "mask = torch.triu(torch.ones((5, 5)))\n",
    "A = mha(Q, K, V, mask)\n",
    "print(A)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Encoder Layer\n",
    "\n",
    "Now, we will write the code for a single Encoder layer.\n",
    "\n",
    "<img src=\"img/encoder_layer.png\" width=200>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EncoderLayer(nn.Module):\n",
    "    def __init__(self, d_model, n_heads, d_ff=1024, dropout=0.1):\n",
    "        \n",
    "        super().__init__()\n",
    "        \n",
    "        # from bottom up --> \n",
    "        \n",
    "        # mha & norm\n",
    "        self.mha = MultiHeadAttention(n_heads, d_model)\n",
    "        self.norm1 = ## Define a layernorm here: refer https://pytorch.org/docs/stable/generated/torch.nn.LayerNorm.html\n",
    "        \n",
    "        # feed forward & norm\n",
    "        self.ff = nn.Sequential(\n",
    "            ### Define a  single layer feed forward network here\n",
    "            ## Use linear --> activation --> dropout --> linear\n",
    "        )\n",
    "        self.norm2 =  ## Define a layernorm here: refer https://pytorch.org/docs/stable/generated/torch.nn.LayerNorm.html\n",
    "        \n",
    "        # dropout regularization\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "    def forward(self, x, mask):\n",
    "        ## What should be the query, key and values\n",
    "        z = self.mha(?, ?, ? mask) # mha \n",
    "        z = self.dropout(z) # refer to the section 5.4 of Vaswani et al. \n",
    "        x = ## Write code to add & norm1\n",
    "        \n",
    "        z = self.ff(x) # feed forward\n",
    "        z = self.dropout(z) # refer to the section 5.4 of Vaswani et al. \n",
    "        x = ## Write code to add & norm2\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# example call \n",
    "d = 6\n",
    "enc = EncoderLayer(d, 2, 12)\n",
    "x = torch.randn((2, 5, d)) #  batches,  max length, dimension\n",
    "mask = torch.triu(torch.ones((5, 5)))\n",
    "enc(x, mask)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Decoder Layer\n",
    "\n",
    "Now we will write the code for the decoder layer.\n",
    "\n",
    "<img src=\"img/decoder_layer.png\" width=200>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DecoderLayer(nn.Module):\n",
    "    def __init__(self, d_model, n_heads, d_ff=1024, dropout=0.1):\n",
    "        \n",
    "        super().__init__()\n",
    "        \n",
    "        # from bottom-up ---> \n",
    "        \n",
    "        # masked-mha\n",
    "        self.mha1 = MultiHeadAttention(n_heads, d_model)\n",
    "        self.norm1 = ## Define a layernorm here: refer https://pytorch.org/docs/stable/generated/torch.nn.LayerNorm.html\n",
    "        \n",
    "        # mha\n",
    "        self.mha2 = MultiHeadAttention(n_heads, d_model)\n",
    "        self.norm2 = ## Define a layernorm here: refer https://pytorch.org/docs/stable/generated/torch.nn.LayerNorm.html\n",
    "                \n",
    "        # feed-forward\n",
    "        self.ff = nn.Sequential(\n",
    "            ### Define a  single layer feed forward network here\n",
    "            ## Use linear --> activation --> dropout --> linear\n",
    "        )\n",
    "        self.norm3 = ## Define a layernorm here: refer https://pytorch.org/docs/stable/generated/torch.nn.LayerNorm.html\n",
    "        \n",
    "        # dropout regularization\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x, e_x, src_mask, tgt_mask):\n",
    "        ## What should be the query, key and values\n",
    "        z = self.mha1(?, ?, ?, tgt_mask) # masked mha  \n",
    "        z = self.dropout(z) # refer to the section 5.4 of Vaswani et al. \n",
    "        x = ## Write code to add & norm1\n",
    "        \n",
    "        ## What should be the query, key and values\n",
    "        z = self.mha2(?, ?, ?, src_mask) # mha \n",
    "        z = self.dropout(z) # refer to the section 5.4 of Vaswani et al. \n",
    "        x = ## Write code to add & norm2\n",
    "        \n",
    "        z = self.ff(x)\n",
    "        z = self.dropout(z) # refer to the section 5.4 of Vaswani et al. \n",
    "        x = ## Write code to add & norm3\n",
    "        \n",
    "        return x  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# example call \n",
    "d = 6\n",
    "dec = DecoderLayer(d, 2, 12)\n",
    "x = torch.randn((2, 5, d)) #  batches,  max length, dimension\n",
    "e_x =  torch.randn((2, 5, d))\n",
    "mask = torch.triu(torch.ones((5, 5)))\n",
    "dec(x, e_x, mask, mask)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Input Embedding\n",
    "\n",
    "Now we will write a class that can handle learnable word embeedings for us. \n",
    "This class will hold word vectors of appropriate dimensions, where values in each dimensions are learned via backpropagation. \n",
    "Checkout [`torch.nn.Embedding`](https://pytorch.org/docs/stable/generated/torch.nn.Embedding.html) to do this.\n",
    "\n",
    "<img src=\"img/embedding.png\" width=500>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TokenEmbedding(nn.Module):\n",
    "    def __init__(self, n_tokens, emb_size):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.embedding = ## Define nn.Embedding of appropriate dimensions\n",
    "        self.emb_size = emb_size \n",
    "    \n",
    "    def forward(self, tokens):\n",
    "        return self.embedding(tokens.long()) * math.sqrt(self.emb_size)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The reasons for multiplying embeddings with `math.sqrt(self.emb_size)` are hypothesized here: https://datascience.stackexchange.com/a/87909/1165\n",
    "\n",
    "A plausible explanation seems like - \n",
    "\n",
    "*It is to make the positional encoding relatively smaller. This means the original meaning in the embedding vector won’t be lost when we add them together* because positional encoding can't have values greater than 1 (see below). "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Positional Encoding\n",
    "\n",
    "Poitional Encodings are used to inject dependence on word ordering. This is done via adding a position dependent vector. Such a vector is constructed as follows - \n",
    "\n",
    "A word at $pos$ position has $i^{th}$ dimension encoded as following\n",
    "\n",
    "$$\n",
    "\\begin{equation*}\n",
    "PE(pos, i) = \\begin{cases}\n",
    "\\sin\\Big(\\frac{pos}{10000^{\\frac{i}{d_{model}}}}\\Big), &\\text{if i is even} \\\\ \\\\\n",
    "\\cos\\Big(\\frac{pos}{10000^{\\frac{i-1}{d_{model}}}}\\Big), &\\text{if i is odd} \\\\\n",
    "\\end{cases}\n",
    "\\end{equation*}\n",
    "$$\n",
    "\n",
    "Note that the exponent of $10000$ will be same for consecutive even and odd numbers.\n",
    "\n",
    "\n",
    "<img src=\"img/pe.png\" width=500>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_position_encodings(max_len, d):\n",
    "    \"\"\"\n",
    "    Creates encoding\n",
    "    \n",
    "    Args:\n",
    "        max_len (int): number of positions to encode\n",
    "        d (int): dimension of encoding\n",
    "    \n",
    "    Returns:\n",
    "        (torch.tensor): [max_len, d] each row is the encoding for that row position\n",
    "    \"\"\"\n",
    "    exp_coeff = torch.arange(0, d, 2) / d\n",
    "    arg = ## Write code to compute 1/denominator in the above formula  # take exp of log for numerical stability\n",
    "    \n",
    "    pos = torch.arange(0, max_len).reshape(max_len, 1)\n",
    "    pos_embedding = torch.zeros((max_len, d))\n",
    "    pos_embedding[:, 0::2] = # take sine; use torch.sin\n",
    "    pos_embedding[:, 1::2] = # take cos; use torch.cos\n",
    "    return pos_embedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "enc_size = 200\n",
    "n_pos = 100\n",
    "pos_enc = get_position_encodings(n_pos, enc_size)\n",
    "\n",
    "# check whether all rows are unique `torch.unique`\n",
    "print(f\"number of unique rows in encoding: {pos_enc.unique(dim=0).shape[0]}\")\n",
    "print(f\"number of positions: {n_pos}\")\n",
    "\n",
    "\n",
    "# how does the encoding looks like?\n",
    "fig, ax = plt.subplots(nrows=1, ncols=1, figsize=(10, 5), dpi=100)\n",
    "ax.matshow(pos_enc)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Masking for Machine Translation Task (as required by Transformer)\n",
    "\n",
    "To simplify our discussion here, we will consider attention scores for a single sentence and one head. \n",
    "Let's denote the attention scores by $A$ (dimension: `query_length x key_length`) such that $A_{qk}$ denotes how much weight does the $q^{th}$ query gives to the $k^{th}$ key.\n",
    "When query, keys, and values are computed from the same input, it is termed as self-attention.\n",
    "Masking lets us nullify the weights of certain queries on certain keys. \n",
    "\n",
    "There are three MHA units in Transformer and the corresponding masks required by them are as follows:\n",
    "\n",
    "\n",
    "**(A) Source padding mask**:\n",
    "\n",
    "The MHA unit in the Encoder uses self-attention on source sentence, therefore, `query_length` and `key_length` are the same as number of positions in the source sentence. \n",
    "For this MHA unit, $A_{qk}$ denotes the weight of the $q^{th}$-token in the sentence on the $k^{th}$-token in the same sentence. \n",
    "Naturally, we don't want any $q^{th}$-token to attend to irrelevant padding-tokens, i.e., $\\forall k \\text{  s.t.  } src[k] == PAD\\text{_}IDX$.\n",
    "Thus, such a mask should contain 0 wherever either of the query or key token is a padding token.\n",
    "\n",
    "\n",
    "**(B) Target-to-Source padding mask**:\n",
    "\n",
    "The MHA unit that connects encoder to decoders act upon target sentence's embeddings as queries and encoded embeddings as keys and values. \n",
    "The idea is to let the (partially) translated embedding (query) attend to the embeddings of all the tokens (keys and values) in the source sentence.\n",
    "Therefore, $A_{qk}$ denotes the weight of $q^{th}$-query in the translated sentence on the $k^{th}$-key in the source sentence. \n",
    "Once again, we want to prevent queries from attending to keys that correspond to the padding tokens in the source sentence.\n",
    "Thus, such a mask should contain 0 wherever there is a padding token in the source sentence. \n",
    "This will look exactly as the mask explained in (A) above. \n",
    "\n",
    "**(C) Target no-peak mask**: \n",
    "\n",
    "Machine translation task is mathematically modeled as Conditional Language Models (CLM). We are interested in modeling the likelihood of an output sequence $\\{y_i\\}_{i=1}^{t}$ conditioned on an input $x$, i.e, $P(y_1, y_2, ..., y_t | x)$. This is in contrast to language models (LM) where we are interested in modeling the unconditional probability $p(y_1, y_2, ..., y_t) = \\Pi_{t=1}^{n}p(y_t | y_{<t})$ of a sequence $\\{y_i\\}_{i=1}^{t}$. CLMs are modeled similarly as $P(y_1, y_2, ..., y_t | x) = \\Pi_{t=1}^{n}p(y_t | y_{<t}, x)$. Therefore, we need to predict one token at a time, $y_t$, conditioned on the past outputs $\\{y_{<t}\\}$ and the source $x$. \n",
    "\n",
    "\n",
    "[*Teacher forcing*](https://machinelearningmastery.com/teacher-forcing-for-recurrent-neural-networks/): Transformer naturally uses teacher forcing mechanism, wherein the ground truth for $\\{y_{<t}\\}$ are used as an input, and the task is to predict $y_t$ conditioned on $\\{y_{<t}\\}$  and $x$.\n",
    "\n",
    "The output of the Masked-MHA unit in the decoder is such that, for each sentence, the embeddings at the $q^{th}$ index is the weighted (by attention scores) sum of values (target embeddings). \n",
    "For example, the Masked-MHA's embedding at $q=0$  will represent value-embedding corresponding to the `<BOS>` token in the target `TokenEmbeddings`. \n",
    "Note that the $q=0$ embedding will be same for all the sentences. \n",
    "With the help of next MHA unit that connects encoder and the decoder, the output embeddings corresponding to $q=0$ will be used to predict the next token in the target language. \n",
    "\n",
    "Similarly, the Masked-MHA's embedding at $q=1$  will represent weighted sum of value-embeddings corresponding to the target `TokenEmbeddings` of `<BOS>` token and the ground-truth 1st token.\n",
    "The subsequent MHA unit will output embeddings for $q=1$ that can be used to predict the next token (i.e, 2nd token) in the target language. \n",
    "\n",
    "Thus, Masked-MHA's mask will be such that $Mask_{qk} = 0 \\ \\ \\forall \\ \\ k \\text{  s.t.  } k > q$.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "src, tgt = get_sample_batch('train', 10)\n",
    "\n",
    "src_padding_mask =  (src != PAD_IDX).unsqueeze(1).unsqueeze(2)\n",
    "print(src_padding_mask.shape)\n",
    "\n",
    "N, L = tgt.shape\n",
    "tgt_mask = torch.tril(torch.ones((L, L)), diagonal=0).expand(N, 1, L, L)\n",
    "print(tgt_mask.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_masks(src, tgt):\n",
    "    src_padding_mask =  (src != PAD_IDX).unsqueeze(1).unsqueeze(2)\n",
    "    \n",
    "    N, L = tgt.shape\n",
    "    tgt_mask = torch.tril(torch.ones((L, L)), diagonal=0).expand(N, 1, L, L)\n",
    "    \n",
    "    return src_padding_mask, tgt_mask\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Transformer\n",
    "\n",
    "Putting it all together ... \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "d_model = 6\n",
    "\n",
    "src_embedding = TokenEmbedding(SRC_N_TOKENS, d_model)\n",
    "tgt_embedding = TokenEmbedding(TGT_N_TOKENS, d_model)\n",
    "pos_encoding = get_position_encodings(20, d_model)\n",
    "\n",
    "\n",
    "src, tgt = get_sample_batch('train', 10)\n",
    "src_padding_mask, tgt_mask = create_masks(src, tgt)\n",
    "\n",
    "print(\"Source sentences shape: \", src.shape)\n",
    "print(\"Target sentences shape: \", tgt.shape)\n",
    "print(\"Source embedding shape: \", src_embedding(src).shape)\n",
    "print(\"source padding mask shape:\", src_padding_mask.shape)\n",
    "print(\"target mask shape:\", tgt_mask.shape)\n",
    "\n",
    "x = src_embedding(src) + ## Add appropriate pos encodings here\n",
    "print(\"Source embeddings' shape after adding positional encodings:\", x.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = src_embedding(src) + ## Add appropriate pos encodings here\n",
    "z = tgt_embedding(tgt) + ## Add appropriate pos encodings here\n",
    "\n",
    "n_heads = 2\n",
    "encoders = [\n",
    "        EncoderLayer(d_model, n_heads), \n",
    "        EncoderLayer(d_model, n_heads)\n",
    "]\n",
    "\n",
    "x = encoders[0](x, src_padding_mask)\n",
    "x = encoders[1](x, src_padding_mask)\n",
    "\n",
    "decoders = [\n",
    "    DecoderLayer(d_model, n_heads),\n",
    "    DecoderLayer(d_model, n_heads)\n",
    "]\n",
    "\n",
    "z = decoders[0](z, x, src_padding_mask, tgt_mask)\n",
    "z = decoders[1](z, x, src_padding_mask, tgt_mask)\n",
    "\n",
    "print(\"pre-softmax outputs' shape: \", z.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Transformer(nn.Module):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "       src_n_tokens (int): vocabulary size or number of unique tokens in the source language\n",
    "       tgt_n_tokens (int): vocabulary size or number of unique tokens in the target language\n",
    "       n_encoders (int): number of encoders to stack on top of each other\n",
    "       n_decoders (int): number of decoders to stack on top of each other\n",
    "       d_model (int): embedding dimension size\n",
    "       n_heads (int): number of heads in multi-head attention\n",
    "       d_ff (int): hidden dimension of feed forward networks\n",
    "       dropout (float): dropout probability\n",
    "       max_length (int): maximum length for positional encodings (optional).\n",
    "    \"\"\"\n",
    "    def __init__(self, src_n_tokens, tgt_n_tokens, n_encoders, n_decoders,  d_model, n_heads, d_ff=1024, dropout=0.1, max_length=1000):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.input_args = [src_n_tokens, tgt_n_tokens, n_encoders, n_decoders,  d_model, n_heads, d_ff, dropout, max_length]\n",
    "        \n",
    "        # embedddings + positional encoding\n",
    "        self.src_embeddings = ## DEFINE TOKEN EMBEDDINGS\n",
    "        self.tgt_embeddings = ## DEFINE TOKEN EMBEDDINGS\n",
    "        self.register_buffer('positional_encodings', nn.Parameter(get_position_encodings(max_length, d_model)))\n",
    "        \n",
    "        # encoders\n",
    "        self.encoders = nn.ModuleList([\n",
    "                ## Define n_encoders EncoderLayers here\n",
    "        ])\n",
    "        \n",
    "        # decoders\n",
    "        self.decoders = nn.ModuleList([\n",
    "            ## Define n_decoders DecoderLayer here\n",
    "        ])\n",
    "        \n",
    "        # final linear (pre-softmax)\n",
    "        self.out = nn.Sequential(\n",
    "            ## difine a linear transformation with appropriate dimensions without softmax\n",
    "            ## softmax is implicit in the loss function\n",
    "        )\n",
    "        \n",
    "        # dropout regularization\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, src, tgt, src_padding_mask, tgt_mask):\n",
    "        # encode\n",
    "        e_x = self.encode(src, src_padding_mask)\n",
    "        \n",
    "        # decode \n",
    "        z = self.decode(tgt, e_x, src_padding_mask, tgt_mask)\n",
    "        \n",
    "        # linear (pre-softmax)\n",
    "        out = self.out(z)\n",
    "        \n",
    "        return out\n",
    "\n",
    "    def encode(self, src, src_padding_mask):\n",
    "        \"\"\"\n",
    "        Returns the encodings for src sentence\n",
    "        \"\"\"\n",
    "        # encoder inputs\n",
    "        src_emb = ## define src embeddings\n",
    "        src_emb = self.dropout(src_emb) # dropout described in section 5.4 of Vaswani et al.\n",
    "        \n",
    "        # encoder\n",
    "        for encoder in self.encoders:\n",
    "            e_x = encoder(?, ?)\n",
    "        \n",
    "        return e_x\n",
    "    \n",
    "    def decode(self, tgt, e_x, src_padding_mask, tgt_mask):\n",
    "        \"\"\"\n",
    "        Returns the pre-softmax decoder outputs\n",
    "        \"\"\"\n",
    "        # decoder inputs \n",
    "        tgt_emb = ## define tgt embeddings\n",
    "        tgt_emb = self.dropout(tgt_emb) # dropout described in section 5.4 of Vaswani et al.\n",
    "        \n",
    "        # decoder \n",
    "        z = tgt_emb\n",
    "        for decoder in self.decoders:\n",
    "            z = decoder(?, ?, ?, ?)\n",
    "        \n",
    "        return z\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# example call\n",
    "model = Transformer(SRC_N_TOKENS, TGT_N_TOKENS, 2, 2, 6, 2)\n",
    "\n",
    "src, tgt = get_sample_batch('train', 10)\n",
    "src_padding_mask, tgt_mask = create_masks(src, tgt)\n",
    "\n",
    "out = model(src, tgt, src_padding_mask, tgt_mask)\n",
    "print(\"shape of output: \", out.shape)\n",
    "print(\"number of tokens in the target language:\", TGT_N_TOKENS)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Number of parameters\n",
    "\n",
    "Let's look at the number of learnable parameters in the Transformer with the configurations same as in  [Vaswani et al.](https://arxiv.org/abs/1706.03762)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def mem_size(model):\n",
    "    \"\"\"\n",
    "    Get model size in GB (as str: \"N GB\")\n",
    "    \"\"\"\n",
    "    mem_params = sum(\n",
    "        [param.nelement() * param.element_size() for param in model.parameters()]\n",
    "    )\n",
    "    mem_bufs = sum([buf.nelement() * buf.element_size() for buf in model.buffers()])\n",
    "    mem = mem_params + mem_bufs\n",
    "    return f\"{mem / 1e9:.4f} GB\"\n",
    "\n",
    "def num_params(model):\n",
    "    \"\"\"\n",
    "    Print number of parameters in model's named children\n",
    "    and total\n",
    "    \"\"\"\n",
    "    s = \"Number of parameters:\\n\"\n",
    "    n_params = 0\n",
    "    for name, child in model.named_children():\n",
    "        n = sum(p.numel() for p in child.parameters())\n",
    "        s += f\"  • {name:<15}: {n}\\n\"\n",
    "        n_params += n\n",
    "    s += f\"{'total':<19}: {n_params}\"\n",
    "\n",
    "    return s\n",
    "\n",
    "def pp_model_summary(model):\n",
    "    print(num_params(model))\n",
    "    print(f\"{'Total memory':<18} : {mem_size(model)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N_ENCODERS = 6\n",
    "N_DECODERS = 6\n",
    "EMB_SIZE = 512\n",
    "N_HEADS = 8\n",
    "FFN_HID_DIM = 512\n",
    "\n",
    "model = Transformer(SRC_N_TOKENS, TGT_N_TOKENS, N_ENCODERS, N_DECODERS, EMB_SIZE, N_HEADS, FFN_HID_DIM)\n",
    "pp_model_summary(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training the transformer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Scaling down the transformer \n",
    "\n",
    "The original configuration is meant to be for a bigger dataset (refer section 5.1 and 5.2 of [Vaswani et al.](https://arxiv.org/abs/1706.03762)). \n",
    "Since we are dealing with a toy problem we will scale the transformer down by reducing the number of encoders and embedding dimensions.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "N_ENCODERS = 3\n",
    "N_DECODERS = 3\n",
    "EMB_SIZE = 128\n",
    "N_HEADS = 8\n",
    "FFN_HID_DIM = 256\n",
    "\n",
    "model = Transformer(SRC_N_TOKENS, TGT_N_TOKENS, N_ENCODERS, N_DECODERS, EMB_SIZE, N_HEADS, FFN_HID_DIM)\n",
    "pp_model_summary(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Loss function\n",
    "\n",
    "The authors use Label Smoothing with $\\epsilon_{ls} = 0.1$, which makes the prediction targets less peaky as compared to one-hot targets. \n",
    "Thus, instead of measuring our loss against one-hot targets, we will be measuring it against a smoothed distribution where the value at the target class will be  confidence,$C = 1 - \\epsilon_{ls}$ and smotthing, $S = \\frac{\\epsilon_{ls}}{n_{classes} - 1}$ otherwise.\n",
    "\n",
    "Refer to [Muller et al. to](https://arxiv.org/abs/1906.02629) understand - Why does Label Smoothing help?\n",
    "\n",
    "\n",
    "Assuming that the model outputs softmax outputs for a single sentence as $\\hat{y}$ and the prediction is the $t^{th}$ class such that $y_t = 1$ and 0 otherwise.\n",
    "\n",
    "$$\n",
    "\\begin{equation*}\n",
    "\\begin{split}\n",
    "\\mathcal{L}_{ls}(\\hat{y}, y) &= CE(\\ C \\times y + S \\times (1-y)\\ |\\ \\hat{y}\\ ) \\\\\n",
    "& = CE(\\ (C - S) y + S \\ \\mathbf{1} \\ |\\   \\hat{y}\\ ) \\\\\n",
    "& = CE(\\ (C - S) y \\ | \\ \\hat{y}\\ ) + CE(\\ S \\ \\mathbf{1} \\ | \\ \\hat{y}) \\\\\n",
    "& = -(C - S) \\times \\log{\\hat{y}_t} - S \\times \\sum_{i} \\log{\\hat{y}_i} \\\\ \n",
    "\\end{split}\n",
    "\\end{equation*}\n",
    "$$\n",
    "\n",
    "\n",
    "Here $CE(p, q) = -\\sum_{x}p(x)\\log{q(x)}$ is the standard cross-entropy loss.\n",
    "\n",
    "As $\\log{\\hat{y}_i}$ is a log of softmax operation, and [`nn.LogSoftmax`](https://pytorch.org/docs/stable/generated/torch.nn.LogSoftmax.html) is a numerically stable operation to compute log-softmax, we will use that instead.\n",
    "To do this, we will make the model output pre-softmax logits only. \n",
    "\n",
    "Such a modified loss function can be implemented as an `nn.Module` with the output of its `forward` function as a scalar value. \n",
    "Refer to [this stackoverflow post](https://stackoverflow.com/a/66773267/3413239) on various ways to implement this loss function.\n",
    "\n",
    "[[Muller et al. 2019] Why does Label Smoothing help?](https://arxiv.org/abs/1906.02629)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# adapted from https://stackoverflow.com/a/66773267/3413239\n",
    "\n",
    "class LabelSmoothingLoss(nn.Module):\n",
    "    def __init__(self, classes, epsilon=0.0):\n",
    "        super().__init__()\n",
    "        self.smoothing = epsilon / (classes - 2) # classes to not consider for smoothing:  PAD_IDX class, target class\n",
    "        self.confidence = 1 - epsilon \n",
    "    \n",
    "    def forward(self, x, target):\n",
    "        \n",
    "        # mask out PAD_IDX\n",
    "        tgt_padding_mask = (target != PAD_IDX).unsqueeze(2)\n",
    "        tgt_padding_mask = tgt_padding_mask.to(x.device)\n",
    "        x = x.masked_fill(tgt_padding_mask==0, 0)\n",
    "    \n",
    "        # compute CE loss on smoothed distribution\n",
    "        logprobs = torch.nn.functional.log_softmax(x, dim=-1)\n",
    "        target_loss = -logprobs.gather(dim=-1, index=target.unsqueeze(2)).squeeze(2)\n",
    "        smooth_loss = -logprobs.sum(dim=-1)\n",
    "        loss = (self.confidence - self.smoothing) * target_loss + self.smoothing * smooth_loss\n",
    "        return loss.mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# example call\n",
    "loss = LabelSmoothingLoss(TGT_N_TOKENS, 0.1)\n",
    "\n",
    "model = Transformer(SRC_N_TOKENS, TGT_N_TOKENS, 2, 2, 6, 2)\n",
    "\n",
    "src, tgt = get_sample_batch('train', 10)\n",
    "src_padding_mask, tgt_mask = create_masks(src, tgt)\n",
    "out = model(src, tgt, src_padding_mask, tgt_mask)\n",
    "\n",
    "print(\"shape of output: \", out.shape)\n",
    "print(\"number of tokens in the target language:\", TGT_N_TOKENS)\n",
    "print(\"tgt shape: \", tgt.shape)\n",
    "\n",
    "loss(out, tgt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Iterate over a dataloader\n",
    "\n",
    "We will write a function that iterates over a batch of inputs and computes an appropriate loss over that dataset. \n",
    "This will help us save from writing redundant code for validation and test datasets. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process(model, data_loader, loss_fn, optimizer=None):\n",
    "    \n",
    "    running_loss, n_samples = 0, 0\n",
    "    \n",
    "    for src, tgt in data_loader:\n",
    "        \n",
    "        n_samples += src.shape[0] # number of sentences \n",
    "        \n",
    "        # \n",
    "        src = src.to(device)\n",
    "        tgt = tgt.to(device)\n",
    "\n",
    "        # ignore last token in the tgt because at the time of predicting the last token, available inputs will be everything before that token\n",
    "        tgt_input = ### Define target input\n",
    "\n",
    "        # create masks\n",
    "        src_padding_mask, tgt_mask = ## create masks\n",
    "        src_padding_mask = src_padding_mask.to(device)\n",
    "        tgt_mask = tgt_mask.to(device)\n",
    "        \n",
    "        # forward pass\n",
    "        logits = ## call model\n",
    "        \n",
    "        # target predictions are made for 1st token onwards. 0th token is assumed to be <BOS> by design of the model. Read no-peak masking (above) for more details.\n",
    "        tgt_out = ## define target outputs\n",
    "    \n",
    "        # compute loss\n",
    "        loss = loss_fn(logits, tgt_out)\n",
    "        \n",
    "        # backward pass \n",
    "        if optimizer is not None:\n",
    "            ## Perform backward pass\n",
    "        \n",
    "        running_loss += loss.item()\n",
    "    \n",
    "    return running_loss / n_samples\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define Optimizer\n",
    "Section 5.3 of [Vaswani et al.](https://arxiv.org/abs/1706.03762) discusses the the learning rate (LR) scheduling for Adam optimizer. \n",
    "We will implement such a LR scheduler now. \n",
    "\n",
    "For clarity, the suggested learning rate at step $t$ is given by,\n",
    "\n",
    "$$lr(t) = \\frac{1}{\\sqrt{d_{model}}} \\times min(\\frac{1}{\\sqrt{t}}, \\ \\ t \\times \\frac{1}{WS^{1.5}}), $$\n",
    "\n",
    "where $WS$ is the number of warm up steps before which the learning rate increases linearly. \n",
    "Thereafter, it decreases proportionally to the \"inverse square root of the the step number\".\n",
    "\n",
    "\n",
    "Following code is motivated from [The Annotated Transformer](https://nlp.seas.harvard.edu/2018/04/03/attention.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class WarmupOpt:\n",
    "    def __init__(self, optimizer, d_model, warmup_steps):\n",
    "        \n",
    "        self.optimizer = optimizer\n",
    "        self.d_model = d_model\n",
    "        self.ws = warmup_steps\n",
    "        \n",
    "        self.t = 0\n",
    "        self._rate = 0\n",
    "        \n",
    "    def step(self):\n",
    "        \"\"\"\n",
    "        the call to optimizer.step comes here first\n",
    "        \"\"\"\n",
    "        self.t += 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",
    "        if step is None:\n",
    "            step = self.t\n",
    "\n",
    "        return ## implement the learning rate equation above\n",
    "    \n",
    "    def zero_grad(self):\n",
    "        self.optimizer.zero_grad()\n",
    "\n",
    "        \n",
    "\n",
    "# refer to section 5.3 of Vaswani et al. \n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.0001, betas=(0.9, 0.98), eps=1e-9)\n",
    "optimizer = WarmupOpt(optimizer, d_model=EMB_SIZE, warmup_steps=4000)\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train\n",
    "\n",
    "Finally, lets write the whole training loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(model, loss_fn):\n",
    "    model.to(device)\n",
    "    \n",
    "    # fix seed for reproducibility \n",
    "    rng = np.random.RandomState(1)\n",
    "    torch.manual_seed(rng.randint(np.iinfo(int).max))\n",
    "\n",
    "    # create a model directory to store the best model\n",
    "    model_dir = pathlib.Path(\"./models\").resolve()\n",
    "    if not model_dir.exists():\n",
    "        model_dir.mkdir()\n",
    "\n",
    "    # optimizer\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=0.0001, betas=(0.9, 0.98), eps=1e-9)\n",
    "    optimizer = WarmupOpt(optimizer, d_model=EMB_SIZE, warmup_steps=4000)\n",
    "\n",
    "    # validation dataloader \n",
    "    val_dataloader = torch.utils.data.DataLoader(list(get_data_iter('valid')), batch_size=BATCH_SIZE, collate_fn=collate_fn)\n",
    "    \n",
    "    \n",
    "    # logging\n",
    "    best_val_loss = np.inf\n",
    "    train_losses, val_losses = [], []\n",
    "    no_improvement_cnt = 0\n",
    "    \n",
    "    for epoch in range(NUM_EPOCHS):\n",
    "        print(f\"@ epoch {epoch}\", end=\"\")\n",
    "\n",
    "        # training loss\n",
    "        train_dataloader = torch.utils.data.DataLoader(get_data_iter('train'), batch_size=BATCH_SIZE, collate_fn=collate_fn)\n",
    "        train_loss = process(model, train_dataloader, loss_fn, optimizer)\n",
    "\n",
    "        # validation loss\n",
    "        with torch.no_grad():\n",
    "            val_loss = process(model, val_dataloader, loss_fn)\n",
    "\n",
    "        # save the best model\n",
    "        if val_loss < best_val_loss:\n",
    "            best_val_loss = val_loss\n",
    "            torch.save(model.state_dict(), model_dir / \"best.ckpt\")\n",
    "            no_improvement_cnt = 0\n",
    "        else:\n",
    "\n",
    "            # if there has been no improvement in validation loss, stop early\n",
    "            no_improvement_cnt += 1\n",
    "\n",
    "            if no_improvement_cnt % 10 == 0:\n",
    "                print(\"\\nEarly stopping!\")\n",
    "                break\n",
    "\n",
    "        # logging\n",
    "        train_losses.append(train_loss)\n",
    "        val_losses.append(val_loss)\n",
    "        print(f\"\\ttrain_loss: {train_loss: .5f}, val_loss: {val_loss:.5f}\")\n",
    "\n",
    "    print(f\"best val loss: {best_val_loss:.5f}\")\n",
    "\n",
    "    # load the best model\n",
    "    model = model.__class__(*model.input_args)\n",
    "    model.load_state_dict(torch.load(model_dir / \"best.ckpt\"))\n",
    "    model = model.to(device) \n",
    "    \n",
    "    metrics = {\n",
    "        'train_losses': train_losses,\n",
    "        'val_losses': val_losses,\n",
    "    }\n",
    "    return model, metrics  \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "NUM_EPOCHS = 50\n",
    "BATCH_SIZE = 128 \n",
    "\n",
    "N_ENCODERS = 3\n",
    "N_DECODERS = 3\n",
    "EMB_SIZE = 128\n",
    "N_HEADS = 8\n",
    "FFN_HID_DIM = 256\n",
    "model = Transformer(SRC_N_TOKENS, TGT_N_TOKENS, N_ENCODERS, N_DECODERS, EMB_SIZE, N_HEADS, FFN_HID_DIM)\n",
    "pp_model_summary(model)\n",
    "\n",
    "# initialization as in https://pytorch.org/tutorials/beginner/translation_transformer.html\n",
    "for p in model.parameters():\n",
    "    if p.dim() > 1:\n",
    "        nn.init.xavier_uniform_(p)\n",
    "\n",
    "        \n",
    "# loss function \n",
    "loss_fn = LabelSmoothingLoss(classes=TGT_N_TOKENS, epsilon=0.1) # refer to section 5.4 of Vaswani et al.\n",
    "\n",
    "model, metrics = train(model, loss_fn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot losses\n",
    "fig, axs = plt.subplots(nrows=1, ncols=1, figsize=(10,5), dpi=100)\n",
    "\n",
    "## PLot losses\n",
    "axs.plot(?, color=\"#BDD9BF\", marker=\"o\", label=\"Train loss\", linestyle=\":\", linewidth=2)\n",
    "axs.plot(?, color=\"#A997DF\", marker=\"o\", label=\"Val loss\", linestyle=\":\", linewidth=2)\n",
    "axs.set_ylabel(\"loss\", fontsize=20)\n",
    "\n",
    "axs.set_xlabel(\"Epochs\", fontsize=20)\n",
    "\n",
    "# tick size\n",
    "for tick in axs.xaxis.get_major_ticks():\n",
    "    tick.label.set_fontsize(15)\n",
    "\n",
    "for tick in axs.yaxis.get_major_ticks():\n",
    "    tick.label.set_fontsize(15)\n",
    "    \n",
    "# legend\n",
    "axs.legend(prop={\"size\":15})\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inference\n",
    "\n",
    "Now that we have a trained model, we need to use it to perform translation. \n",
    "For this, we will need to go back to the definition of CLMs, which models the likelihood of observing a sequence of words, i.e. $P(y_1, y_2, ..., y_t \\ |\\ x) = \\prod\\limits_{i=1}^{t} p(y_{i} \\ | \\ y_{<i}, x)$.\n",
    "Thus, the inference procedure aims at finding the sequence $\\{\\hat{y}_i\\}_{i=1}^t$ with the maximum likelihood conditioned on $x$, i.e., $\\hat{y} =\\underset{y_1, y_2, ..., y_t}{\\operatorname{arg max}}P(y_1, y_2, ..., y_t \\ |\\ x) = \\underset{y_1, y_2, ..., y_t}{\\operatorname{arg max}}\\ \\prod\\limits_{i=1}^{t} p(y_{i} \\ | \\ y_{<i}, x)$.\n",
    "\n",
    "Since we formulated the problem as sequence prediction, i.e., our model yields $p(y_{i} \\ | \\ y_{<i}, x)$, we will make this inference sequentially.\n",
    "Hence, finding the best sequence $\\{y_i\\}_{i=1}^{t}$ is NP-Hard as one will have to check all the possibilities before finding the sequence with the maximum likelihood.\n",
    "Note that for a sequence of length $L$ if there are $N$ possible values, there will be a total of $N^{L}$ possibile sequences.\n",
    "For example, if there are $N=10^5$ unique tokens and a sequence is of length $L=5$, there will be $10^{25}$ possible sequences. \n",
    "As a result we will have to fall back on approximation algorithms to find a maximum likelihood sequence. \n",
    "\n",
    "One can view the above search problem as drawing inference on a tree of depth $L$ with the root node as `<BOS>` token. \n",
    "Each node will have $N$ children, i.e., $N$ possible next token. \n",
    "\n",
    "<img src=\"img/tree.png\" width=1000>\n",
    "\n",
    "Finding the most likely sequence then boils down to moving from a node to another.\n",
    "Some of the ways to go about it are:\n",
    "- **Greedy Decoding**: Move to the most likely node given the current node in the tree. This results in the sequence $\\prod\\limits_{i=1}^{t}\\underset{y_i}{\\operatorname{arg max}}p(y_i \\ | \\ y_{<i}, x)$. Note it is **not equivalent to finidng the maximum likleihood sequence** because $\\underset{y_1, y_2, ..., y_t}{\\operatorname{arg max}}\\ \\prod\\limits_{i=1}^{t} p(y_{i} \\ | \\ y_{<i}, x) \\neq \\prod\\limits_{i=1}^{t}\\underset{y_i}{\\operatorname{arg max}}p(y_i \\ | \\ y_{<i}, x)$\n",
    "\n",
    "- **Beam Search (k)**: At each iteration, **extend only the top $k$ likely sequences by $k$ most likely extensions**. $k$ is also termed at beam size. At each step, we will have $k^2$ sequences and we will keep only the top-$k$ sequences. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def greedy_decoding(model, src, max_len):\n",
    "    \"\"\"\n",
    "    Implements greedy search for a most likely translation of `src`. It accepts a batch of source sentences. \n",
    "    \n",
    "    Args:\n",
    "        model (Transformer): model to extract p(y_t | y_{<t}, x)\n",
    "        src (torch.tensor): [n_sentences x max_sentence_length] source sentence with paddings, BOS and EOS tokens (as expected by the model)\n",
    "        max_len (int): maximum length of the translated sentence.\n",
    "    \n",
    "    Returns:\n",
    "        (torch.tensor): indices of tokens in a target sentence\n",
    "    \"\"\"\n",
    "    model.to(device)\n",
    "    model.eval() # no use of dropout from here on\n",
    "    \n",
    "    src = src.to(device)\n",
    "    N  = src.shape[0]  # n_sentences\n",
    "    eos_tracker = set()\n",
    "    with torch.no_grad():\n",
    "        \n",
    "        # src padding\n",
    "        src_padding_mask = (src != PAD_IDX).unsqueeze(1).unsqueeze(2)\n",
    "        src_padding_mask = src_padding_mask.to(device)\n",
    "        \n",
    "        # source encoding (x)\n",
    "        e_x = # compute source encodings\n",
    "        \n",
    "        # output decoding (y_i)\n",
    "        outputs = torch.empty((N, max_len + 1), dtype = torch.long)\n",
    "        outputs[:] = PAD_IDX\n",
    "        outputs = outputs.to(device)\n",
    "        outputs[:, 0] = # define a starting token to begin decoding\n",
    "        \n",
    "        for i in range(1, max_len+1):\n",
    "            tgt_mask = # define an appropriate mask to hide all tokens ahead of the token at i\n",
    "            tgt_mask = tgt_mask.to(device)\n",
    "            \n",
    "            out = model.out(model.decode(outputs[:, :i], e_x, src_padding_mask, tgt_mask))\n",
    "            ix = # last position's predictions only; use out.data.topk with appropriate indexing on out\n",
    "            \n",
    "            rem_sentences = [i not in eos_tracker for i in range(N)]\n",
    "            outputs[rem_sentences, i] = ix[rem_sentences, 0]\n",
    "            \n",
    "            eos_tracker.update(set(torch.where(ix[:, 0] == EOS_IDX)[0].tolist())) # these sentences have finished\n",
    "            \n",
    "            if len(eos_tracker) == N:\n",
    "                break\n",
    "    \n",
    "    return outputs.to('cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# token to text conversion\n",
    "\n",
    "def get_readable_tokens(tokens, language):\n",
    "    assert len(tokens.shape) == 1, f\"unrecognized input shape: {tokens.shape}\"\n",
    "    return vocab_transform[language].lookup_tokens(tokens.numpy())\n",
    "    \n",
    "def get_readable_text(readable_tokens):\n",
    "    return \" \".join(readable_tokens).replace(\"<bos>\", \"\").replace(\"<eos>\", \"\").replace(\"<pad>\", \"\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# check the predictions\n",
    "src, tgt = get_sample_batch('valid', 10)\n",
    "pred_sentences = greedy_decoding(model, src, 40)\n",
    "\n",
    "for i in range(src.shape[0]):\n",
    "    s = get_readable_tokens(src[i, :], SRC_LANGUAGE)\n",
    "    y = get_readable_tokens(tgt[i, :], TGT_LANGUAGE)\n",
    "    y_hat = get_readable_tokens(pred_sentences[i, :], TGT_LANGUAGE)\n",
    "    meteor = nltk.translate.meteor_score.meteor_score([get_readable_text(y)], get_readable_text(y_hat))\n",
    "    print(get_readable_text(s))\n",
    "    print(get_readable_text(y))\n",
    "    print(get_readable_text(y_hat))\n",
    "    print(\"METEOR:\", meteor)\n",
    "    print(\"-\"*50)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### METEOR: Translation performance\n",
    "\n",
    "Finally, we will compute the performance on translation task using METEOR metric.\n",
    "Checkout how to use this [score here](`nltk.translate.meteor_score`](https://www.nltk.org/_modules/nltk/translate/meteor_score.html)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_meteor(model, data_loader, strategy):\n",
    "    running_meteor = 0\n",
    "    n_samples = 0\n",
    "    for src, tgt in data_loader:\n",
    "        \n",
    "        n_samples += src.shape[0] # number of sentences\n",
    "        \n",
    "        # inference\n",
    "        preds = strategy(model, src, src.shape[1] + 50)\n",
    "        \n",
    "        for i in range(src.shape[0]):\n",
    "            y = # get readable tokens\n",
    "            y_hat = # get readable tokens\n",
    "            meteor = # compute scores using nltk.translate.meteor_score\n",
    "            running_meteor += meteor\n",
    "    \n",
    "    return running_meteor / n_samples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Validation Performance\n",
    "\n",
    "What's the average meteor score on validation dataset?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "val_dataloader = # instantiate validation dataloader here\n",
    "val_meteor = compute_meteor(model, val_dataloader, greedy_decoding)\n",
    "print(\"Validation performance:\")\n",
    "print(f\"meteor: {val_meteor: 0.5f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test Performance\n",
    "\n",
    "Finally, we want to evaluate model's performance on test and validation dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_dataloader = # instantiate test dataloader here\n",
    "test_loss = process(model, test_dataloader, loss_fn)\n",
    "test_meteor = compute_meteor(model, test_dataloader, greedy_decoding)\n",
    "\n",
    "print(\"Test performance: \\n\")\n",
    "print(f\"loss: {test_loss:0.5f}\")\n",
    "print(f\"meteor: {test_meteor:0.5f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### (Advanced) Inference-II: Beam Search\n",
    "\n",
    "Algorithm is described [here](https://d2l.ai/chapter_recurrent-modern/beam-search.html).\n",
    "Below is an inefficient method for beam search as it performs beam search one sentence at a time.\n",
    "For a more efficient procedure, refer [OpennMT-Py's implementation](https://opennmt.net/OpenNMT-py/options/translate.html?highlight=beam%20search#Beam%20Search)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _beam_search(model, src, max_len, k=5):\n",
    "    \"\"\"\n",
    "    Implements beam search for a most likely translation of `src`.\n",
    "    \n",
    "    Args:\n",
    "        model (Transformer): model to extract p(y_t | y_{<t}, x)\n",
    "        src (torch.tensor): [n_sentences x max_sentence_length] source sentence with paddings, BOS and EOS tokens (as expected by the model)\n",
    "        max_len (int): maximum length of the translated sentence.\n",
    "        k (int): beam size or number of most likely sequences to consider\n",
    "    \n",
    "    Returns:\n",
    "        (torch.tensor): indices of tokens in a target sentence\n",
    "    \"\"\"\n",
    "    def _compute_probability(tgts, e_x, src_padding_mask):\n",
    "        \"\"\"\n",
    "        Computes prediction for next token\n",
    "        \"\"\"\n",
    "        tgt_mask = torch.tril(torch.ones(tgts.shape[0], 1, tgts.shape[1], tgts.shape[1]), diagonal=0)\n",
    "        tgt_mask = tgt_mask.to(device)\n",
    "        \n",
    "        in_e_x = e_x.repeat((tgts.shape[0], 1, 1))\n",
    "        out = model.out(model.decode(tgts, in_e_x, src_padding_mask, tgt_mask))\n",
    "        logprobs = torch.nn.functional.log_softmax(out, dim=-1)\n",
    "        values, ix = logprobs[:, -1, :].data.topk(k=k)\n",
    "        \n",
    "        return values, ix\n",
    "\n",
    "    \n",
    "    model.to(device)\n",
    "    model.eval() # no use of dropout from here on\n",
    "    \n",
    "    src = src.to(device)\n",
    "    N  = src.shape[0]  # n_sentences\n",
    "    predictions = [None for _ in range(N)]\n",
    "    for j in range(N):\n",
    "        finished_candidates = []\n",
    "        in_src = src[j, :].unsqueeze(0)\n",
    "        \n",
    "        # src padding\n",
    "        src_padding_mask = (in_src != PAD_IDX).unsqueeze(1).unsqueeze(2)\n",
    "        src_padding_mask = src_padding_mask.to(device)\n",
    "\n",
    "        # source encoding (x)\n",
    "        e_x = model.encode(in_src, src_padding_mask) # source encodings\n",
    "\n",
    "        # decoding (y_i)\n",
    "        tgts = torch.tensor([[BOS_IDX]], dtype = torch.long)\n",
    "        tgts = tgts.to(device)\n",
    "        \n",
    "        ## generate first k-hypothesis for the first token \n",
    "        values, ix = _compute_probability(tgts, e_x, src_padding_mask)\n",
    "        \n",
    "        ## kick start with k first tokens\n",
    "        candidates = [\n",
    "            [[BOS_IDX, token.item()], score.item()] for token, score in zip(ix[0, :], values[0, :])\n",
    "        ]\n",
    "        \n",
    "        # beam search starting from 2nd token\n",
    "        for i in range(2, max_len+1):\n",
    "            # condition on the k-th hypothesis \n",
    "            tgts = torch.tensor([x[0] for x in candidates], dtype=torch.long)\n",
    "            tgts = tgts.to(device)\n",
    "            \n",
    "            # compute k next likely token\n",
    "            values, ix = _compute_probability(tgts, e_x, src_padding_mask)\n",
    "            \n",
    "            # extend each hypothesis with their k most likely extensions\n",
    "            all_candidates = []\n",
    "            for c, cand in enumerate(candidates):\n",
    "                for idx, score in zip(ix[c, :], values[c, :]):\n",
    "                    all_candidates.append([cand[0] + [idx.item()], cand[1] + score.item()])\n",
    "            \n",
    "            # keep top-k candidates\n",
    "            _candidates = sorted(all_candidates, key=lambda x:-x[1])[:k]\n",
    "            \n",
    "            # if the candidate is finished, store it\n",
    "            finished_candidates += [x for x in _candidates if x[0][-1] == EOS_IDX]\n",
    "            \n",
    "            # continue with rest of the candidates\n",
    "            candidates = [x for x in _candidates if x[0][-1] != EOS_IDX]\n",
    "            \n",
    "            if len(candidates) == 0:\n",
    "                break\n",
    "        \n",
    "        # select the best one \n",
    "        # refer to equation 9.8.4 of https://d2l.ai/chapter_recurrent-modern/beam-search.html\n",
    "        finished_candidates += candidates\n",
    "        predictions[j] = max(finished_candidates, key=lambda x: x[1]/pow(len(x[0]), 0.75))[0] \n",
    "    \n",
    "    return pad_sequence([torch.tensor(x) for x in predictions], padding_value=PAD_IDX).transpose(1,0).cpu()\n",
    "\n",
    "\n",
    "def beam_search_strategy(k=5, *args):\n",
    "    def func(*args):\n",
    "        return _beam_search(*args, k)\n",
    "    return func"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# check the predictions\n",
    "src, tgt = get_sample_batch('valid', 10)\n",
    "pred_sentences = _beam_search(model, src, 30, 5)\n",
    "\n",
    "for i in range(src.shape[0]):\n",
    "    s = get_readable_tokens(src[i, :], SRC_LANGUAGE)\n",
    "    y = get_readable_tokens(tgt[i, :], TGT_LANGUAGE)\n",
    "    y_hat = get_readable_tokens(pred_sentences[i, :], TGT_LANGUAGE)\n",
    "    meteor = nltk.translate.meteor_score.meteor_score([get_readable_text(y)], get_readable_text(y_hat))\n",
    "    print(get_readable_text(s))\n",
    "    print(get_readable_text(y))\n",
    "    print(get_readable_text(y_hat))\n",
    "    print(\"METEOR:\", meteor)\n",
    "    print(\"-\"*50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# CAUTION: It takes a LOT OF time\n",
    "\n",
    "for k in range(1, 4):\n",
    "    # define strategy\n",
    "    strategy = beam_search_strategy(k)\n",
    "    \n",
    "    # val\n",
    "    val_dataloader = torch.utils.data.DataLoader(list(get_data_iter('valid')), batch_size=BATCH_SIZE, collate_fn=collate_fn)\n",
    "    val_meteor = compute_meteor(model, val_dataloader, strategy)\n",
    "   \n",
    "\n",
    "    # test\n",
    "    test_dataloader = torch.utils.data.DataLoader(list(get_data_iter('test')), batch_size=BATCH_SIZE, collate_fn=collate_fn)\n",
    "    test_meteor = compute_meteor(model, test_dataloader, strategy)\n",
    "    \n",
    "    print(f\"k = {k} \\t val meteor: {val_meteor: 0.5f} \\t test_meteor: {test_meteor: 0.5f}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Greedy vs Beam Search\n",
    "Is beam search always better than greedy search? Maybe not! [Read here](https://discuss.huggingface.co/t/is-beam-search-always-better-than-greedy-search/2943/3)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dlc",
   "language": "python",
   "name": "dlc"
  },
  "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}