The Annotated Transformer ++.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The Annotated Transformer ++\n",
    "\n",
    "The idea of this notebook is to make it easier even for non-researchers to understand the omnipresent transformer model!\n",
    "\n",
    "In this notebook you'll learn:\n",
    "\n",
    "✅ What are transformers exactly? <br/>\n",
    "✅ How to train them? <br/>\n",
    "✅ How to use them? (machine translation example) <br/>\n",
    "\n",
    "After you complete this one you'll have a much better understanding of transformers!\n",
    "\n",
    "So, let's start!\n",
    "\n",
    "---\n",
    "\n",
    "## What the heck are transformers and how they came to be?\n",
    "\n",
    "Transformers were originally proposed by Vaswani et al. in a seminal paper called [Attention Is All You Need](https://arxiv.org/pdf/1706.03762.pdf).\n",
    "\n",
    "You probably heard of transformers one way or another. GPT-3 and BERT to name a few well known ones 🦄. The main idea is that they showed that you don't have to use recurrent or convolutional layers and that simple architecture coupled with attention is super powerful. It gave the benefit of much better long-range dependency modeling and the architecture itself is highly parallelizable (💻💻💻) which leads to better compute efficiency!\n",
    "\n",
    "Here is how their beautifully simple architecture looks like:\n",
    "\n",
    "<img src=\"data/readme_pics/transformer_architecture.PNG\" alt=\"transformer architecture\" align=\"center\" style=\"width: 350px;\"/> <br/>\n",
    "\n",
    "---\n",
    "\n",
    "That was everything you need to know for now! If you need further help understanding all of the details I created this [in-depth overview of the paper:](https://www.youtube.com/watch?v=cbYxHkgkSVs)\n",
    "\n",
    "<a href=\"https://www.youtube.com/watch?v=cbYxHkgkSVs\" target=\"_blank\"><img src=\"https://img.youtube.com/vi/cbYxHkgkSVs/0.jpg\" \n",
    "alt=\"An in-depth overview of the Attention Is All You Need paper\" width=\"480\" align=\"left\" height=\"360\" border=\"10\" /></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# I always like to structure my imports into Python's native libs,\n",
    "# stuff I installed via conda/pip and local file imports (we don't have those here)\n",
    "\n",
    "import math\n",
    "import copy\n",
    "import os\n",
    "import time\n",
    "import enum\n",
    "import argparse\n",
    "\n",
    "# Visualization related imports\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn\n",
    "\n",
    "# Deep learning related imports\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.optim import Adam\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "from torch.hub import download_url_to_file\n",
    "\n",
    "# Data manipulation related imports\n",
    "from torchtext.data import Dataset, BucketIterator, Field, Example\n",
    "from torchtext.data.utils import interleave_keys\n",
    "from torchtext import datasets\n",
    "from torchtext.data import Example\n",
    "import spacy\n",
    "\n",
    "# BLEU (the metric we'll be using for the machine translation task)\n",
    "from nltk.translate.bleu_score import corpus_bleu"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's create some constants to make the code a bit cleaner\n",
    "\n",
    "# Architecture related constants taken from the paper\n",
    "BASELINE_MODEL_NUMBER_OF_LAYERS = 6\n",
    "BASELINE_MODEL_DIMENSION = 512\n",
    "BASELINE_MODEL_NUMBER_OF_HEADS = 8\n",
    "BASELINE_MODEL_DROPOUT_PROB = 0.1\n",
    "BASELINE_MODEL_LABEL_SMOOTHING_VALUE = 0.1\n",
    "\n",
    "CHECKPOINTS_PATH = os.path.join(os.getcwd(), 'models', 'checkpoints') # semi-trained models during training will be dumped here\n",
    "BINARIES_PATH = os.path.join(os.getcwd(), 'models', 'binaries') # location where trained models are located\n",
    "DATA_DIR_PATH = os.path.join(os.getcwd(), 'data') # training data will be stored here\n",
    "\n",
    "os.makedirs(CHECKPOINTS_PATH, exist_ok=True)\n",
    "os.makedirs(BINARIES_PATH, exist_ok=True)\n",
    "os.makedirs(DATA_DIR_PATH, exist_ok=True)\n",
    "\n",
    "# Special token symbols used later in the data section\n",
    "BOS_TOKEN = '<s>'\n",
    "EOS_TOKEN = '</s>'\n",
    "PAD_TOKEN = \"<pad>\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Part 1: Understanding the model (transformer)\n",
    "\n",
    "There are 2 really important parts of any deep learning pipeline:\n",
    "1. Model 🦄\n",
    "2. Data 📜\n",
    "\n",
    "Let's start with the model (🦄) first and understand how the transformer is architectured.\n",
    "\n",
    "I'll take a top-down approach here.\n",
    "\n",
    "We should first understand how all of the modules fit together (**level 0**) and then we'll start digging deeper into the nitty-gritty details!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Transformer(nn.Module):\n",
    "\n",
    "    def __init__(self, model_dimension, src_vocab_size, trg_vocab_size, number_of_heads, number_of_layers, dropout_probability, log_attention_weights=False):\n",
    "        super().__init__()\n",
    "\n",
    "        # Embeds source/target token ids into embedding vectors\n",
    "        self.src_embedding = Embedding(src_vocab_size, model_dimension)\n",
    "        self.trg_embedding = Embedding(trg_vocab_size, model_dimension)\n",
    "\n",
    "        # Adds positional information to source/target token's embedding vector\n",
    "        # (otherwise we'd lose the positional information which is important in human languages)\n",
    "        self.src_pos_embedding = PositionalEncoding(model_dimension, dropout_probability)\n",
    "        self.trg_pos_embedding = PositionalEncoding(model_dimension, dropout_probability)\n",
    "\n",
    "        # All of these will get deep-copied multiple times internally\n",
    "        mha = MultiHeadedAttention(model_dimension, number_of_heads, dropout_probability, log_attention_weights)\n",
    "        pwn = PositionwiseFeedForwardNet(model_dimension, dropout_probability)\n",
    "        encoder_layer = EncoderLayer(model_dimension, dropout_probability, mha, pwn)\n",
    "        decoder_layer = DecoderLayer(model_dimension, dropout_probability, mha, pwn)\n",
    "\n",
    "        # Encoder and Decoder stacks\n",
    "        self.encoder = Encoder(encoder_layer, number_of_layers)\n",
    "        self.decoder = Decoder(decoder_layer, number_of_layers)\n",
    "\n",
    "        # Converts final target token representations into log probability vectors of trg_vocab_size dimensionality\n",
    "        # Why log? -> PyTorch's nn.KLDivLoss expects log probabilities\n",
    "        self.decoder_generator = DecoderGenerator(model_dimension, trg_vocab_size)\n",
    "        \n",
    "        self.init_params()\n",
    "\n",
    "    # This part wasn't mentioned in the paper, but it's super important!\n",
    "    def init_params(self):\n",
    "        # I tested both PyTorch's default initialization and this, and xavier has tremendous impact! I didn't expect\n",
    "        # that the model's perf, with normalization layers, is so dependent on the choice of weight initialization.\n",
    "        for name, p in self.named_parameters():\n",
    "            if p.dim() > 1:\n",
    "                nn.init.xavier_uniform_(p)\n",
    "\n",
    "    def forward(self, src_token_ids_batch, trg_token_ids_batch, src_mask, trg_mask):\n",
    "        src_representations_batch = self.encode(src_token_ids_batch, src_mask)\n",
    "        trg_log_probs = self.decode(trg_token_ids_batch, src_representations_batch, trg_mask, src_mask)\n",
    "        \n",
    "        return trg_log_probs\n",
    "\n",
    "    # Modularize into encode/decode functions for optimizing the decoding/translation process (we'll get to it later)\n",
    "    def encode(self, src_token_ids_batch, src_mask):\n",
    "        # Shape = (B, S, D) , where B - batch size, S - longest source token-sequence length and D - model dimension\n",
    "        # The whole encoder stack perserves this shape\n",
    "        src_embeddings_batch = self.src_embedding(src_token_ids_batch)  # get embedding vectors for src token ids\n",
    "        src_embeddings_batch = self.src_pos_embedding(src_embeddings_batch)  # add positional embedding\n",
    "        src_representations_batch = self.encoder(src_embeddings_batch, src_mask)  # forward pass through the encoder\n",
    "\n",
    "        return src_representations_batch\n",
    "\n",
    "    def decode(self, trg_token_ids_batch, src_representations_batch, trg_mask, src_mask):\n",
    "        trg_embeddings_batch = self.trg_embedding(trg_token_ids_batch)  # get embedding vectors for trg token ids\n",
    "        trg_embeddings_batch = self.trg_pos_embedding(trg_embeddings_batch)  # add positional embedding\n",
    "        \n",
    "        # Shape (B, T, D), where B - batch size, T - longest target token-sequence length and D - model dimension\n",
    "        trg_representations_batch = self.decoder(trg_embeddings_batch, src_representations_batch, trg_mask, src_mask)\n",
    "\n",
    "        # After this line we'll have a shape (B, T, V), where V - target vocab size, \n",
    "        # decoder generator does a simple linear projection followed by log softmax\n",
    "        trg_log_probs = self.decoder_generator(trg_representations_batch)\n",
    "\n",
    "        # Reshape into (B*T, V) as that's a suitable format for passing it into KL div loss\n",
    "        trg_log_probs = trg_log_probs.reshape(-1, trg_log_probs.shape[-1])\n",
    "\n",
    "        return trg_log_probs  # the reason I use log here is that PyTorch's nn.KLDivLoss expects log probabilities"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Encoder and Decoder\n",
    "\n",
    "Stay with me here! \n",
    "\n",
    "These initial parts will probably be the most confusing as you're not yet familiar with terminology and I'll reference many things yet to be defined.\n",
    "\n",
    "Get used to it. Whether you're reading research papers or code things can't always get linearized. <br/>\n",
    "Think of it as a graph which can't get converted into a [DAG](https://en.wikipedia.org/wiki/Directed_acyclic_graph) (direct acyclic graph) because we have circular dependencies!\n",
    "\n",
    "Let's go **1 level** deeper - and see exactly how encoder and decoder are structured.\n",
    "\n",
    "<img src=\"data/readme_pics/deeper_inception.jpg\" alt=\"we need to go deeper\" align=\"center\" style=\"width: 350px;\"/> <br/>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "High level overview of the functionality:\n",
    "\n",
    "1. The encoder takes as the input a batch of source sequences whose tokens have already been embedded.\n",
    "2. It then does 6 iterations (6 layers for the base transformer) of mixing those via attention.\n",
    "3. The final output will later get consumed by the decoder. That's it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Encoder(nn.Module):\n",
    "\n",
    "    def __init__(self, encoder_layer, number_of_layers):\n",
    "        super().__init__()\n",
    "        assert isinstance(encoder_layer, EncoderLayer), f'Expected EncoderLayer got {type(encoder_layer)}.'\n",
    "\n",
    "        # Get a list of 'number_of_layers' independent encoder layers\n",
    "        self.encoder_layers = get_clones(encoder_layer, number_of_layers)  \n",
    "        self.norm = nn.LayerNorm(encoder_layer.model_dimension)\n",
    "\n",
    "    def forward(self, src_embeddings_batch, src_mask):\n",
    "        # Just update the naming so as to reflect the semantics of what this var will become (the initial encoder layer\n",
    "        # has embedding vectors as input but later layers have richer token representations)\n",
    "        src_representations_batch = src_embeddings_batch\n",
    "\n",
    "        # Forward pass through the encoder stack\n",
    "        for encoder_layer in self.encoder_layers:\n",
    "            # src_mask's role is to mask/ignore padded token representations in the multi-headed self-attention module\n",
    "            src_representations_batch = encoder_layer(src_representations_batch, src_mask)\n",
    "\n",
    "        # Not mentioned explicitly in the paper \n",
    "        # (a consequence of using LayerNorm before instead of after the SublayerLogic module)\n",
    "        return self.norm(src_representations_batch)\n",
    "\n",
    "\n",
    "class EncoderLayer(nn.Module):\n",
    "\n",
    "    def __init__(self, model_dimension, dropout_probability, multi_headed_attention, pointwise_net):\n",
    "        super().__init__()\n",
    "        num_of_sublayers_encoder = 2\n",
    "        self.sublayers = get_clones(SublayerLogic(model_dimension, dropout_probability), num_of_sublayers_encoder)\n",
    "\n",
    "        self.multi_headed_attention = multi_headed_attention\n",
    "        self.pointwise_net = pointwise_net\n",
    "\n",
    "        self.model_dimension = model_dimension\n",
    "\n",
    "    def forward(self, src_representations_batch, src_mask):\n",
    "        # Define an anonymous (lambda) function which only takes src_representations_batch (srb) as input,\n",
    "        # this way we have a uniform interface for the sublayer logic.\n",
    "        encoder_self_attention = lambda srb: self.multi_headed_attention(query=srb, key=srb, value=srb, mask=src_mask)\n",
    "\n",
    "        # Self-attention MHA sublayer followed by point-wise feed forward net sublayer\n",
    "        # SublayerLogic takes as the input the data and the logic it should execute (attention/feedforward)\n",
    "        src_representations_batch = self.sublayers[0](src_representations_batch, encoder_self_attention)\n",
    "        src_representations_batch = self.sublayers[1](src_representations_batch, self.pointwise_net)\n",
    "\n",
    "        return src_representations_batch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The decoder has a similar structure to encoder.\n",
    "\n",
    "1. It again starts with a batch of target sequences whose tokens have already been embedded.\n",
    "2. It then does 6 iterations (6 layers for the base transformer) of mixing those via attention (this time also attending to source token representations)\n",
    "3. The final output i.e. target token representations go into the DecoderGenerator module which will convert them\n",
    "into log probabilities\n",
    "\n",
    "A really important difference is that the decoder uses **causal masking** so as to prevent tokens from looking into the future.\n",
    "\n",
    "Here is an example of a causal mask used in decoder's self attention module:\n",
    "<img src=\"data/readme_pics/causal_mask.PNG\" alt=\"casual mask\" align=\"center\"/> <br/>\n",
    "\n",
    "What this mask tells is that tokens on the left side of the chart can only attend to tokens (columns) whose values are set to 1.0 (True).\n",
    "\n",
    "As you can see the token can also attend to itself and to tokens that come before - hence **causal**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Decoder(nn.Module):\n",
    "\n",
    "    def __init__(self, decoder_layer, number_of_layers):\n",
    "        super().__init__()\n",
    "        assert isinstance(decoder_layer, DecoderLayer), f'Expected DecoderLayer got {type(decoder_layer)}.'\n",
    "\n",
    "        self.decoder_layers = get_clones(decoder_layer, number_of_layers)\n",
    "        self.norm = nn.LayerNorm(decoder_layer.model_dimension)\n",
    "\n",
    "    def forward(self, trg_embeddings_batch, src_representations_batch, trg_mask, src_mask):\n",
    "        # Just update the naming so as to reflect the semantics of what this var will become\n",
    "        trg_representations_batch = trg_embeddings_batch\n",
    "\n",
    "        # Forward pass through the decoder stack\n",
    "        for decoder_layer in self.decoder_layers:\n",
    "            # Target mask masks pad tokens as well as future tokens (current target token can't look forward)\n",
    "            trg_representations_batch = decoder_layer(trg_representations_batch, src_representations_batch, trg_mask, src_mask)\n",
    "\n",
    "        # Not mentioned explicitly in the paper \n",
    "        # (a consequence of using LayerNorm before instead of after the SublayerLogic module)\n",
    "        return self.norm(trg_representations_batch)\n",
    "\n",
    "\n",
    "class DecoderLayer(nn.Module):\n",
    "\n",
    "    def __init__(self, model_dimension, dropout_probability, multi_headed_attention, pointwise_net):\n",
    "        super().__init__()\n",
    "        num_of_sublayers_decoder = 3\n",
    "        self.sublayers = get_clones(SublayerLogic(model_dimension, dropout_probability), num_of_sublayers_decoder)\n",
    "\n",
    "        self.trg_multi_headed_attention = copy.deepcopy(multi_headed_attention)\n",
    "        self.src_multi_headed_attention = copy.deepcopy(multi_headed_attention)\n",
    "        self.pointwise_net = pointwise_net\n",
    "\n",
    "        self.model_dimension = model_dimension\n",
    "\n",
    "    def forward(self, trg_representations_batch, src_representations_batch, trg_mask, src_mask):\n",
    "        # Define an anonymous (lambda) function which only takes trg_representations_batch (trb - funny name I know)\n",
    "        # as input - this way we have a uniform interface for the sublayer logic.\n",
    "        # The inputs which are not passed into lambdas (masks/srb) are \"cached\" here that's why the thing works.\n",
    "        srb = src_representations_batch  # simple/short alias\n",
    "        decoder_trg_self_attention = lambda trb: self.trg_multi_headed_attention(query=trb, key=trb, value=trb, mask=trg_mask)\n",
    "        decoder_src_attention = lambda trb: self.src_multi_headed_attention(query=trb, key=srb, value=srb, mask=src_mask)\n",
    "\n",
    "        # Self-attention MHA sublayer followed by a source-attending MHA and point-wise feed forward net sublayer\n",
    "        trg_representations_batch = self.sublayers[0](trg_representations_batch, decoder_trg_self_attention)\n",
    "        trg_representations_batch = self.sublayers[1](trg_representations_batch, decoder_src_attention)\n",
    "        trg_representations_batch = self.sublayers[2](trg_representations_batch, self.pointwise_net)\n",
    "\n",
    "        return trg_representations_batch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And the final **level 1** module is the DecoderGenerator (level 0 was the topmost Transformer class).\n",
    "\n",
    "This module does 2 things:\n",
    "1. Projects the final decoder token representations of dimension D into V dimensions (target vocabulary size)\n",
    "2. Applies a log softmax ([PyTorch's KL div loss](https://pytorch.org/docs/stable/generated/torch.nn.KLDivLoss.html) expects log probabilities again)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DecoderGenerator(nn.Module):\n",
    "    def __init__(self, model_dimension, vocab_size):\n",
    "        super().__init__()\n",
    "\n",
    "        self.linear = nn.Linear(model_dimension, vocab_size)\n",
    "\n",
    "        # -1 stands for apply the log-softmax along the last dimension i.e. over the vocab dimension as the output from\n",
    "        # the linear layer has shape (B, T, V), B - batch size, T - max target token-sequence, V - target vocab size\n",
    "        \n",
    "        # again using log softmax as PyTorch's nn.KLDivLoss expects log probabilities (just a technical detail)\n",
    "        self.log_softmax = nn.LogSoftmax(dim=-1)\n",
    "\n",
    "    def forward(self, trg_representations_batch):\n",
    "        # Project from D (model dimension) into V (target vocab size) and apply the log softmax along V dimension\n",
    "        return self.log_softmax(self.linear(trg_representations_batch))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Going even deeper\n",
    "\n",
    "Phew, give yourself a pat on the back and take a deep breath as we're diving into final details related to transformer's architecture.\n",
    "\n",
    "Here is how the positional encodings look like:\n",
    "\n",
    "<img src=\"data/readme_pics/positional_encoding_visualized.jpg\"/>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SublayerLogic(nn.Module):\n",
    "    def __init__(self, model_dimension, dropout_probability):\n",
    "        super().__init__()\n",
    "        self.norm = nn.LayerNorm(model_dimension)\n",
    "        self.dropout = nn.Dropout(p=dropout_probability)\n",
    "\n",
    "    # Note: the original paper had LayerNorm AFTER the residual connection and addition operation\n",
    "    # multiple experiments I found showed that it's more effective to do it BEFORE, how did they figure out which one is\n",
    "    # better? Experiments! There is a similar thing in DCGAN and elsewhere.\n",
    "    def forward(self, representations_batch, sublayer_module):\n",
    "        # Residual connection between input and sublayer output, details: Page 7, Chapter 5.4 \"Regularization\",\n",
    "        return representations_batch + self.dropout(sublayer_module(self.norm(representations_batch)))\n",
    "\n",
    "\n",
    "class PositionwiseFeedForwardNet(nn.Module):\n",
    "    \"\"\"\n",
    "        It's position-wise because this feed forward net will be independently applied to every token's representation.\n",
    "\n",
    "        This net will basically be applied independently to every token's representation (you can think of it as if\n",
    "        there was a nested for-loop going over the batch size and max token sequence length dimensions\n",
    "        and applied this network to token representations. PyTorch does this auto-magically behind the scenes.\n",
    "\n",
    "    \"\"\"\n",
    "    def __init__(self, model_dimension, dropout_probability, width_mult=4):\n",
    "        super().__init__()\n",
    "\n",
    "        self.linear1 = nn.Linear(model_dimension, width_mult * model_dimension)\n",
    "        self.linear2 = nn.Linear(width_mult * model_dimension, model_dimension)\n",
    "\n",
    "        # This dropout layer is not explicitly mentioned in the paper but it's common to use to avoid over-fitting\n",
    "        self.dropout = nn.Dropout(p=dropout_probability)\n",
    "        self.relu = nn.ReLU()\n",
    "\n",
    "    # Representations_batch's shape = (B - batch size, S/T - max token sequence length, D- model dimension).\n",
    "    def forward(self, representations_batch):\n",
    "        return self.linear2(self.dropout(self.relu(self.linear1(representations_batch))))\n",
    "\n",
    "\n",
    "#\n",
    "# Input modules\n",
    "#\n",
    "\n",
    "\n",
    "class Embedding(nn.Module):\n",
    "\n",
    "    def __init__(self, vocab_size, model_dimension):\n",
    "        super().__init__()\n",
    "        self.embeddings_table = nn.Embedding(vocab_size, model_dimension)\n",
    "        self.model_dimension = model_dimension\n",
    "\n",
    "    def forward(self, token_ids_batch):\n",
    "        assert token_ids_batch.ndim == 2, f'Expected: (batch size, max token sequence length), got {token_ids_batch.shape}'\n",
    "\n",
    "        # token_ids_batch has shape (B, S/T), where B - batch size, S/T max src/trg token-sequence length\n",
    "        # Final shape will be (B, S/T, D) where D is the model dimension, every token id has associated vector\n",
    "        embeddings = self.embeddings_table(token_ids_batch)\n",
    "\n",
    "        # (stated in the paper) multiply the embedding weights by the square root of model dimension\n",
    "        # Page 5, Chapter 3.4 \"Embeddings and Softmax\"\n",
    "        return embeddings * math.sqrt(self.model_dimension)\n",
    "\n",
    "\n",
    "class PositionalEncoding(nn.Module):\n",
    "\n",
    "    def __init__(self, model_dimension, dropout_probability, expected_max_sequence_length=5000):\n",
    "        super().__init__()\n",
    "        self.dropout = nn.Dropout(p=dropout_probability)\n",
    "\n",
    "        # (stated in the paper) Use sine functions whose frequencies form a geometric progression as position encodings,\n",
    "        # (learning encodings will also work so feel free to change it!). Page 6, Chapter 3.5 \"Positional Encoding\"\n",
    "        position_id = torch.arange(0, expected_max_sequence_length).unsqueeze(1)\n",
    "        frequencies = torch.pow(10000., -torch.arange(0, model_dimension, 2, dtype=torch.float) / model_dimension)\n",
    "\n",
    "        # Checkout playground.py for visualization of how these look like (it's super simple don't get scared)\n",
    "        positional_encodings_table = torch.zeros(expected_max_sequence_length, model_dimension)\n",
    "        positional_encodings_table[:, 0::2] = torch.sin(position_id * frequencies)  # sine on even positions\n",
    "        positional_encodings_table[:, 1::2] = torch.cos(position_id * frequencies)  # cosine on odd positions\n",
    "\n",
    "        # Register buffer because we want to save the positional encodings table inside state_dict even though\n",
    "        # these are not trainable (not model's parameters) so they otherwise would be excluded from the state_dict\n",
    "        self.register_buffer('positional_encodings_table', positional_encodings_table)\n",
    "\n",
    "    def forward(self, embeddings_batch):\n",
    "        assert embeddings_batch.ndim == 3 and embeddings_batch.shape[-1] == self.positional_encodings_table.shape[1], \\\n",
    "            f'Expected (batch size, max token sequence length, model dimension) got {embeddings_batch.shape}'\n",
    "\n",
    "        # embedding_batch's shape = (B, S/T, D), where S/T max src/trg token-sequence length, D - model dimension\n",
    "        # So here we get (S/T, D) shape which will get broad-casted to (B, S/T, D) when we try and add it to embeddings\n",
    "        positional_encodings = self.positional_encodings_table[:embeddings_batch.shape[1]]\n",
    "\n",
    "        # (stated in the paper) Applying dropout to the sum of positional encodings and token embeddings\n",
    "        # Page 7, Chapter 5.4 \"Regularization\"\n",
    "        return self.dropout(embeddings_batch + positional_encodings)\n",
    "\n",
    "\n",
    "#\n",
    "# Helper model functions\n",
    "#\n",
    "\n",
    "\n",
    "def get_clones(module, num_of_deep_copies):\n",
    "    # Create deep copies so that we can tweak each module's weights independently\n",
    "    return nn.ModuleList([copy.deepcopy(module) for _ in range(num_of_deep_copies)])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# The attention part of \"Attention Is All You Need\"\n",
    "\n",
    "And last but definitely not least - the multi headed attention module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiHeadedAttention(nn.Module):\n",
    "    \"\"\"\n",
    "        This module already exists in PyTorch. The reason I implemented it here from scratch is that\n",
    "        PyTorch implementation is super complicated as they made it as generic/robust as possible whereas\n",
    "        on the other hand I only want to support a limited use-case.\n",
    "\n",
    "        Also this is arguable the most important architectural component in the Transformer model.\n",
    "\n",
    "        Additional note:\n",
    "        This is conceptually super easy stuff. It's just that matrix implementation makes things a bit less intuitive.\n",
    "        If you take your time and go through the code and figure out all of the dimensions + write stuff down on paper\n",
    "        you'll understand everything. Also do check out this amazing blog for conceptual understanding:\n",
    "\n",
    "        https://jalammar.github.io/illustrated-transformer/\n",
    "\n",
    "        Optimization notes:\n",
    "\n",
    "        qkv_nets could be replaced by Parameter(torch.empty(3 * model_dimension, model_dimension)) and one more matrix\n",
    "        for bias, which would make the implementation a bit more optimized. For the sake of easier understanding though,\n",
    "        I'm doing it like this - using 3 \"feed forward nets\" (without activation/identity hence the quotation marks).\n",
    "        Conceptually both implementations are the same.\n",
    "\n",
    "        PyTorch's query/key/value are of different shape namely (max token sequence length, batch size, model dimension)\n",
    "        whereas I'm using (batch size, max token sequence length, model dimension) because it's easier to understand\n",
    "        and consistent with computer vision apps (batch dimension is always first followed by the number of channels (C)\n",
    "        and image's spatial dimensions height (H) and width (W) -> (B, C, H, W).\n",
    "\n",
    "        This has an important optimization implication, they can reshape their matrix into (B*NH, S/T, HD)\n",
    "        (where B - batch size, S/T - max src/trg sequence length, NH - number of heads, HD - head dimension)\n",
    "        in a single step and I can only get to (B, NH, S/T, HD) in single step\n",
    "        (I could call contiguous() followed by view but that's expensive as it would incur additional matrix copy)\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, model_dimension, number_of_heads, dropout_probability, log_attention_weights):\n",
    "        super().__init__()\n",
    "        assert model_dimension % number_of_heads == 0, f'Model dimension must be divisible by the number of heads.'\n",
    "\n",
    "        self.head_dimension = int(model_dimension / number_of_heads)\n",
    "        self.number_of_heads = number_of_heads\n",
    "\n",
    "        self.qkv_nets = get_clones(nn.Linear(model_dimension, model_dimension), 3)  # identity activation hence \"nets\"\n",
    "        self.out_projection_net = nn.Linear(model_dimension, model_dimension)\n",
    "\n",
    "        self.attention_dropout = nn.Dropout(p=dropout_probability)  # no pun intended, not explicitly mentioned in paper\n",
    "        self.softmax = nn.Softmax(dim=-1)  # -1 stands for apply the softmax along the last dimension\n",
    "\n",
    "        self.log_attention_weights = log_attention_weights  # should we log attention weights\n",
    "        self.attention_weights = None  # for visualization purposes, I cache the weights here (translation_script.py)\n",
    "\n",
    "    def attention(self, query, key, value, mask):\n",
    "        # Step 1: Scaled dot-product attention, Page 4, Chapter 3.2.1 \"Scaled Dot-Product Attention\"\n",
    "        # Notation: B - batch size, S/T max src/trg token-sequence length, NH - number of heads, HD - head dimension\n",
    "        # query/key/value shape = (B, NH, S/T, HD), scores shape = (B, NH, S, S), (B, NH, T, T) or (B, NH, T, S)\n",
    "        # scores have different shapes as MHA is used in 3 contexts, self attention for src/trg and source attending MHA\n",
    "        scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(self.head_dimension)\n",
    "\n",
    "        # Step 2: Optionally mask tokens whose representations we want to ignore by setting a big negative number\n",
    "        # to locations corresponding to those tokens (force softmax to output 0 probability on those locations).\n",
    "        # mask shape = (B, 1, 1, S) or (B, 1, T, T) will get broad-casted (copied) as needed to match scores shape\n",
    "        if mask is not None:\n",
    "            scores.masked_fill_(mask == torch.tensor(False), float(\"-inf\"))\n",
    "\n",
    "        # Step 3: Calculate the attention weights - how much should we attend to surrounding token representations\n",
    "        attention_weights = self.softmax(scores)\n",
    "\n",
    "        # Step 4: Not defined in the original paper apply dropout to attention weights as well\n",
    "        attention_weights = self.attention_dropout(attention_weights)\n",
    "\n",
    "        # Step 5: based on attention weights calculate new token representations\n",
    "        # attention_weights shape = (B, NH, S, S)/(B, NH, T, T) or (B, NH, T, S), value shape = (B, NH, S/T, HD)\n",
    "        # Final shape (B, NH, S, HD) for source MHAs or (B, NH, T, HD) target MHAs (again MHAs are used in 3 contexts)\n",
    "        intermediate_token_representations = torch.matmul(attention_weights, value)\n",
    "\n",
    "        return intermediate_token_representations, attention_weights  # attention weights for visualization purposes\n",
    "\n",
    "    def forward(self, query, key, value, mask):\n",
    "        batch_size = query.shape[0]\n",
    "\n",
    "        # Step 1: Input linear projection\n",
    "        # Notation: B - batch size, NH - number of heads, S/T - max src/trg token-sequence length, HD - head dimension\n",
    "        # Shape goes from (B, S/T, NH*HD) over (B, S/T, NH, HD) to (B, NH, S/T, HD) (NH*HD=D where D is model dimension)\n",
    "        query, key, value = [net(x).view(batch_size, -1, self.number_of_heads, self.head_dimension).transpose(1, 2)\n",
    "                             for net, x in zip(self.qkv_nets, (query, key, value))]\n",
    "\n",
    "        # Step 2: Apply attention - compare query with key and use that to combine values (see the function for details)\n",
    "        intermediate_token_representations, attention_weights = self.attention(query, key, value, mask)\n",
    "\n",
    "        # Potentially, for visualization purposes, log the attention weights, turn off during training though!\n",
    "        # I had memory problems when I leave this on by default\n",
    "        if self.log_attention_weights:\n",
    "            self.attention_weights = attention_weights\n",
    "\n",
    "        # Step 3: Reshape from (B, NH, S/T, HD) over (B, S/T, NH, HD) (via transpose) into (B, S/T, NHxHD) which is\n",
    "        # the same shape as in the beginning of this forward function i.e. input to MHA (multi-head attention) module\n",
    "        reshaped = intermediate_token_representations.transpose(1, 2).reshape(batch_size, -1, self.number_of_heads * self.head_dimension)\n",
    "\n",
    "        # Step 4: Output linear projection\n",
    "        token_representations = self.out_projection_net(reshaped)\n",
    "\n",
    "        return token_representations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's it! You've done it, you now have a better understanding of how the transformer is structured. \n",
    "\n",
    "Level 2 unlocked. 😍 Let's proceed.\n",
    "\n",
    "**Let me stress this - it takes time to understand this project don't get frustrated! ❤️ (it took me ~3 weeks to code this up!)**\n",
    "\n",
    "# Part 2: Understanding the data pipeline\n",
    "\n",
    "The second important part of any deep learning pipeline is the data engine or pipeline. 📜\n",
    "\n",
    "It's the infrastructure responsible for feeding the data to your transformer (or to ML model in general).\n",
    "\n",
    "We'll be using PyTorch's torchtext which has it's problems but it gets the job done so let's continue! <br/>\n",
    "\n",
    "Let's first define functions which will directly feed the data to the transformer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's define a couple of useful enum classes\n",
    "class DatasetType(enum.Enum):\n",
    "    IWSLT = 0,\n",
    "    WMT14 = 1\n",
    "\n",
    "\n",
    "class LanguageDirection(enum.Enum):\n",
    "    E2G = 0,\n",
    "    G2E = 1\n",
    "    \n",
    "\n",
    "def get_datasets_and_vocabs(dataset_path, language_direction, use_iwslt=True, use_caching_mechanism=True):\n",
    "    german_to_english = language_direction == LanguageDirection.G2E.name\n",
    "    spacy_de = spacy.load('de_core_news_sm')\n",
    "    spacy_en = spacy.load('en_core_web_sm')\n",
    "\n",
    "    def tokenize_de(text):\n",
    "        return [tok.text for tok in spacy_de.tokenizer(text)]\n",
    "\n",
    "    def tokenize_en(text):\n",
    "        return [tok.text for tok in spacy_en.tokenizer(text)]\n",
    "\n",
    "    src_tokenizer = tokenize_de if german_to_english else tokenize_en\n",
    "    trg_tokenizer = tokenize_en if german_to_english else tokenize_de\n",
    "    \n",
    "    # These will tokenize the source/target sentences\n",
    "    # batch first set to true as my transformer is expecting that format (that's consistent with the format\n",
    "    # used in  computer vision), namely (B, C, H, W) -> batch size, number of channels, height and width\n",
    "    src_field_processor = Field(tokenize=src_tokenizer, pad_token=PAD_TOKEN, batch_first=True)\n",
    "    trg_field_processor = Field(tokenize=trg_tokenizer, init_token=BOS_TOKEN, eos_token=EOS_TOKEN, pad_token=PAD_TOKEN, batch_first=True)\n",
    "\n",
    "    fields = [('src', src_field_processor), ('trg', trg_field_processor)]\n",
    "    MAX_LEN = 100  # filter out examples that have more than MAX_LEN tokens\n",
    "    filter_pred = lambda x: len(x.src) <= MAX_LEN and len(x.trg) <= MAX_LEN\n",
    "\n",
    "    # Caching mechanism stuff - avoid redoing the tokenization over and over again\n",
    "    prefix = 'de_en' if german_to_english else 'en_de'\n",
    "    prefix += '_iwslt' if use_iwslt else '_wmt14'\n",
    "    train_cache_path = os.path.join(dataset_path, f'{prefix}_train_cache.csv')\n",
    "    val_cache_path = os.path.join(dataset_path, f'{prefix}_val_cache.csv')\n",
    "    test_cache_path = os.path.join(dataset_path, f'{prefix}_test_cache.csv')\n",
    "\n",
    "    # This simple caching mechanism gave me ~30x speedup on my machine! From ~70s -> ~2.5s!\n",
    "    ts = time.time()\n",
    "    if not use_caching_mechanism or not (os.path.exists(train_cache_path) and os.path.exists(val_cache_path)):\n",
    "        # dataset objects have a list of examples where example is simply an empty Python Object that has\n",
    "        # .src and .trg attributes which contain a tokenized list of strings (created by tokenize_en and tokenize_de).\n",
    "        # It's that simple, we can consider our datasets as a table with 2 columns 'src' and 'trg'\n",
    "        # each containing fields with tokenized strings from source and target languages\n",
    "        src_ext = '.de' if german_to_english else '.en'\n",
    "        trg_ext = '.en' if german_to_english else '.de'\n",
    "        dataset_split_fn = datasets.IWSLT.splits if use_iwslt else datasets.WMT14.splits\n",
    "        \n",
    "        train_dataset, val_dataset, test_dataset = dataset_split_fn(\n",
    "            exts=(src_ext, trg_ext),\n",
    "            fields=fields,\n",
    "            root=dataset_path,\n",
    "            filter_pred=filter_pred\n",
    "        )\n",
    "\n",
    "        save_cache(train_cache_path, train_dataset)\n",
    "        save_cache(val_cache_path, val_dataset)\n",
    "        save_cache(test_cache_path, test_dataset)\n",
    "    else:\n",
    "        # it's actually better to load from cache as we'll get rid of '\\xa0', '\\xa0 ' and '\\x85' unicode characters\n",
    "        # which we don't need and which SpaCy unfortunately includes as tokens.\n",
    "        train_dataset, val_dataset = DatasetWrapper.get_train_and_val_datasets(\n",
    "            train_cache_path,\n",
    "            val_cache_path,\n",
    "            fields,\n",
    "            filter_pred=filter_pred\n",
    "        )\n",
    "\n",
    "    print(f'Time it took to prepare the data: {time.time() - ts:3f} seconds.')\n",
    "\n",
    "    MIN_FREQ = 2\n",
    "    # __getattr__ implementation in the base Dataset class enables us to call .src on Dataset objects even though\n",
    "    # we only have a list of examples in the Dataset object and the example itself had .src attribute.\n",
    "    # Implementation will yield examples and call .src/.trg attributes on them (and those contain tokenized lists)\n",
    "    src_field_processor.build_vocab(train_dataset.src, min_freq=MIN_FREQ)\n",
    "    trg_field_processor.build_vocab(train_dataset.trg, min_freq=MIN_FREQ)\n",
    "\n",
    "    return train_dataset, val_dataset, src_field_processor, trg_field_processor\n",
    "\n",
    "\n",
    "# https://github.com/pytorch/text/issues/536#issuecomment-719945594 <- there is a \"bug\" in BucketIterator i.e. it's\n",
    "# description is misleading as it won't group examples of similar length unless you set sort_within_batch to True!\n",
    "def get_data_loaders(dataset_path, language_direction, dataset_name, batch_size, device):\n",
    "    train_dataset, val_dataset, src_field_processor, trg_field_processor = get_datasets_and_vocabs(dataset_path, language_direction, dataset_name == DatasetType.IWSLT.name)\n",
    "\n",
    "    train_token_ids_loader, val_token_ids_loader = BucketIterator.splits(\n",
    "     datasets=(train_dataset, val_dataset),\n",
    "     batch_size=batch_size,\n",
    "     device=device,\n",
    "     sort_within_batch=True,  # this part is really important otherwise we won't group similar length sentences\n",
    "     batch_size_fn=batch_size_fn  # this helps us max out GPU's VRAM\n",
    "    )\n",
    "\n",
    "    return train_token_ids_loader, val_token_ids_loader, src_field_processor, trg_field_processor\n",
    "\n",
    "\n",
    "global longest_src_sentence, longest_trg_sentence\n",
    "\n",
    "\n",
    "def batch_size_fn(new_example, count, sofar):\n",
    "    \"\"\"\n",
    "        If we use this function in the BucketIterator the batch_size is no longer the number of examples/sentences\n",
    "        in a batch but a number of tokens in a batch - which allows us to max out VRAM on a given GPU.\n",
    "\n",
    "        Example: if we don't use this function and we set batch size to say 10 we will sometimes end up with\n",
    "        a tensor of size (10, 100) because the longest sentence had a size of 100 tokens but other times we'll end\n",
    "        up with a size of (10, 5) because the longest sentence had only 5 tokens!\n",
    "\n",
    "        With this function what we do is we specify that source and target tensors can't go over a certain number\n",
    "        of tokens like 1000. So usually either source or target tensors will contain around 1000 tokens and\n",
    "        in worst case both will be really close to a 1000 tokens each. If that is still below max VRAM availabe on\n",
    "        the system we're using the max potential of our GPU w.r.t. VRAM.\n",
    "\n",
    "        Note: to understand this function you unfortunately would probably have to dig deeper into torch text's\n",
    "        source code.\n",
    "\n",
    "    \"\"\"\n",
    "    global longest_src_sentence, longest_trg_sentence\n",
    "\n",
    "    if count == 1:\n",
    "        longest_src_sentence = 0\n",
    "        longest_trg_sentence = 0\n",
    "\n",
    "    longest_src_sentence = max(longest_src_sentence, len(new_example.src))\n",
    "    # 2 because of start/end of sentence tokens (<s> and </s>)\n",
    "    longest_trg_sentence = max(longest_trg_sentence, len(new_example.trg) + 2)\n",
    "\n",
    "    num_of_tokens_in_src_tensor = count * longest_src_sentence\n",
    "    num_of_tokens_in_trg_tensor = count * longest_trg_sentence\n",
    "\n",
    "    return max(num_of_tokens_in_src_tensor, num_of_tokens_in_trg_tensor)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Those were the main functions we need to later feed the data to the transformer inside the training loop!\n",
    "\n",
    "Now let's just define some functions which will get called in the training loop and will provide us with:\n",
    "1. correct source/target masks\n",
    "2. correct batches of token ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_masks_and_count_tokens_src(src_token_ids_batch, pad_token_id):\n",
    "    batch_size = src_token_ids_batch.shape[0]\n",
    "\n",
    "    # src_mask shape = (B, 1, 1, S) check out attention function in transformer_model.py where masks are applied\n",
    "    # src_mask only masks pad tokens as we want to ignore their representations (no information in there...)\n",
    "    src_mask = (src_token_ids_batch != pad_token_id).view(batch_size, 1, 1, -1)\n",
    "    num_src_tokens = torch.sum(src_mask.long())\n",
    "\n",
    "    return src_mask, num_src_tokens\n",
    "\n",
    "\n",
    "def get_masks_and_count_tokens_trg(trg_token_ids_batch, pad_token_id):\n",
    "    batch_size = trg_token_ids_batch.shape[0]\n",
    "    device = trg_token_ids_batch.device\n",
    "\n",
    "    # Same as src_mask but we additionally want to mask tokens from looking forward into the future tokens\n",
    "    # Note: wherever the mask value is true we want to attend to that token, otherwise we mask (ignore) it.\n",
    "    sequence_length = trg_token_ids_batch.shape[1]  # trg_token_ids shape = (B, T) where T max trg token-sequence length\n",
    "    trg_padding_mask = (trg_token_ids_batch != pad_token_id).view(batch_size, 1, 1, -1)  # shape = (B, 1, 1, T)\n",
    "    trg_no_look_forward_mask = torch.triu(torch.ones((1, 1, sequence_length, sequence_length), device=device) == 1).transpose(2, 3)\n",
    "\n",
    "    # logic AND operation (both padding mask and no-look-forward must be true to attend to a certain target token)\n",
    "    trg_mask = trg_padding_mask & trg_no_look_forward_mask  # final shape = (B, 1, T, T)\n",
    "    num_trg_tokens = torch.sum(trg_padding_mask.long())\n",
    "\n",
    "    return trg_mask, num_trg_tokens\n",
    "\n",
    "\n",
    "def get_masks_and_count_tokens(src_token_ids_batch, trg_token_ids_batch, pad_token_id, device):\n",
    "    src_mask, num_src_tokens = get_masks_and_count_tokens_src(src_token_ids_batch, pad_token_id)\n",
    "    trg_mask, num_trg_tokens = get_masks_and_count_tokens_trg(trg_token_ids_batch, pad_token_id)\n",
    "\n",
    "    return src_mask, trg_mask, num_src_tokens, num_trg_tokens\n",
    "\n",
    "\n",
    "def get_src_and_trg_batches(token_ids_batch):\n",
    "    src_token_ids_batch, trg_token_ids_batch = token_ids_batch.src, token_ids_batch.trg\n",
    "\n",
    "    # Target input should be shifted by 1 compared to the target output tokens\n",
    "    # Example: if we had a sentence like: [<s>,what,is,up,</s>] then to train the NMT model what we do is we pass\n",
    "    # [<s>,what,is,up] to the input as set [what,is,up,</s>] as the expected output.\n",
    "    trg_token_ids_batch_input = trg_token_ids_batch[:, :-1]\n",
    "\n",
    "    # We reshape from (B, S) into (BxS, 1) as that's the the shape expected by LabelSmoothing which will produce\n",
    "    # the shape (BxS, V) where V is the target vocab size which is the same shape as the one that comes out\n",
    "    # from the transformer so we can directly pass them into the KL divergence loss\n",
    "    trg_token_ids_batch_gt = trg_token_ids_batch[:, 1:].reshape(-1, 1)\n",
    "\n",
    "    return src_token_ids_batch, trg_token_ids_batch_input, trg_token_ids_batch_gt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Caching mechanism\n",
    "\n",
    "Optional section, just make sure to run the cell but it's not crucial for understanding how transformers work.\n",
    "\n",
    "We'll need to add some caching mechanism otherwise running the training will always take lots of time before the actual training loop starts running."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "#\n",
    "# Caching mechanism datasets and functions (you don't need this but it makes things a lot faster!)\n",
    "#\n",
    "\n",
    "\n",
    "class FastTranslationDataset(Dataset):\n",
    "    \"\"\"\n",
    "        After understanding the source code of torch text's IWSLT, TranslationDataset and Dataset I realized how I\n",
    "        can make data preparation much faster (tokenization was taking a lot of time and there is no need to redo it\n",
    "        every time) by using a simple caching mechanism.\n",
    "\n",
    "        This dataset leverages that caching mechanism which reduced loading time from ~70s -> 2.5s (massive!)\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    @staticmethod\n",
    "    def sort_key(ex):\n",
    "        # What this does is basically it takes a 16-bit binary representation of lengths and interleaves them.\n",
    "        # Example: lengths len(ex.src)=5 and len(ex.trg)=3 result in f(101, 011)=100111, 7 and 1 in f(111, 001)=101011\n",
    "        # It's basically a heuristic that helps the BucketIterator sort bigger batches first\n",
    "        return interleave_keys(len(ex.src), len(ex.trg))\n",
    "\n",
    "    def __init__(self, cache_path, fields, **kwargs):\n",
    "        # save_cache interleaves src and trg examples so here we read the cache file having that format in mind\n",
    "        cached_data = [line.split() for line in open(cache_path, encoding='utf-8')]\n",
    "\n",
    "        cached_data_src = cached_data[0::2]  # Even lines contain source examples\n",
    "        cached_data_trg = cached_data[1::2]  # Odd lines contain target examples\n",
    "\n",
    "        assert len(cached_data_src) == len(cached_data_trg), f'Source and target data should be of the same length.'\n",
    "\n",
    "        examples = []\n",
    "        src_dataset_total_number_of_tokens = 0\n",
    "        trg_dataset_total_number_of_tokens = 0\n",
    "        for src_tokenized_data, trg_tokenized_data in zip(cached_data_src, cached_data_trg):\n",
    "            ex = Example()\n",
    "\n",
    "            setattr(ex, 'src', src_tokenized_data)\n",
    "            setattr(ex, 'trg', trg_tokenized_data)\n",
    "\n",
    "            examples.append(ex)\n",
    "\n",
    "            # Update the number of tokens\n",
    "            src_dataset_total_number_of_tokens += len(src_tokenized_data)\n",
    "            trg_dataset_total_number_of_tokens += len(trg_tokenized_data)\n",
    "\n",
    "        # Print relevant information about the dataset (parsing the cache file name)\n",
    "        filename_parts = os.path.split(cache_path)[1].split('_')\n",
    "        src_language, trg_language = ('English', 'German') if filename_parts[0] == 'en' else ('German', 'English')\n",
    "        dataset_name = 'IWSLT' if filename_parts[2] == 'iwslt' else 'WMT-14'\n",
    "        dataset_type = 'train' if filename_parts[3] == 'train' else 'val'\n",
    "        print(f'{dataset_type} dataset ({dataset_name}) has {src_dataset_total_number_of_tokens} tokens in the source language ({src_language}) corpus.')\n",
    "        print(f'{dataset_type} dataset ({dataset_name}) has {trg_dataset_total_number_of_tokens} tokens in the target language ({trg_language}) corpus.')\n",
    "\n",
    "        # Call the parent class Dataset's constructor\n",
    "        super().__init__(examples, fields, **kwargs)\n",
    "\n",
    "\n",
    "class DatasetWrapper(FastTranslationDataset):\n",
    "    \"\"\"\n",
    "        Just a wrapper around the FastTranslationDataset.\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    @classmethod\n",
    "    def get_train_and_val_datasets(cls, train_cache_path, val_cache_path, fields, **kwargs):\n",
    "\n",
    "        train_dataset = cls(train_cache_path, fields, **kwargs)\n",
    "        val_dataset = cls(val_cache_path, fields, **kwargs)\n",
    "\n",
    "        return train_dataset, val_dataset\n",
    "\n",
    "\n",
    "def save_cache(cache_path, dataset):\n",
    "    with open(cache_path, 'w', encoding='utf-8') as cache_file:\n",
    "        # Interleave source and target tokenized examples, source is on even lines, target is on odd lines\n",
    "        for ex in dataset.examples:\n",
    "            cache_file.write(' '.join(ex.src) + '\\n')\n",
    "            cache_file.write(' '.join(ex.trg) + '\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Phew, we've got our data pipeline 📜 and model 🦄 all ready to go!\n",
    "\n",
    "Level 3 unlocked. 😍\n",
    "\n",
    "**Let me again stress this - it takes time to understand this project don't get frustrated! ❤️ (it took me ~3 weeks to code this up!)**\n",
    "\n",
    "Now there are a couple more things we need to take care of.\n",
    "\n",
    "# Part 3: Training loop utilities\n",
    "\n",
    "So we're done setting up the model and data! Now let's add a couple of other things we'll need in the training loop.\n",
    "\n",
    "First, let's add a wrapper around Adam which will give it a custom learning rate.\n",
    "\n",
    "<img src=\"data/readme_pics/custom_learning_rate_schedule.PNG\" alt=\"Custom LR Adam\"/>\n",
    "\n",
    "It's also a good design to make it behave as PyTorch's native optimizers (by defining `step()` and `zero_grad()` functions)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "class CustomLRAdamOptimizer:\n",
    "    \"\"\"\n",
    "        Linear ramp learning rate for the warm-up number of steps and then start decaying\n",
    "        according to the inverse square root law of the current training step number.\n",
    "\n",
    "        Check out playground.py for visualization of the learning rate (visualize_custom_lr_adam).\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, optimizer, model_dimension, num_of_warmup_steps):\n",
    "        self.optimizer = optimizer\n",
    "        self.model_size = model_dimension\n",
    "        self.num_of_warmup_steps = num_of_warmup_steps\n",
    "\n",
    "        self.current_step_number = 0\n",
    "\n",
    "    def step(self):\n",
    "        self.current_step_number += 1\n",
    "        current_learning_rate = self.get_current_learning_rate()\n",
    "\n",
    "        for p in self.optimizer.param_groups:\n",
    "            p['lr'] = current_learning_rate\n",
    "\n",
    "        self.optimizer.step()  # apply gradients\n",
    "\n",
    "    # Check out the formula at Page 7, Chapter 5.3 \"Optimizer\" and playground.py for visualization\n",
    "    def get_current_learning_rate(self):\n",
    "        # For readability purpose\n",
    "        step = self.current_step_number\n",
    "        warmup = self.num_of_warmup_steps\n",
    "\n",
    "        return self.model_size ** (-0.5) * min(step ** (-0.5), step * warmup ** (-1.5))\n",
    "\n",
    "    def zero_grad(self):\n",
    "        self.optimizer.zero_grad()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Label smoothing, what's that? Well you usually set your target vocabulary distribution to a `one-hot`. Meaning 1 position out of 30k (or whatever your vocab size is) is set to 1. probability and everything else to 0.\n",
    "\n",
    "In `label smoothing` instead of placing 1. on that particular position you place say 0.9 and you evenly distribute the rest of the \"probability mass\" (0.1 in this case) over the other positions (that's visualized as a different shade of purple on the image above in a fictional vocab of size 4 - hence 4 columns):\n",
    "\n",
    "<img src=\"data/readme_pics/label_smoothing.PNG\" width=\"700\" alt=\"Label smoothing example\"/>\n",
    "\n",
    "*Note: Pad token's distribution is set to all zeros as we don't want our model to predict those!*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LabelSmoothingDistribution(nn.Module):\n",
    "    \"\"\"\n",
    "        Instead of one-hot target distribution set the target word's probability to \"confidence_value\" (usually 0.9)\n",
    "        and distribute the rest of the \"smoothing_value\" mass (usually 0.1) over the rest of the vocab.\n",
    "\n",
    "        Check out playground.py for visualization of how the smooth target distribution looks like compared to one-hot.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, smoothing_value, pad_token_id, trg_vocab_size, device):\n",
    "        assert 0.0 <= smoothing_value <= 1.0\n",
    "\n",
    "        super(LabelSmoothingDistribution, self).__init__()\n",
    "\n",
    "        self.confidence_value = 1.0 - smoothing_value\n",
    "        self.smoothing_value = smoothing_value\n",
    "\n",
    "        self.pad_token_id = pad_token_id\n",
    "        self.trg_vocab_size = trg_vocab_size\n",
    "        self.device = device\n",
    "\n",
    "    def forward(self, trg_token_ids_batch):\n",
    "\n",
    "        batch_size = trg_token_ids_batch.shape[0]\n",
    "        smooth_target_distributions = torch.zeros((batch_size, self.trg_vocab_size), device=self.device)\n",
    "\n",
    "        # -2 because we are not distributing the smoothing mass over the pad token index and over the ground truth index\n",
    "        # those 2 values will be overwritten by the following 2 lines with confidence_value and 0 (for pad token index)\n",
    "        smooth_target_distributions.fill_(self.smoothing_value / (self.trg_vocab_size - 2))\n",
    "\n",
    "        smooth_target_distributions.scatter_(1, trg_token_ids_batch, self.confidence_value)\n",
    "        smooth_target_distributions[:, self.pad_token_id] = 0.\n",
    "\n",
    "        # If we had a pad token as a target we set the distribution to all 0s instead of smooth labeled distribution\n",
    "        smooth_target_distributions.masked_fill_(trg_token_ids_batch == self.pad_token_id, 0.)\n",
    "\n",
    "        return smooth_target_distributions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally we'll need a function that will help us evaluate our model in the training loop.\n",
    "\n",
    "We'll be using the [BLEU metric](https://en.wikipedia.org/wiki/BLEU), which is a common metric for the machine translation task.\n",
    "\n",
    "Here is an example of how using the Xavier initialization dramatically improved BLEU metric:\n",
    "\n",
    "<img src=\"data/readme_pics/bleu_score_xavier_vs_default_pt_init.PNG\" width=\"450\"/>\n",
    "\n",
    "You can see here 3 runs, the 2 lower ones used PyTorch default initialization (one used `mean` for KL divergence loss and the better one used `batchmean`), whereas the upper one used **Xavier uniform** initialization!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Calculate the BLEU-4 score\n",
    "def calculate_bleu_score(transformer, token_ids_loader, trg_field_processor):\n",
    "    with torch.no_grad():\n",
    "        pad_token_id = trg_field_processor.vocab.stoi[PAD_TOKEN]\n",
    "\n",
    "        gt_sentences_corpus = []\n",
    "        predicted_sentences_corpus = []\n",
    "\n",
    "        ts = time.time()\n",
    "        for batch_idx, token_ids_batch in enumerate(token_ids_loader):\n",
    "            src_token_ids_batch, trg_token_ids_batch = token_ids_batch.src, token_ids_batch.trg\n",
    "            if batch_idx % 10 == 0:\n",
    "                print(f'batch={batch_idx}, time elapsed = {time.time()-ts} seconds.')\n",
    "\n",
    "            # Optimization - compute the source token representations only once\n",
    "            src_mask, _ = get_masks_and_count_tokens_src(src_token_ids_batch, pad_token_id)\n",
    "            src_representations_batch = transformer.encode(src_token_ids_batch, src_mask)\n",
    "\n",
    "            predicted_sentences = greedy_decoding(transformer, src_representations_batch, src_mask, trg_field_processor)\n",
    "            predicted_sentences_corpus.extend(predicted_sentences)  # add them to the corpus of translations\n",
    "\n",
    "            # Get the token and not id version of GT (ground-truth) sentences\n",
    "            trg_token_ids_batch = trg_token_ids_batch.cpu().numpy()\n",
    "            for target_sentence_ids in trg_token_ids_batch:\n",
    "                target_sentence_tokens = [trg_field_processor.vocab.itos[id] for id in target_sentence_ids if id != pad_token_id]\n",
    "                gt_sentences_corpus.append([target_sentence_tokens])  # add them to the corpus of GT translations\n",
    "\n",
    "        bleu_score = corpus_bleu(gt_sentences_corpus, predicted_sentences_corpus)\n",
    "        print(f'BLEU-4 corpus score = {bleu_score}, corpus length = {len(gt_sentences_corpus)}, time elapsed = {time.time()-ts} seconds.')\n",
    "        return bleu_score"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see to calculate the BLEU score we need a decoding function.\n",
    "\n",
    "The original paper used **beam search** but for now I've only got **greedy** implemented. <br/> It's not crucial to have beam search to understand how transformers work so this is fine.\n",
    "\n",
    "*Note: The BLEU metric will give us a pessimistic estimate of our model's perf as the beam search usually gives better results.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "def greedy_decoding(baseline_transformer, src_representations_batch, src_mask, trg_field_processor, max_target_tokens=100):\n",
    "    \"\"\"\n",
    "    Supports batch (decode multiple source sentences) greedy decoding.\n",
    "\n",
    "    Decoding could be further optimized to cache old token activations because they can't look ahead and so\n",
    "    adding a newly predicted token won't change old token's activations.\n",
    "\n",
    "    Example: we input <s> and do a forward pass. We get intermediate activations for <s> and at the output at position\n",
    "    0, after the doing linear layer we get e.g. token <I>. Now we input <s>,<I> but <s>'s activations will remain\n",
    "    the same. Similarly say we now got <am> at output position 1, in the next step we input <s>,<I>,<am> and so <I>'s\n",
    "    activations will remain the same as it only looks at/attends to itself and to <s> and so forth.\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    device = next(baseline_transformer.parameters()).device\n",
    "    pad_token_id = trg_field_processor.vocab.stoi[PAD_TOKEN]\n",
    "\n",
    "    # Initial prompt is the beginning/start of the sentence token. Make it compatible shape with source batch => (B,1)\n",
    "    target_sentences_tokens = [[BOS_TOKEN] for _ in range(src_representations_batch.shape[0])]\n",
    "    trg_token_ids_batch = torch.tensor([[trg_field_processor.vocab.stoi[tokens[0]]] for tokens in target_sentences_tokens], device=device)\n",
    "\n",
    "    # Set to true for a particular target sentence once it reaches the EOS (end-of-sentence) token\n",
    "    is_decoded = [False] * src_representations_batch.shape[0]\n",
    "\n",
    "    while True:\n",
    "        trg_mask, _ = get_masks_and_count_tokens_trg(trg_token_ids_batch, pad_token_id)\n",
    "        # Shape = (B*T, V) where T is the current token-sequence length and V target vocab size\n",
    "        predicted_log_distributions = baseline_transformer.decode(trg_token_ids_batch, src_representations_batch, trg_mask, src_mask)\n",
    "\n",
    "        # Extract only the indices of last token for every target sentence (we take every T-th token)\n",
    "        num_of_trg_tokens = len(target_sentences_tokens[0])\n",
    "        predicted_log_distributions = predicted_log_distributions[num_of_trg_tokens-1::num_of_trg_tokens]\n",
    "\n",
    "        # This is the \"greedy\" part of the greedy decoding:\n",
    "        # We find indices of the highest probability target tokens and discard every other possibility\n",
    "        most_probable_last_token_indices = torch.argmax(predicted_log_distributions, dim=-1).cpu().numpy()\n",
    "\n",
    "        # Find target tokens associated with these indices\n",
    "        predicted_words = [trg_field_processor.vocab.itos[index] for index in most_probable_last_token_indices]\n",
    "\n",
    "        for idx, predicted_word in enumerate(predicted_words):\n",
    "            target_sentences_tokens[idx].append(predicted_word)\n",
    "\n",
    "            if predicted_word == EOS_TOKEN:  # once we find EOS token for a particular sentence we flag it\n",
    "                is_decoded[idx] = True\n",
    "\n",
    "        if all(is_decoded) or num_of_trg_tokens == max_target_tokens:\n",
    "            break\n",
    "\n",
    "        # Prepare the input for the next iteration (merge old token ids with the new column of most probable token ids)\n",
    "        trg_token_ids_batch = torch.cat((trg_token_ids_batch, torch.unsqueeze(torch.tensor(most_probable_last_token_indices, device=device), 1)), 1)\n",
    "\n",
    "    # Post process the sentences - remove everything after the EOS token\n",
    "    target_sentences_tokens_post = []\n",
    "    for target_sentence_tokens in target_sentences_tokens:\n",
    "        try:\n",
    "            target_index = target_sentence_tokens.index(EOS_TOKEN) + 1\n",
    "        except:\n",
    "            target_index = None\n",
    "\n",
    "        target_sentence_tokens = target_sentence_tokens[:target_index]\n",
    "        target_sentences_tokens_post.append(target_sentence_tokens)\n",
    "\n",
    "    return target_sentences_tokens_post"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A couple of more functions which will keep things nice and clean. \n",
    "\n",
    "Feel free to ignore these as they are not crucial for understanding how the transformer works."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_available_binary_name():\n",
    "    prefix = 'transformer'\n",
    "\n",
    "    def valid_binary_name(binary_name):\n",
    "        # First time you see raw f-string? Don't worry the only trick is to double the brackets.\n",
    "        pattern = re.compile(rf'{prefix}_[0-9]{{6}}\\.pth')\n",
    "        return re.fullmatch(pattern, binary_name) is not None\n",
    "\n",
    "    # Just list the existing binaries so that we don't overwrite them but write to a new one\n",
    "    valid_binary_names = list(filter(valid_binary_name, os.listdir(BINARIES_PATH)))\n",
    "    if len(valid_binary_names) > 0:\n",
    "        last_binary_name = sorted(valid_binary_names)[-1]\n",
    "        new_suffix = int(last_binary_name.split('.')[0][-6:]) + 1  # increment by 1\n",
    "        return f'{prefix}_{str(new_suffix).zfill(6)}.pth'\n",
    "    else:\n",
    "        return f'{prefix}_000000.pth'\n",
    "\n",
    "\n",
    "def get_training_state(training_config, model):\n",
    "    training_state = {\n",
    "        # \"commit_hash\": git.Repo(search_parent_directories=True).head.object.hexsha,\n",
    "        \"dataset_name\": training_config['dataset_name'],\n",
    "        \"language_direction\": training_config['language_direction'],\n",
    "\n",
    "        \"num_of_epochs\": training_config['num_of_epochs'],\n",
    "        \"batch_size\": training_config['batch_size'],\n",
    "\n",
    "        \"state_dict\": model.state_dict()\n",
    "    }\n",
    "\n",
    "    return training_state\n",
    "\n",
    "\n",
    "def print_model_metadata(training_state):\n",
    "    header = f'\\n{\"*\"*5} Model training metadata: {\"*\"*5}'\n",
    "    print(header)\n",
    "\n",
    "    for key, value in training_state.items():\n",
    "        if key != 'state_dict':  # don't print state_dict it's a bunch of numbers...\n",
    "            if key == 'language_direction':  # convert into human readable format\n",
    "                value = 'English to German' if value == 'E2G' else 'German to English'\n",
    "            print(f'{key}: {value}')\n",
    "    print(f'{\"*\" * len(header)}\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Level 4 unlocked. 😍\n",
    "\n",
    "**Let me again^2 stress this - it takes time to understand this project don't get frustrated! ❤️ (it took me ~3 weeks to code this up!)**\n",
    "\n",
    "# Part 4: Training loop\n",
    "\n",
    "Now let us combine all of the above sections together and finally train our transformer!\n",
    "\n",
    "Here are the functions we'll be using:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Global vars for logging purposes\n",
    "num_of_trg_tokens_processed = 0\n",
    "bleu_scores = []\n",
    "global_train_step, global_val_step = [0, 0]\n",
    "writer = SummaryWriter()  # (tensorboard) writer will output to ./runs/ directory by default\n",
    "\n",
    "\n",
    "# Simple decorator function so that I don't have to pass these arguments every time I call get_train_val_loop\n",
    "def get_train_val_loop(baseline_transformer, custom_lr_optimizer, kl_div_loss, label_smoothing, pad_token_id, time_start):\n",
    "\n",
    "    def train_val_loop(is_train, token_ids_loader, epoch):\n",
    "        global num_of_trg_tokens_processed, global_train_step, global_val_step, writer\n",
    "\n",
    "        if is_train:\n",
    "            baseline_transformer.train()\n",
    "        else:\n",
    "            baseline_transformer.eval()\n",
    "\n",
    "        device = next(baseline_transformer.parameters()).device\n",
    "\n",
    "        #\n",
    "        # Main loop - start of the CORE PART\n",
    "        #\n",
    "        for batch_idx, token_ids_batch in enumerate(token_ids_loader):\n",
    "            src_token_ids_batch, trg_token_ids_batch_input, trg_token_ids_batch_gt = get_src_and_trg_batches(token_ids_batch)\n",
    "            src_mask, trg_mask, num_src_tokens, num_trg_tokens = get_masks_and_count_tokens(src_token_ids_batch, trg_token_ids_batch_input, pad_token_id, device)\n",
    "\n",
    "            # log because the KL loss expects log probabilities (just an implementation detail)\n",
    "            predicted_log_distributions = baseline_transformer(src_token_ids_batch, trg_token_ids_batch_input, src_mask, trg_mask)\n",
    "            smooth_target_distributions = label_smoothing(trg_token_ids_batch_gt)  # these are regular probabilities\n",
    "\n",
    "            if is_train:\n",
    "                custom_lr_optimizer.zero_grad()  # clean the trainable weights gradients in the computational graph\n",
    "\n",
    "            loss = kl_div_loss(predicted_log_distributions, smooth_target_distributions)\n",
    "\n",
    "            if is_train:\n",
    "                loss.backward()  # compute the gradients for every trainable weight in the computational graph\n",
    "                custom_lr_optimizer.step()  # apply the gradients to weights\n",
    "\n",
    "            # End of CORE PART\n",
    "\n",
    "            #\n",
    "            # Logging and metrics\n",
    "            #\n",
    "\n",
    "            if is_train:\n",
    "                global_train_step += 1\n",
    "                num_of_trg_tokens_processed += num_trg_tokens\n",
    "\n",
    "                if training_config['enable_tensorboard']:\n",
    "                    writer.add_scalar('training_loss', loss.item(), global_train_step)\n",
    "\n",
    "                if training_config['console_log_freq'] is not None and batch_idx % training_config['console_log_freq'] == 0:\n",
    "                    print(f'Transformer training: time elapsed= {(time.time() - time_start):.2f} [s] '\n",
    "                          f'| epoch={epoch + 1} | batch= {batch_idx + 1} '\n",
    "                          f'| target tokens/batch= {num_of_trg_tokens_processed / training_config[\"console_log_freq\"]}')\n",
    "\n",
    "                    num_of_trg_tokens_processed = 0\n",
    "\n",
    "                # Save model checkpoint\n",
    "                if training_config['checkpoint_freq'] is not None and (epoch + 1) % training_config['checkpoint_freq'] == 0 and batch_idx == 0:\n",
    "                    ckpt_model_name = f\"transformer_ckpt_epoch_{epoch + 1}.pth\"\n",
    "                    torch.save(get_training_state(training_config, baseline_transformer), os.path.join(CHECKPOINTS_PATH, ckpt_model_name))\n",
    "            else:\n",
    "                global_val_step += 1\n",
    "\n",
    "                if training_config['enable_tensorboard']:\n",
    "                    writer.add_scalar('val_loss', loss.item(), global_val_step)\n",
    "\n",
    "    return train_val_loop\n",
    "\n",
    "\n",
    "def train_transformer(training_config):\n",
    "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")  # checking whether you have a GPU, I hope so!\n",
    "\n",
    "    # Step 1: Prepare data loaders\n",
    "    train_token_ids_loader, val_token_ids_loader, src_field_processor, trg_field_processor = get_data_loaders(\n",
    "        training_config['dataset_path'],\n",
    "        training_config['language_direction'],\n",
    "        training_config['dataset_name'],\n",
    "        training_config['batch_size'],\n",
    "        device)\n",
    "\n",
    "    pad_token_id = src_field_processor.vocab.stoi[PAD_TOKEN]  # pad token id is the same for target as well\n",
    "    src_vocab_size = len(src_field_processor.vocab)\n",
    "    trg_vocab_size = len(trg_field_processor.vocab)\n",
    "\n",
    "    # Step 2: Prepare the model (original transformer) and push to GPU\n",
    "    baseline_transformer = Transformer(\n",
    "        model_dimension=BASELINE_MODEL_DIMENSION,\n",
    "        src_vocab_size=src_vocab_size,\n",
    "        trg_vocab_size=trg_vocab_size,\n",
    "        number_of_heads=BASELINE_MODEL_NUMBER_OF_HEADS,\n",
    "        number_of_layers=BASELINE_MODEL_NUMBER_OF_LAYERS,\n",
    "        dropout_probability=BASELINE_MODEL_DROPOUT_PROB\n",
    "    ).to(device)\n",
    "\n",
    "    # Step 3: Prepare other training related utilities\n",
    "    kl_div_loss = nn.KLDivLoss(reduction='batchmean')  # gives better BLEU score than \"mean\"\n",
    "\n",
    "    # Makes smooth target distributions as opposed to conventional one-hot distributions\n",
    "    # My feeling is that this is a really dummy and arbitrary heuristic but time will tell.\n",
    "    label_smoothing = LabelSmoothingDistribution(BASELINE_MODEL_LABEL_SMOOTHING_VALUE, pad_token_id, trg_vocab_size, device)\n",
    "\n",
    "    # Check out playground.py for an intuitive visualization of how the LR changes with time/training steps, easy stuff.\n",
    "    custom_lr_optimizer = CustomLRAdamOptimizer(\n",
    "                Adam(baseline_transformer.parameters(), betas=(0.9, 0.98), eps=1e-9),\n",
    "                BASELINE_MODEL_DIMENSION,\n",
    "                training_config['num_warmup_steps']\n",
    "            )\n",
    "\n",
    "    # The decorator function makes things cleaner since there is a lot of redundancy between the train and val loops\n",
    "    train_val_loop = get_train_val_loop(baseline_transformer, custom_lr_optimizer, kl_div_loss, label_smoothing, pad_token_id, time.time())\n",
    "\n",
    "    # Step 4: Start the training\n",
    "    for epoch in range(training_config['num_of_epochs']):\n",
    "        # Training loop\n",
    "        train_val_loop(is_train=True, token_ids_loader=train_token_ids_loader, epoch=epoch)\n",
    "\n",
    "        # Validation loop\n",
    "        with torch.no_grad():\n",
    "            train_val_loop(is_train=False, token_ids_loader=val_token_ids_loader, epoch=epoch)\n",
    "\n",
    "            bleu_score = calculate_bleu_score(baseline_transformer, val_token_ids_loader, trg_field_processor)\n",
    "            if training_config['enable_tensorboard']:\n",
    "                writer.add_scalar('bleu_score', bleu_score, epoch)\n",
    "\n",
    "    # Save the latest transformer in the binaries directory\n",
    "    torch.save(get_training_state(training_config, baseline_transformer), os.path.join(BINARIES_PATH, get_available_binary_name()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Let's train it!\n",
    "\n",
    "And finally run this cell to start training your transformer! \n",
    "\n",
    "Feel free to skip it and proceed to the next section where I've added the option of automatic download of pretrained models (in case you don't want or you don't have enough resources to train it yourself)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#\n",
    "# Fixed args - don't change these unless you have a good reason\n",
    "#\n",
    "num_warmup_steps = 4000\n",
    "\n",
    "#\n",
    "# Modifiable args - feel free to play with these (only small subset is exposed by design to avoid cluttering)\n",
    "#\n",
    "parser = argparse.ArgumentParser()\n",
    "# According to the paper I infered that the baseline was trained for ~19 epochs on the WMT-14 dataset and I got\n",
    "# nice returns up to epoch ~20 on IWSLT as well (nice round number)\n",
    "parser.add_argument(\"--num_of_epochs\", type=int, help=\"number of training epochs\", default=20)\n",
    "# You should adjust this for your particular machine (I have RTX 2080 with 8 GBs of VRAM so 1500 fits nicely!)\n",
    "parser.add_argument(\"--batch_size\", type=int, help=\"target number of tokens in a src/trg batch\", default=500)\n",
    "\n",
    "# Data related args\n",
    "parser.add_argument(\"--dataset_name\", choices=[el.name for el in DatasetType], help='which dataset to use for training', default=DatasetType.IWSLT.name)\n",
    "parser.add_argument(\"--language_direction\", choices=[el.name for el in LanguageDirection], help='which direction to translate', default=LanguageDirection.E2G.name)\n",
    "parser.add_argument(\"--dataset_path\", type=str, help='download dataset to this path', default=DATA_DIR_PATH)\n",
    "\n",
    "# Logging/debugging/checkpoint related (helps a lot with experimentation)\n",
    "parser.add_argument(\"--enable_tensorboard\", type=bool, help=\"enable tensorboard logging\", default=True)\n",
    "parser.add_argument(\"--console_log_freq\", type=int, help=\"log to output console (batch) freq\", default=10)\n",
    "parser.add_argument(\"--checkpoint_freq\", type=int, help=\"checkpoint model saving (epoch) freq\", default=1)\n",
    "args = parser.parse_args(\"\")\n",
    "\n",
    "# Wrapping training configuration into a dictionary\n",
    "training_config = dict()\n",
    "for arg in vars(args):\n",
    "    training_config[arg] = getattr(args, arg)\n",
    "training_config['num_warmup_steps'] = num_warmup_steps\n",
    "\n",
    "# Train the original transformer model\n",
    "train_transformer(training_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Level 5 unlocked!!! 😍\n",
    "\n",
    "**Let me again^3 stress this - it takes time to understand this project don't get frustrated! ❤️ (it took me ~3 weeks to code this up!)**\n",
    "\n",
    "# Part 5: Let's translate! (machine translation example)\n",
    "\n",
    "Now for the interesting part (well depends whom you ask!)! Let's translate between **German** and **English**.\n",
    "\n",
    "But before we do that we'll want to add functions which will help us visualize the attention mechanism!\n",
    "\n",
    "As an example here are the attentions I get for the input sentence `Ich bin ein guter Mensch, denke ich.`\n",
    "\n",
    "These belong to layer 6 of the encoder. You can see all of the 8 multi-head attention heads.\n",
    "\n",
    "<img src=\"data/readme_pics/attention_enc_self.PNG\" width=\"850\"/>\n",
    "\n",
    "These help us understand as to which part of the source/target sentence is the transformer attending to."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_attention_heatmap(data, x, y, head_id, ax):\n",
    "    seaborn.heatmap(data, xticklabels=x, yticklabels=y, square=True, vmin=0.0, vmax=1.0, cbar=False, annot=True, fmt=\".2f\", ax=ax)\n",
    "    ax.set_title(f'MHA head id = {head_id}')\n",
    "\n",
    "\n",
    "def visualize_attention_helper(attention_weights, source_sentence_tokens=None, target_sentence_tokens=None, title=''):\n",
    "    num_columns = 4\n",
    "    num_rows = 2\n",
    "    fig, axs = plt.subplots(num_rows, num_columns, figsize=(20, 10))  # prepare the figure and axes\n",
    "\n",
    "    assert source_sentence_tokens is not None or target_sentence_tokens is not None, \\\n",
    "        f'Either source or target sentence must be passed in.'\n",
    "\n",
    "    target_sentence_tokens = source_sentence_tokens if target_sentence_tokens is None else target_sentence_tokens\n",
    "    source_sentence_tokens = target_sentence_tokens if source_sentence_tokens is None else source_sentence_tokens\n",
    "\n",
    "    for head_id, head_attention_weights in enumerate(attention_weights):\n",
    "        row_index = int(head_id / num_columns)\n",
    "        column_index = head_id % num_columns\n",
    "        plot_attention_heatmap(head_attention_weights, source_sentence_tokens, target_sentence_tokens if head_id % num_columns == 0 else [], head_id, axs[row_index, column_index])\n",
    "\n",
    "    fig.suptitle(title)\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "def visualize_attention(baseline_transformer, source_sentence_tokens, target_sentence_tokens):\n",
    "    encoder = baseline_transformer.encoder\n",
    "    decoder = baseline_transformer.decoder\n",
    "\n",
    "    # Remove the end of sentence token </s> as we never attend to it, it's produced at the output and we stop\n",
    "    target_sentence_tokens = target_sentence_tokens[0][:-1]\n",
    "\n",
    "    # Visualize encoder attention weights\n",
    "    for layer_id, encoder_layer in enumerate(encoder.encoder_layers):\n",
    "        mha = encoder_layer.multi_headed_attention  # Every encoder layer has 1 MHA module\n",
    "\n",
    "        # attention_weights shape = (B, NH, S, S), extract 0th batch and loop over NH (number of heads) MHA heads\n",
    "        # S stands for maximum source token-sequence length\n",
    "        attention_weights = mha.attention_weights.cpu().numpy()[0]\n",
    "\n",
    "        title = f'Encoder layer {layer_id + 1}'\n",
    "        visualize_attention_helper(attention_weights, source_sentence_tokens, title=title)\n",
    "\n",
    "    # Visualize decoder attention weights\n",
    "    for layer_id, decoder_layer in enumerate(decoder.decoder_layers):\n",
    "        mha_trg = decoder_layer.trg_multi_headed_attention  # Extract the self-attention MHA\n",
    "        mha_src = decoder_layer.src_multi_headed_attention  # Extract the source attending MHA\n",
    "\n",
    "        # attention_weights shape = (B, NH, T, T), T stands for maximum target token-sequence length\n",
    "        attention_weights_trg = mha_trg.attention_weights.cpu().numpy()[0]\n",
    "        # shape = (B, NH, T, S), target token representations create queries and keys/values come from the encoder\n",
    "        attention_weights_src = mha_src.attention_weights.cpu().numpy()[0]\n",
    "\n",
    "        title = f'Decoder layer {layer_id + 1}, self-attention MHA'\n",
    "        visualize_attention_helper(attention_weights_trg, target_sentence_tokens=target_sentence_tokens, title=title)\n",
    "\n",
    "        title = f'Decoder layer {layer_id + 1}, source-attending MHA'\n",
    "        visualize_attention_helper(attention_weights_src, source_sentence_tokens, target_sentence_tokens, title)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download pretrained transformers automatically\n",
    "\n",
    "In case you skipped training your own model this code will automatically download pretrained models for you once you start the next cell (`translate_a_single_sentence()` function)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "IWSLT_ENGLISH_TO_GERMAN_MODEL_URL = r'https://www.dropbox.com/s/a6pfo6t9m2dh1jq/iwslt_e2g.pth?dl=1'\n",
    "IWSLT_GERMAN_TO_ENGLISH_MODEL_URL = r'https://www.dropbox.com/s/dgcd4xhwig7ygqd/iwslt_g2e.pth?dl=1'\n",
    "\n",
    "\n",
    "# Not yet trained\n",
    "WMT14_ENGLISH_TO_GERMAN_MODEL_URL = None\n",
    "WMT14_GERMAN_TO_ENGLISH_MODEL_URL = None\n",
    "\n",
    "\n",
    "DOWNLOAD_DICT = {\n",
    "    'iwslt_e2g': IWSLT_ENGLISH_TO_GERMAN_MODEL_URL,\n",
    "    'iwslt_g2e': IWSLT_GERMAN_TO_ENGLISH_MODEL_URL,\n",
    "    'wmt14_e2g': WMT14_ENGLISH_TO_GERMAN_MODEL_URL,\n",
    "    'wmt14_g2e': WMT14_GERMAN_TO_ENGLISH_MODEL_URL\n",
    "}\n",
    "\n",
    "\n",
    "download_choices = list(DOWNLOAD_DICT.keys())\n",
    "\n",
    "\n",
    "def download_models(translation_config):\n",
    "    # Step 1: Form the key\n",
    "    language_direction = translation_config['language_direction'].lower()\n",
    "    dataset_name = translation_config['dataset_name'].lower()\n",
    "    key = f'{dataset_name}_{language_direction}'\n",
    "\n",
    "    # Step 2: Check whether this model already exists\n",
    "    model_name = f'{key}.pth'\n",
    "    model_path = os.path.join(BINARIES_PATH, model_name)\n",
    "    if os.path.exists(model_path):\n",
    "        print(f'No need to download, found model {model_path} that was trained on {dataset_name} for language direction {language_direction}.')\n",
    "        return model_path\n",
    "\n",
    "    # Step 3: Download the resource to local filesystem\n",
    "    remote_resource_path = DOWNLOAD_DICT[key]\n",
    "    if remote_resource_path is None:  # handle models which I've not provided URLs for yet\n",
    "        print(f'No model found that was trained on {dataset_name} for language direction {language_direction}.')\n",
    "        exit(0)\n",
    "\n",
    "    print(f'Downloading from {remote_resource_path}. This may take a while.')\n",
    "    download_url_to_file(remote_resource_path, model_path)\n",
    "\n",
    "    return model_path"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Neural Machine Translation ❤️ \n",
    "\n",
    "Finally this is where the magic happens!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train dataset (IWSLT) has 3937527 tokens in the source language (English) corpus.\n",
      "train dataset (IWSLT) has 3634135 tokens in the target language (German) corpus.\n",
      "val dataset (IWSLT) has 20911 tokens in the source language (English) corpus.\n",
      "val dataset (IWSLT) has 19540 tokens in the target language (German) corpus.\n",
      "Time it took to prepare the data: 2.328090 seconds.\n",
      "\n",
      "***** Model training metadata: *****\n",
      "dataset_name: IWSLT\n",
      "language_direction: English to German\n",
      "num_of_epochs: 20\n",
      "batch_size: 1500\n",
      "*************************************\n",
      "\n",
      "Source sentence tokens = ['How', 'are', 'you', 'doing', 'today', '?']\n",
      "Translation | Target sentence tokens = [['<s>', 'Wie', 'geht', 'es', 'Ihnen', 'heute', '?', '</s>']]\n"
     ]
    }
   ],
   "source": [
    "# Super easy to add translation for a batch of sentences passed as a .txt file for example\n",
    "def translate_a_single_sentence(translation_config):\n",
    "    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")  # checking whether you have a GPU\n",
    "\n",
    "    # Step 1: Prepare the field processor (tokenizer, numericalizer)\n",
    "    _, _, src_field_processor, trg_field_processor = get_datasets_and_vocabs(\n",
    "        translation_config['dataset_path'],\n",
    "        translation_config['language_direction'],\n",
    "        translation_config['dataset_name'] == DatasetType.IWSLT.name\n",
    "    )\n",
    "    assert src_field_processor.vocab.stoi[PAD_TOKEN] == trg_field_processor.vocab.stoi[PAD_TOKEN]\n",
    "    pad_token_id = src_field_processor.vocab.stoi[PAD_TOKEN]  # needed for constructing masks\n",
    "\n",
    "    # Step 2: Prepare the model\n",
    "    baseline_transformer = Transformer(\n",
    "        model_dimension=BASELINE_MODEL_DIMENSION,\n",
    "        src_vocab_size=len(src_field_processor.vocab),\n",
    "        trg_vocab_size=len(trg_field_processor.vocab),\n",
    "        number_of_heads=BASELINE_MODEL_NUMBER_OF_HEADS,\n",
    "        number_of_layers=BASELINE_MODEL_NUMBER_OF_LAYERS,\n",
    "        dropout_probability=BASELINE_MODEL_DROPOUT_PROB,\n",
    "        log_attention_weights=True\n",
    "    ).to(device)\n",
    "\n",
    "    model_path = os.path.join(BINARIES_PATH, translation_config['model_name'])\n",
    "    if not os.path.exists(model_path):\n",
    "        print(f'Model {model_path} does not exist, attempting to download.')\n",
    "        model_path = download_models(translation_config)\n",
    "\n",
    "    model_state = torch.load(model_path)\n",
    "    print_model_metadata(model_state)\n",
    "    baseline_transformer.load_state_dict(model_state[\"state_dict\"], strict=True)\n",
    "    baseline_transformer.eval()\n",
    "\n",
    "    # Step 3: Prepare the input sentence\n",
    "    source_sentence = translation_config['source_sentence']\n",
    "    ex = Example.fromlist([source_sentence], fields=[('src', src_field_processor)])  # tokenize the sentence\n",
    "\n",
    "    source_sentence_tokens = ex.src\n",
    "    print(f'Source sentence tokens = {source_sentence_tokens}')\n",
    "\n",
    "    # Numericalize and convert to cuda tensor\n",
    "    src_token_ids_batch = src_field_processor.process([source_sentence_tokens], device)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        # Step 4: Optimization - compute the source token representations only once\n",
    "        src_mask, _ = get_masks_and_count_tokens_src(src_token_ids_batch, pad_token_id)\n",
    "        src_representations_batch = baseline_transformer.encode(src_token_ids_batch, src_mask)\n",
    "\n",
    "        # Step 5: Decoding process\n",
    "        target_sentence_tokens = greedy_decoding(baseline_transformer, src_representations_batch, src_mask, trg_field_processor)\n",
    "        print(f'Translation | Target sentence tokens = {target_sentence_tokens}')\n",
    "\n",
    "        # Step 6: Potentially visualize the encoder/decoder attention weights\n",
    "        if translation_config['visualize_attention']:\n",
    "            visualize_attention(baseline_transformer, source_sentence_tokens, target_sentence_tokens)\n",
    "\n",
    "\n",
    "#\n",
    "# modifiable args - feel free to play with these (only small subset is exposed by design to avoid cluttering)\n",
    "#\n",
    "parser = argparse.ArgumentParser()\n",
    "parser.add_argument(\"--source_sentence\", type=str, help=\"source sentence to translate into target\", default=\"How are you doing today?\")\n",
    "parser.add_argument(\"--model_name\", type=str, help=\"transformer model name\", default=r'iwslt_e2g.pth')\n",
    "\n",
    "# Keep these 2 in sync with the model you pick via model_name\n",
    "parser.add_argument(\"--dataset_name\", type=str, choices=['IWSLT', 'WMT14'], help='which dataset to use for training', default=DatasetType.IWSLT.name)\n",
    "parser.add_argument(\"--language_direction\", type=str, choices=[el.name for el in LanguageDirection], help='which direction to translate', default=LanguageDirection.E2G.name)\n",
    "\n",
    "# Cache files and datasets are downloaded here during training, keep them in sync for speed\n",
    "parser.add_argument(\"--dataset_path\", type=str, help='download dataset to this path', default=DATA_DIR_PATH)\n",
    "\n",
    "# Decoding related args\n",
    "parser.add_argument(\"--beam_size\", type=int, help=\"used only in case decoding method is chosen\", default=4)\n",
    "parser.add_argument(\"--length_penalty_coefficient\", type=int, help=\"length penalty for the beam search\", default=0.6)\n",
    "\n",
    "parser.add_argument(\"--visualize_attention\", type=bool, help=\"should visualize encoder/decoder attention\", default=False)\n",
    "args = parser.parse_args(\"\")\n",
    "\n",
    "# Wrapping training configuration into a dictionary\n",
    "translation_config = dict()\n",
    "for arg in vars(args):\n",
    "    translation_config[arg] = getattr(args, arg)\n",
    "\n",
    "# Translate the given source sentence\n",
    "translate_a_single_sentence(translation_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Phew!!! That was a mouthful! If you stayed with me until here, **congrats!** (achievement unlocked - transformer master 😍)\n",
    "\n",
    "Take your time to analyze this notebook. This is not a toy project, it took me ~3 weeks to finish it <br/> \n",
    "so don't expect to understand everything in 30 minutes unless you're really familiar with most of the concepts mentioned here.\n",
    "\n",
    "And last but not least!\n",
    "\n",
    "# Connect with me\n",
    "\n",
    "I share lots of useful (I hope so at least!) content on LinkedIn, Twitter, YouTube and Medium. <br/>\n",
    "So feel free to connect with me there:\n",
    "1. My [LinkedIn](https://www.linkedin.com/in/aleksagordic) and [Twitter](https://twitter.com/gordic_aleksa) profiles\n",
    "2. My YouTube channel - [The AI Epiphany](https://www.youtube.com/c/TheAiEpiphany)\n",
    "3. My [Medium](https://gordicaleksa.medium.com/) profile\n",
    "\n",
    "Also do drop me a message if you found this useful or if you think I could've done something better! <br/>\n",
    "I always like getting some feedback on the work I do.\n",
    "\n",
    "If you notice some bugs/errors feel free to **open up an issue** or even **submit a pull request**.\n",
    "\n",
    "# Additional resources\n",
    "\n",
    "I highly recommend [this blog by Jay Alammar](https://jalammar.github.io/illustrated-transformer/) to everyone learning about transformers.\n",
    "\n",
    "\n",
    "Also do check out my [YouTube video](https://www.youtube.com/watch?v=bvBK-coXf9I) on how I structured my transformers learning journey!\n",
    "\n",
    "<a href=\"https://www.youtube.com/watch?v=bvBK-coXf9I\" target=\"_blank\"><img src=\"https://img.youtube.com/vi/bvBK-coXf9I/0.jpg\" \n",
    "alt=\"learning about transformers\" width=\"480\" align=\"left\" height=\"360\" border=\"10\" /></a>"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}