Attention_is_all_you_need.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Attention is all you need.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "authorship_tag": "ABX9TyOBsHwoVyiWgRsuD6XwzGwR",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/akashe/NLP/blob/main/Attention_is_all_you_need.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LFNnw8Zfm5PF"
      },
      "source": [
        "[The paper that changed NLP.](https://arxiv.org/pdf/1706.03762.pdf)\n",
        "\n",
        "Before Transformer, langauge models employed LSTMs and GRUs. Though both of them broke many SOA results, Transformers took language models to another level. Transformers have few advantages over traditional RNN:\n",
        "\n",
        "1. Longer dependencies: LSTMs can't work with very long sequences.[Explanation](https://akashe.io/blog/2020/12/03/rnn-lstm-gru-and-attention/#how-LSTM-solves-vanishing-exploding-gradients). Transformers don't unroll in time. Transformers depend only on the max_len of the sequence to find dependency. Max_len can be 100 or 10000.\n",
        "2. Faster computation: Transformer are highly parallelizable. LSTMs have to unroll in sequence dim. This limits how fast you can train a LSTM.\n",
        "\n",
        "What transformers lack:\n",
        "\n",
        "1. Notion of position: Transformers lack the notion of sequential data, so they are given an additional input of positional embeddings to sequential nature of the data.\n",
        "2. They don't have an internal state.[Paper](https://arxiv.org/pdf/2002.09402.pdf)\n",
        "\n",
        "\n",
        "The task:\n",
        "\n",
        "We will make and train a transformer architecture to translate from german to english.\n",
        "\n",
        "Note: Code heavily inspired from other places."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "awrKgBLktFvB"
      },
      "source": [
        "Import libs"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "BgnqrVC_tMRi"
      },
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "\n",
        "import torchtext\n",
        "from torchtext.datasets import Multi30k\n",
        "from torchtext.data import Field, BucketIterator\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import matplotlib.ticker as ticker\n",
        "\n",
        "import spacy\n",
        "import numpy as np\n",
        "\n",
        "import random\n",
        "import math\n",
        "import time"
      ],
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lfz2OUGetOkl"
      },
      "source": [
        "Set Seed"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eQ_GZ1n6tSY1"
      },
      "source": [
        "SEED = 1007\n",
        "\n",
        "random.seed(SEED)\n",
        "np.random.seed(SEED)\n",
        "torch.manual_seed(SEED)\n",
        "torch.cuda.manual_seed(SEED)\n",
        "torch.backends.cudnn.deterministic = True"
      ],
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RPF1KM5HtViX"
      },
      "source": [
        "Download spacy English and German models"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "a2-QLvEptcVp",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "85ffbf3a-aa85-40df-dc8f-43f680ed1c0c"
      },
      "source": [
        "!python -m spacy download en\n",
        "!python -m spacy download de"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Requirement already satisfied: en_core_web_sm==2.2.5 from https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-2.2.5/en_core_web_sm-2.2.5.tar.gz#egg=en_core_web_sm==2.2.5 in /usr/local/lib/python3.6/dist-packages (2.2.5)\n",
            "Requirement already satisfied: spacy>=2.2.2 in /usr/local/lib/python3.6/dist-packages (from en_core_web_sm==2.2.5) (2.2.4)\n",
            "Requirement already satisfied: blis<0.5.0,>=0.4.0 in /usr/local/lib/python3.6/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (0.4.1)\n",
            "Requirement already satisfied: preshed<3.1.0,>=3.0.2 in /usr/local/lib/python3.6/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (3.0.5)\n",
            "Requirement already satisfied: thinc==7.4.0 in /usr/local/lib/python3.6/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (7.4.0)\n",
            "Requirement already satisfied: wasabi<1.1.0,>=0.4.0 in /usr/local/lib/python3.6/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (0.8.2)\n",
            "Requirement already satisfied: tqdm<5.0.0,>=4.38.0 in /usr/local/lib/python3.6/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (4.41.1)\n",
            "Requirement already satisfied: cymem<2.1.0,>=2.0.2 in /usr/local/lib/python3.6/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (2.0.5)\n",
            "Requirement already satisfied: murmurhash<1.1.0,>=0.28.0 in /usr/local/lib/python3.6/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (1.0.5)\n",
            "Requirement already satisfied: plac<1.2.0,>=0.9.6 in /usr/local/lib/python3.6/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (1.1.3)\n",
            "Requirement already satisfied: srsly<1.1.0,>=1.0.2 in /usr/local/lib/python3.6/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (1.0.5)\n",
            "Requirement already satisfied: requests<3.0.0,>=2.13.0 in /usr/local/lib/python3.6/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (2.23.0)\n",
            "Requirement already satisfied: catalogue<1.1.0,>=0.0.7 in /usr/local/lib/python3.6/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (1.0.0)\n",
            "Requirement already satisfied: numpy>=1.15.0 in /usr/local/lib/python3.6/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (1.19.5)\n",
            "Requirement already satisfied: setuptools in /usr/local/lib/python3.6/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (53.0.0)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.6/dist-packages (from requests<3.0.0,>=2.13.0->spacy>=2.2.2->en_core_web_sm==2.2.5) (2.10)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.6/dist-packages (from requests<3.0.0,>=2.13.0->spacy>=2.2.2->en_core_web_sm==2.2.5) (2020.12.5)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.6/dist-packages (from requests<3.0.0,>=2.13.0->spacy>=2.2.2->en_core_web_sm==2.2.5) (3.0.4)\n",
            "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.6/dist-packages (from requests<3.0.0,>=2.13.0->spacy>=2.2.2->en_core_web_sm==2.2.5) (1.24.3)\n",
            "Requirement already satisfied: importlib-metadata>=0.20; python_version < \"3.8\" in /usr/local/lib/python3.6/dist-packages (from catalogue<1.1.0,>=0.0.7->spacy>=2.2.2->en_core_web_sm==2.2.5) (3.4.0)\n",
            "Requirement already satisfied: typing-extensions>=3.6.4; python_version < \"3.8\" in /usr/local/lib/python3.6/dist-packages (from importlib-metadata>=0.20; python_version < \"3.8\"->catalogue<1.1.0,>=0.0.7->spacy>=2.2.2->en_core_web_sm==2.2.5) (3.7.4.3)\n",
            "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.6/dist-packages (from importlib-metadata>=0.20; python_version < \"3.8\"->catalogue<1.1.0,>=0.0.7->spacy>=2.2.2->en_core_web_sm==2.2.5) (3.4.0)\n",
            "\u001b[38;5;2m✔ Download and installation successful\u001b[0m\n",
            "You can now load the model via spacy.load('en_core_web_sm')\n",
            "\u001b[38;5;2m✔ Linking successful\u001b[0m\n",
            "/usr/local/lib/python3.6/dist-packages/en_core_web_sm -->\n",
            "/usr/local/lib/python3.6/dist-packages/spacy/data/en\n",
            "You can now load the model via spacy.load('en')\n",
            "Requirement already satisfied: de_core_news_sm==2.2.5 from https://github.com/explosion/spacy-models/releases/download/de_core_news_sm-2.2.5/de_core_news_sm-2.2.5.tar.gz#egg=de_core_news_sm==2.2.5 in /usr/local/lib/python3.6/dist-packages (2.2.5)\n",
            "Requirement already satisfied: spacy>=2.2.2 in /usr/local/lib/python3.6/dist-packages (from de_core_news_sm==2.2.5) (2.2.4)\n",
            "Requirement already satisfied: tqdm<5.0.0,>=4.38.0 in /usr/local/lib/python3.6/dist-packages (from spacy>=2.2.2->de_core_news_sm==2.2.5) (4.41.1)\n",
            "Requirement already satisfied: numpy>=1.15.0 in /usr/local/lib/python3.6/dist-packages (from spacy>=2.2.2->de_core_news_sm==2.2.5) (1.19.5)\n",
            "Requirement already satisfied: setuptools in /usr/local/lib/python3.6/dist-packages (from spacy>=2.2.2->de_core_news_sm==2.2.5) (53.0.0)\n",
            "Requirement already satisfied: catalogue<1.1.0,>=0.0.7 in /usr/local/lib/python3.6/dist-packages (from spacy>=2.2.2->de_core_news_sm==2.2.5) (1.0.0)\n",
            "Requirement already satisfied: srsly<1.1.0,>=1.0.2 in /usr/local/lib/python3.6/dist-packages (from spacy>=2.2.2->de_core_news_sm==2.2.5) (1.0.5)\n",
            "Requirement already satisfied: murmurhash<1.1.0,>=0.28.0 in /usr/local/lib/python3.6/dist-packages (from spacy>=2.2.2->de_core_news_sm==2.2.5) (1.0.5)\n",
            "Requirement already satisfied: thinc==7.4.0 in /usr/local/lib/python3.6/dist-packages (from spacy>=2.2.2->de_core_news_sm==2.2.5) (7.4.0)\n",
            "Requirement already satisfied: blis<0.5.0,>=0.4.0 in /usr/local/lib/python3.6/dist-packages (from spacy>=2.2.2->de_core_news_sm==2.2.5) (0.4.1)\n",
            "Requirement already satisfied: cymem<2.1.0,>=2.0.2 in /usr/local/lib/python3.6/dist-packages (from spacy>=2.2.2->de_core_news_sm==2.2.5) (2.0.5)\n",
            "Requirement already satisfied: plac<1.2.0,>=0.9.6 in /usr/local/lib/python3.6/dist-packages (from spacy>=2.2.2->de_core_news_sm==2.2.5) (1.1.3)\n",
            "Requirement already satisfied: preshed<3.1.0,>=3.0.2 in /usr/local/lib/python3.6/dist-packages (from spacy>=2.2.2->de_core_news_sm==2.2.5) (3.0.5)\n",
            "Requirement already satisfied: requests<3.0.0,>=2.13.0 in /usr/local/lib/python3.6/dist-packages (from spacy>=2.2.2->de_core_news_sm==2.2.5) (2.23.0)\n",
            "Requirement already satisfied: wasabi<1.1.0,>=0.4.0 in /usr/local/lib/python3.6/dist-packages (from spacy>=2.2.2->de_core_news_sm==2.2.5) (0.8.2)\n",
            "Requirement already satisfied: importlib-metadata>=0.20; python_version < \"3.8\" in /usr/local/lib/python3.6/dist-packages (from catalogue<1.1.0,>=0.0.7->spacy>=2.2.2->de_core_news_sm==2.2.5) (3.4.0)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.6/dist-packages (from requests<3.0.0,>=2.13.0->spacy>=2.2.2->de_core_news_sm==2.2.5) (2.10)\n",
            "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.6/dist-packages (from requests<3.0.0,>=2.13.0->spacy>=2.2.2->de_core_news_sm==2.2.5) (1.24.3)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.6/dist-packages (from requests<3.0.0,>=2.13.0->spacy>=2.2.2->de_core_news_sm==2.2.5) (3.0.4)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.6/dist-packages (from requests<3.0.0,>=2.13.0->spacy>=2.2.2->de_core_news_sm==2.2.5) (2020.12.5)\n",
            "Requirement already satisfied: typing-extensions>=3.6.4; python_version < \"3.8\" in /usr/local/lib/python3.6/dist-packages (from importlib-metadata>=0.20; python_version < \"3.8\"->catalogue<1.1.0,>=0.0.7->spacy>=2.2.2->de_core_news_sm==2.2.5) (3.7.4.3)\n",
            "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.6/dist-packages (from importlib-metadata>=0.20; python_version < \"3.8\"->catalogue<1.1.0,>=0.0.7->spacy>=2.2.2->de_core_news_sm==2.2.5) (3.4.0)\n",
            "\u001b[38;5;2m✔ Download and installation successful\u001b[0m\n",
            "You can now load the model via spacy.load('de_core_news_sm')\n",
            "\u001b[38;5;2m✔ Linking successful\u001b[0m\n",
            "/usr/local/lib/python3.6/dist-packages/de_core_news_sm -->\n",
            "/usr/local/lib/python3.6/dist-packages/spacy/data/de\n",
            "You can now load the model via spacy.load('de')\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zG2JU2mftf1q"
      },
      "source": [
        "Load the models"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Hkxg_IJsti5S"
      },
      "source": [
        "spacy_de = spacy.load('de')\n",
        "spacy_en = spacy.load('en')"
      ],
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9acztrSZtnOT"
      },
      "source": [
        "Define tokenizers"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "J2QSYcFVtpk8"
      },
      "source": [
        "def tokenize_de(text):\n",
        "    \"\"\"\n",
        "    Tokenizes German text from a string into a list of strings\n",
        "    \"\"\"\n",
        "    return [tok.text for tok in spacy_de.tokenizer(text)]\n",
        "\n",
        "def tokenize_en(text):\n",
        "    \"\"\"\n",
        "    Tokenizes English text from a string into a list of strings\n",
        "    \"\"\"\n",
        "    return [tok.text for tok in spacy_en.tokenizer(text)]"
      ],
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PdeDjBZitxDl"
      },
      "source": [
        "Define source and target fields"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "G5XhDUvut0aV"
      },
      "source": [
        "SRC = Field(tokenize = tokenize_de, \n",
        "            init_token = '<sos>', \n",
        "            eos_token = '<eos>', \n",
        "            lower = True, \n",
        "            batch_first = True)\n",
        "\n",
        "TRG = Field(tokenize = tokenize_en, \n",
        "            init_token = '<sos>', \n",
        "            eos_token = '<eos>', \n",
        "            lower = True, \n",
        "            batch_first = True)"
      ],
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t7aBOaKXt3LX"
      },
      "source": [
        "Download the data"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qBtBG2Dxt7yQ"
      },
      "source": [
        "train_data, valid_data, test_data = Multi30k.splits(exts = ('.de', '.en'), \n",
        "                                                    fields = (SRC, TRG))"
      ],
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "01Xn9bi-uCdj"
      },
      "source": [
        "Build vocab"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pw_Z9lwRuEMN"
      },
      "source": [
        "SRC.build_vocab(train_data, min_freq = 2)\n",
        "TRG.build_vocab(train_data, min_freq = 2)"
      ],
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Cp9PQAs4uMCV"
      },
      "source": [
        "Set device"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Qb4u7IhsuOH5",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "4bc35f0c-04a5-438a-fde9-ae06ab001096"
      },
      "source": [
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "device"
      ],
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "device(type='cuda')"
            ]
          },
          "metadata": {
            "tags": []
          },
          "execution_count": 9
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7ZJ7vzKPuPyd"
      },
      "source": [
        "Build iterators"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Z0beebsuuS3X"
      },
      "source": [
        "BATCH_SIZE = 128\n",
        "\n",
        "train_iterator, valid_iterator, test_iterator = BucketIterator.splits(\n",
        "    (train_data, valid_data, test_data), \n",
        "     batch_size = BATCH_SIZE,\n",
        "     device = device)"
      ],
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3c_UxDOVvJi0"
      },
      "source": [
        "## Main Architecture:\n",
        "\n",
        "We need:\n",
        "\n",
        "1. Encoder\n",
        "2. Decoder\n",
        "3. Seq2seq\n",
        "4. EncoderLayer\n",
        "5. DecoderLayer\n",
        "6. MultiheadAttentionComponent(with mask)\n",
        "7. FeedForwardComponent\n",
        "8. PositionalEncodingsComponent \n",
        "\n",
        "![](https://raw.githubusercontent.com/bentrevett/pytorch-seq2seq/9479fcb532214ad26fd4bda9fcf081a05e1aaf4e/assets/transformer1.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tni9vPq_ChXE"
      },
      "source": [
        "#### Positional Encoding\n",
        "\n",
        "Positional information is given using sine and cosine functions of different frequencies.\n",
        "\n",
        "$PE_{(pos,2i)} = sin(pos / 10000^{2i/d_{\\text{model}}})$\n",
        "$PE_{(pos,2i+1)} = cos(pos / 10000^{2i/d_{\\text{model}}})$\n",
        "\n",
        "where i lies between 0 to (hid_dim or d_model)//2 and pos refers to the position of the token.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AAsRHQcQwI2X"
      },
      "source": [
        "class PositionalEncodingComponent(nn.Module):\n",
        "  '''\n",
        "  Class to encode positional information to tokens.\n",
        "  \n",
        "\n",
        "  '''\n",
        "  def __init__(self,hid_dim,device,dropout=0.2,max_len=5000):\n",
        "    super().__init__()\n",
        "\n",
        "    assert hid_dim%2==0 # If not, it will result error in allocation to positional_encodings[:,1::2] later\n",
        "\n",
        "    self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "    self.positional_encodings = torch.zeros(max_len,hid_dim)\n",
        "\n",
        "    pos = torch.arange(0,max_len).unsqueeze(1) # pos : [max_len,1]\n",
        "    div_term  = torch.exp(-torch.arange(0,hid_dim,2)*math.log(10000.0)/hid_dim) # Calculating value of 1/(10000^(2i/hid_dim)) in log space and then exponentiating it\n",
        "    # div_term: [hid_dim//2]\n",
        "\n",
        "    self.positional_encodings[:,0::2] = torch.sin(pos*div_term) # pos*div_term [max_len,hid_dim//2]\n",
        "    self.positional_encodings[:,1::2] = torch.cos(pos*div_term) \n",
        "\n",
        "    self.positional_encodings = self.positional_encodings.unsqueeze(0) # To account for batch_size in inputs\n",
        "\n",
        "    self.device = device\n",
        "\n",
        "  def forward(self,x):\n",
        "    x = x + self.positional_encodings[:,:x.size(1)].detach().to(self.device)\n",
        "    return self.dropout(x)\n"
      ],
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NWe37-XCDRhA"
      },
      "source": [
        "#### Pointwise Feed Forward:\n",
        "$\\mathrm{FFN}(x)=\\max(0, xW_1 + b_1) W_2 + b_2$"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "mVh84zjDXr1_"
      },
      "source": [
        "class FeedForwardComponent(nn.Module):\n",
        "  '''\n",
        "  Class for pointwise feed forward connections\n",
        "  '''\n",
        "  def __init__(self,hid_dim,pf_dim,dropout):\n",
        "    super().__init__()\n",
        "\n",
        "    self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "    self.fc1 = nn.Linear(hid_dim,pf_dim)\n",
        "    self.fc2 = nn.Linear(pf_dim,hid_dim)\n",
        "\n",
        "  def forward(self,x):\n",
        "\n",
        "    # x : [batch_size,seq_len,hid_dim]\n",
        "    x = self.dropout(torch.relu(self.fc1(x)))\n",
        "\n",
        "    # x : [batch_size,seq_len,pf_dim]\n",
        "    x = self.fc2(x)\n",
        "\n",
        "    # x : [batch_size,seq_len,hid_dim]\n",
        "    return x"
      ],
      "execution_count": 12,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CVin9xCGDmwZ"
      },
      "source": [
        "#### Attention\n",
        "In transformers, we use self-attention i.e. we using the self value, we learn what parts of it are more important. [Remember for attention](https://akashe.io/blog/2020/12/03/rnn-lstm-gru-and-attention/#Attention), to get attention of $x$ over $y$ we find relative importance $\\alpha_{x,y}$ using a score function and later use $\\alpha_{x,y}$ to get relative parts from $y$.\n",
        "\n",
        "\n",
        "In self attention, x and y are the same. Now some nomenclature,\n",
        "\n",
        "1. Query: We find attention over query. So its similar to $y$ above.\n",
        "2. Key: What we use to find attention over query. Similar to $x$ above.\n",
        "3. Value: What we use to create a final vector using attention values. Similar to $y$ in the expression $\\sum \\alpha_{x,y}y$.\n",
        "\n",
        "\n",
        "In transformers, ~~Query, Key and Value are the same vectors~~. In implementation, the Q,K,V representations are learned using linear transformations but the input are same to the transformations.\n",
        "\n",
        "$\\mathrm{Attention}(Q, K, V) = \\mathrm{softmax}(\\frac{QK^T}{\\sqrt{d_k}})V$\n",
        "\n",
        "\n",
        "\n",
        "#### Multiheaded Attention\n",
        " Instead of calculating attention over the entire src/trg vector, we divide the src/trg vector into multiple smaller heads. We transform them to seperate vectors using learnable matrices($W_i^Q,W_i^K,W_i^V$). Perform self-attention over these transformations and concat them later. Perform another transformation($W^O$) to get the final form.\n",
        "\n",
        " $\\mathrm{MultiHead}(Q, K, V) = \\mathrm{Concat}(\\mathrm{head_1}, ...,\n",
        "\\mathrm{head_h})W^O    \\\\\n",
        "    \\text{where}~\\mathrm{head_i} = \\mathrm{Attention}(QW^Q_i, KW^K_i, VW^V_i)$\n",
        "\n",
        "\n",
        "![](https://raw.githubusercontent.com/bentrevett/pytorch-seq2seq/9479fcb532214ad26fd4bda9fcf081a05e1aaf4e/assets/transformer-attention.png)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jziMEQBXik1n"
      },
      "source": [
        "class MultiHeadedAttentionComponent(nn.Module):\n",
        "  '''\n",
        "  Multiheaded attention Component. This implementation also supports mask. \n",
        "  The reason for mask that in Decoder, we don't want attention mechanism to get\n",
        "  important information from future tokens.\n",
        "  '''\n",
        "  def __init__(self,hid_dim, n_heads, dropout, device):\n",
        "    super().__init__()\n",
        "\n",
        "    assert hid_dim % n_heads == 0 # Since we split hid_dims into n_heads\n",
        "\n",
        "    self.hid_dim = hid_dim\n",
        "    self.n_heads = n_heads # no of heads in 'multiheaded' attention\n",
        "    self.head_dim = hid_dim//n_heads # dims of each head\n",
        "\n",
        "    # Transformation from source vector to query vector\n",
        "    self.fc_q = nn.Linear(hid_dim,hid_dim)\n",
        "\n",
        "    # Transformation from source vector to key vector\n",
        "    self.fc_k = nn.Linear(hid_dim,hid_dim)\n",
        "\n",
        "    # Transformation from source vector to value vector\n",
        "    self.fc_v = nn.Linear(hid_dim,hid_dim)\n",
        "\n",
        "    self.fc_o = nn.Linear(hid_dim,hid_dim)\n",
        "\n",
        "    self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "    # Used in self attention for smoother gradients\n",
        "    self.scale = torch.sqrt(torch.FloatTensor([self.head_dim])).to(device)\n",
        "\n",
        "  def forward(self,query,key,value,mask=None):\n",
        "\n",
        "    #query : [batch_size, query_len, hid_dim]\n",
        "    #key : [batch_size, key_len, hid_dim]\n",
        "    #value : [batch_size, value_len, hid_dim]\n",
        "\n",
        "    batch_size = query.shape[0]\n",
        "\n",
        "    # Transforming quey,key,values\n",
        "    Q = self.fc_q(query)\n",
        "    K = self.fc_k(key)\n",
        "    V = self.fc_v(value)\n",
        "\n",
        "    #Q : [batch_size, query_len, hid_dim]\n",
        "    #K : [batch_size, key_len, hid_dim]\n",
        "    #V : [batch_size, value_len,hid_dim]\n",
        "\n",
        "    # Changing shapes to acocmadate n_heads information\n",
        "    Q = Q.view(batch_size, -1, self.n_heads, self.head_dim).permute(0, 2, 1, 3)\n",
        "    K = K.view(batch_size, -1, self.n_heads, self.head_dim).permute(0, 2, 1, 3)\n",
        "    V = V.view(batch_size, -1, self.n_heads, self.head_dim).permute(0, 2, 1, 3)\n",
        "\n",
        "    #Q : [batch_size, n_heads, query_len, head_dim]\n",
        "    #K : [batch_size, n_heads, key_len, head_dim]\n",
        "    #V : [batch_size, n_heads, value_len, head_dim]\n",
        "\n",
        "    # Calculating alpha\n",
        "    score = torch.matmul(Q,K.permute(0,1,3,2))/self.scale\n",
        "    # score : [batch_size, n_heads, query_len, key_len]\n",
        "\n",
        "    if mask is not None:\n",
        "      score = score.masked_fill(mask==0,-1e10)\n",
        "\n",
        "    alpha = torch.softmax(score,dim=-1)\n",
        "    # alpha : [batch_size, n_heads, query_len, key_len]\n",
        "\n",
        "    # Get the final self-attention  vector\n",
        "    x = torch.matmul(self.dropout(alpha),V)\n",
        "    # x : [batch_size, n_heads, query_len, head_dim]\n",
        "\n",
        "    # Reshaping self attention vector to concatenate\n",
        "    x = x.permute(0,2,1,3).contiguous()\n",
        "    # x : [batch_size, query_len, n_heads, head_dim]\n",
        "\n",
        "    x = x.view(batch_size,-1,self.hid_dim)\n",
        "    # x: [batch_size, query_len, hid_dim]\n",
        "\n",
        "    # Transforming concatenated outputs \n",
        "    x = self.fc_o(x)\n",
        "    #x : [batch_size, query_len, hid_dim] \n",
        "\n",
        "    return x, alpha"
      ],
      "execution_count": 13,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "toWji7OxRyaS"
      },
      "source": [
        "class EncoderLayer(nn.Module):\n",
        "  '''\n",
        "  Operations of a single layer in an Encoder. An Encoder employs multiple such layers. Each layer contains:\n",
        "  1) multihead attention, folllowed by\n",
        "  2) LayerNorm of addition of multihead attention output and input to the layer, followed by\n",
        "  3) FeedForward connections, followed by\n",
        "  4) LayerNorm of addition of FeedForward outputs and output of previous layerNorm.\n",
        "  '''\n",
        "  def __init__(self, hid_dim,n_heads,pf_dim,dropout,device):\n",
        "    super().__init__()\n",
        "    \n",
        "    self.self_attn_layer_norm = nn. LayerNorm(hid_dim) #Layer norm after self-attention\n",
        "    self.ff_layer_norm = nn.LayerNorm(hid_dim) # Layer norm after FeedForward component\n",
        "\n",
        "    self.self_attention = MultiHeadedAttentionComponent(hid_dim,n_heads,dropout,device)\n",
        "    self.feed_forward = FeedForwardComponent(hid_dim,pf_dim,dropout)\n",
        "\n",
        "    self.dropout = nn.Dropout(dropout)\n",
        "    \n",
        "  def forward(self,src,src_mask):\n",
        "    \n",
        "    # src : [batch_size, src_len, hid_dim]\n",
        "    # src_mask : [batch_size, 1, 1, src_len]\n",
        "\n",
        "    # get self-attention\n",
        "    _src, _ = self.self_attention(src,src,src,src_mask)\n",
        "\n",
        "    # LayerNorm after dropout\n",
        "    src = self.self_attn_layer_norm(src + self.dropout(_src))\n",
        "    # src : [batch_size, src_len, hid_dim]\n",
        "\n",
        "    # FeedForward\n",
        "    _src = self.feed_forward(src)\n",
        "\n",
        "    # layerNorm after dropout\n",
        "    src = self.ff_layer_norm(src + self.dropout(_src))\n",
        "    # src: [batch_size, src_len, hid_dim]\n",
        "\n",
        "    return src\n",
        "    "
      ],
      "execution_count": 14,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5haBy4QbUfRq"
      },
      "source": [
        "class DecoderLayer(nn.Module):\n",
        "  '''\n",
        "  Operations of a single layer in an Decoder. An Decoder employs multiple such layers. Each layer contains:\n",
        "  1) masked decoder self attention, followed by\n",
        "  2) LayerNorm of addition of previous attention output and input to the layer,, followed by\n",
        "  3) encoder self attention, followed by\n",
        "  4) LayerNorm of addition of result of encoder self attention and its input, followed by\n",
        "  5) FeedForward connections, followed by\n",
        "  6) LayerNorm of addition of Feedforward results and its input.\n",
        "  '''\n",
        "  def __init__(self,hid_dim,n_heads,pf_dim,dropout,device):\n",
        "    super().__init__()\n",
        "\n",
        "    self.self_attn_layer_norm = nn.LayerNorm(hid_dim)\n",
        "    self.enc_attn_layer_norm = nn.LayerNorm(hid_dim)\n",
        "    self.ff_layer_norm = nn.LayerNorm(hid_dim)\n",
        "\n",
        "    # decoder self attention\n",
        "    self.self_attention = MultiHeadedAttentionComponent(hid_dim,n_heads,dropout,device)\n",
        "\n",
        "    # encoder attention\n",
        "    self.encoder_attention = MultiHeadedAttentionComponent(hid_dim,n_heads,dropout,device)\n",
        "\n",
        "    # FeedForward\n",
        "    self.feed_forward = FeedForwardComponent(hid_dim,pf_dim,dropout)\n",
        "\n",
        "    self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "  def forward(self,trg, enc_src,trg_mask,src_mask):\n",
        "\n",
        "    #trg : [batch_size, trg_len, hid_dim]\n",
        "    #enc_src : [batch_size, src_len, hid_dim]\n",
        "    #trg_mask : [batch_size, 1, trg_len, trg_len]\n",
        "    #src_mask : [batch_size, 1, 1, src_len]\n",
        "\n",
        "    '''\n",
        "    Decoder self-attention\n",
        "    trg_mask is to force decoder to look only into past tokens and not get information from future tokens.\n",
        "    Since we apply mask before doing softmax, the final self attention vector gets no information from future tokens.\n",
        "    '''\n",
        "    _trg, _ = self.self_attention(trg,trg,trg,trg_mask)\n",
        "\n",
        "    # LayerNorm and dropout with resdiual connection\n",
        "    trg = self.self_attn_layer_norm(trg + self.dropout(_trg))\n",
        "    # trg : [batch_size, trg_len, hid_dim]\n",
        "\n",
        "    '''\n",
        "    Encoder attention:\n",
        "    Query: trg\n",
        "    key: enc_src\n",
        "    Value : enc_src\n",
        "    Why? \n",
        "    the idea here is to extract information from encoder outputs. So we use decoder self-attention as a query to find important values from enc_src\n",
        "    and that is why we use src_mask, to avoid getting information from enc_src positions where it is equal to pad-id\n",
        "    After we get necessary infromation from encoder outputs we add them back to decoder self-attention.\n",
        "    '''\n",
        "    _trg, encoder_attn_alpha = self.encoder_attention(trg,enc_src,enc_src,src_mask)\n",
        "\n",
        "    # LayerNorm , residual connection and dropout\n",
        "    trg = self.enc_attn_layer_norm(trg + self.dropout(_trg))\n",
        "    # trg : [ batch_size, trg_len, hid_dim]\n",
        "\n",
        "    # Feed Forward\n",
        "    _trg = self.feed_forward(trg)\n",
        "\n",
        "    # LayerNorm, residual connection and dropout\n",
        "    trg = self.ff_layer_norm(trg + self.dropout(_trg))\n",
        "\n",
        "    return trg, encoder_attn_alpha\n",
        "    "
      ],
      "execution_count": 15,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6RntgytMbM33"
      },
      "source": [
        "class Encoder(nn.Module):\n",
        "  '''\n",
        "  An encoder, creates token embeddings and position embeddings and passes them through multiple encoder layers\n",
        "  '''\n",
        "  def __init__(self,input_dim,hid_dim,n_layers,n_heads,pf_dim,dropout,device,max_length = 5000):\n",
        "    super().__init__()\n",
        "    self.device = device\n",
        "\n",
        "    self.tok_embedding = nn.Embedding(input_dim,hid_dim)\n",
        "    self.pos_embedding = PositionalEncodingComponent(hid_dim,device,dropout,max_length)\n",
        "\n",
        "    # encoder layers\n",
        "    self.layers = nn.ModuleList([EncoderLayer(hid_dim,n_heads,pf_dim,dropout,device) for _ in range(n_layers)])\n",
        "\n",
        "    self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "    self.scale = torch.sqrt(torch.FloatTensor([hid_dim])).to(device)\n",
        "\n",
        "  def forward(self,src,src_mask):\n",
        "\n",
        "    # src : [batch_size, src_len]\n",
        "    # src_mask : [batch_size,1,1,src_len]\n",
        "\n",
        "    batch_size = src.shape[0]\n",
        "    src_len = src.shape[1]\n",
        "\n",
        "    tok_embeddings = self.tok_embedding(src)*self.scale\n",
        "\n",
        "    # token plus position embeddings\n",
        "    src  = self.pos_embedding(tok_embeddings)\n",
        "\n",
        "    for layer in self.layers:\n",
        "      src = layer(src,src_mask)\n",
        "    # src : [batch_size, src_len, hid_dim]\n",
        "\n",
        "    return src"
      ],
      "execution_count": 16,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "j2HncgZfdwQk"
      },
      "source": [
        "class Decoder(nn.Module):\n",
        "  '''\n",
        "  An decoder, creates token embeddings and position embeddings and passes them through multiple decoder layers\n",
        "  '''\n",
        "  def __init__(self,output_dim,hid_dim,n_layers,n_heads,pf_dim,dropout,device,max_length= 5000):\n",
        "    super().__init__()\n",
        "\n",
        "    self.device = device\n",
        "\n",
        "    self.tok_embedding = nn.Embedding(output_dim,hid_dim)\n",
        "    self.pos_embedding = PositionalEncodingComponent(hid_dim,device,dropout,max_length)\n",
        "\n",
        "    # decoder layers\n",
        "    self.layers = nn.ModuleList([DecoderLayer(hid_dim,n_heads,pf_dim,dropout,device) for _ in range(n_layers)])\n",
        "\n",
        "    # convert decoder outputs to real outputs\n",
        "    self.fc_out = nn.Linear(hid_dim,output_dim)\n",
        "\n",
        "    self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "    self.scale = torch.sqrt(torch.FloatTensor([hid_dim])).to(device)\n",
        "\n",
        "  def forward(self, trg, enc_src,trg_mask,src_mask):\n",
        "    \n",
        "    #trg : [batch_size, trg_len]\n",
        "    #enc_src : [batch_size, src_len, hid_dim]\n",
        "    #trg_mask : [batch_size, 1, trg_len, trg_len]\n",
        "    #src_mask : [batch_size, 1, 1, src_len]\n",
        "\n",
        "    batch_size = trg.shape[0]\n",
        "    trg_len = trg.shape[1]\n",
        "\n",
        "    tok_embeddings = self.tok_embedding(trg)*self.scale\n",
        "\n",
        "    # token plus pos embeddings\n",
        "    trg = self.pos_embedding(tok_embeddings)\n",
        "    # trg : [batch_size, trg_len, hid_dim]\n",
        "\n",
        "    # Pass trg thorugh decoder layers\n",
        "    for layer in self.layers:\n",
        "      trg, encoder_attention = layer(trg,enc_src,trg_mask,src_mask)\n",
        "    \n",
        "    # trg : [batch_size,trg_len,hid_dim]\n",
        "    # encoder_attention :  [batch_size, n_head,trg_len, src_len]\n",
        "\n",
        "    # Convert to outputs\n",
        "    output = self.fc_out(trg)\n",
        "    # output : [batch_size, trg_len, output_dim]\n",
        "    \n",
        "    return output, encoder_attention"
      ],
      "execution_count": 17,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "l3TOCxJwgeib"
      },
      "source": [
        "class Seq2Seq(nn.Module):\n",
        "  def __init__(self, encoder, decoder, src_pad_idx, trg_pad_idx, device):\n",
        "    super().__init__()\n",
        "    self.encoder = encoder\n",
        "    self.decoder = decoder\n",
        "    self.src_pad_idx = src_pad_idx\n",
        "    self.trg_pad_idx = trg_pad_idx\n",
        "    self.device = device\n",
        "\n",
        "  def make_src_mask(self,src):\n",
        "    # src : [batch_size, src_len]\n",
        "\n",
        "    # Masking pad values\n",
        "    src_mask = (src != self.src_pad_idx).unsqueeze(1).unsqueeze(2)\n",
        "    # src_mask : [batch_size,1,1,src_len]\n",
        "\n",
        "    return src_mask\n",
        "\n",
        "  def make_trg_mask(self,trg):\n",
        "    # trg : [batch_size, trg_len]\n",
        "\n",
        "    # Masking pad values\n",
        "    trg_pad_mask = (trg != self.trg_pad_idx).unsqueeze(1).unsqueeze(2)\n",
        "    # trg_pad_mask : [batch_size,1,1, trg_len]\n",
        "\n",
        "    # Masking future values\n",
        "    trg_len = trg.shape[1]\n",
        "    trg_sub_mask = torch.tril(torch.ones((trg_len,trg_len),device= self.device)).bool()\n",
        "    # trg_sub_mask : [trg_len, trg_len]\n",
        "\n",
        "    # combine both masks\n",
        "    trg_mask = trg_pad_mask & trg_sub_mask\n",
        "    # trg_mask = [batch_size,1,trg_len,trg_len]\n",
        "\n",
        "    return trg_mask\n",
        "\n",
        "  def forward(self,src,trg):\n",
        "\n",
        "    # src : [batch_size, src_len]\n",
        "    # trg : [batch_size, trg_len]\n",
        "\n",
        "    src_mask = self.make_src_mask(src)\n",
        "    trg_mask = self.make_trg_mask(trg)\n",
        "\n",
        "    # src_mask : [ batch_size, 1,1,src_len]\n",
        "    # trg_mask : [batch_size, 1, trg_len, trg_len]\n",
        "\n",
        "    enc_src = self.encoder(src,src_mask)\n",
        "    #enc_src : [batch_size, src_len, hid_dim]\n",
        "\n",
        "    output, encoder_decoder_attention = self.decoder(trg,enc_src,trg_mask,src_mask)\n",
        "    # output : [batch_size, trg_len, output_dim]\n",
        "    # encoder_decoder_attention : [batch_size, n_heads, trg_len, src_len]\n",
        "\n",
        "    return output, encoder_decoder_attention"
      ],
      "execution_count": 18,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ubtuSyspjqZq"
      },
      "source": [
        "Intializing network"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rE4-1aHajv_0"
      },
      "source": [
        "INPUT_DIM = len(SRC.vocab)\n",
        "OUTPUT_DIM = len(TRG.vocab)\n",
        "HID_DIM = 256\n",
        "ENC_LAYERS = 3\n",
        "DEC_LAYERS = 3\n",
        "ENC_HEADS = 8\n",
        "DEC_HEADS = 8\n",
        "ENC_PF_DIM = 512\n",
        "DEC_PF_DIM = 512\n",
        "ENC_DROPOUT = 0.1\n",
        "DEC_DROPOUT = 0.1\n",
        "\n",
        "enc = Encoder(INPUT_DIM, \n",
        "              HID_DIM, \n",
        "              ENC_LAYERS, \n",
        "              ENC_HEADS, \n",
        "              ENC_PF_DIM, \n",
        "              ENC_DROPOUT, \n",
        "              device)\n",
        "\n",
        "dec = Decoder(OUTPUT_DIM, \n",
        "              HID_DIM, \n",
        "              DEC_LAYERS, \n",
        "              DEC_HEADS, \n",
        "              DEC_PF_DIM, \n",
        "              DEC_DROPOUT, \n",
        "              device)\n",
        "\n",
        "SRC_PAD_IDX = SRC.vocab.stoi[SRC.pad_token]\n",
        "TRG_PAD_IDX = TRG.vocab.stoi[TRG.pad_token]\n",
        "\n",
        "model = Seq2Seq(enc, dec, SRC_PAD_IDX, TRG_PAD_IDX, device).to(device)"
      ],
      "execution_count": 19,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0KezegeUj5QQ"
      },
      "source": [
        "Initialize weights"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "8qdQ3lffj7TX"
      },
      "source": [
        "def initialize_weights(m):\n",
        "    if hasattr(m, 'weight') and m.weight.dim() > 1:\n",
        "        nn.init.xavier_uniform_(m.weight.data)\n",
        "\n",
        "model.apply(initialize_weights);"
      ],
      "execution_count": 20,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1MCldGtbkBac"
      },
      "source": [
        "Total model params"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ZIcfEEPKkDwc",
        "outputId": "51c2084c-47d9-48e1-f4e4-a3c57f4b81c4"
      },
      "source": [
        "def count_parameters(model):\n",
        "    return sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
        "\n",
        "print(f'The model has {count_parameters(model):,} trainable parameters')"
      ],
      "execution_count": 21,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "The model has 8,987,653 trainable parameters\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eBMqUOSNkHAU"
      },
      "source": [
        "Learning rate, criterion and optimizer"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cKANsYKDkMDu"
      },
      "source": [
        "LEARNING_RATE = 0.0005\n",
        "\n",
        "optimizer = torch.optim.Adam(model.parameters(), lr = LEARNING_RATE)\n",
        "\n",
        "criterion = nn.CrossEntropyLoss(ignore_index = TRG_PAD_IDX)"
      ],
      "execution_count": 22,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wDjw8jqikPOt"
      },
      "source": [
        "Train Loop"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3LIDVwh2kSCI"
      },
      "source": [
        "def train(model, iterator, optimizer, criterion, clip):\n",
        "    \n",
        "    model.train()\n",
        "    \n",
        "    epoch_loss = 0\n",
        "    \n",
        "    for i, batch in enumerate(iterator):\n",
        "        \n",
        "        src = batch.src\n",
        "        trg = batch.trg\n",
        "        \n",
        "        optimizer.zero_grad()\n",
        "        \n",
        "        output, _ = model(src, trg[:,:-1])\n",
        "                \n",
        "        #output = [batch size, trg len - 1, output dim]\n",
        "        #trg = [batch size, trg len]\n",
        "            \n",
        "        output_dim = output.shape[-1]\n",
        "            \n",
        "        output = output.contiguous().view(-1, output_dim)\n",
        "        trg = trg[:,1:].contiguous().view(-1)\n",
        "                \n",
        "        #output = [batch size * trg len - 1, output dim]\n",
        "        #trg = [batch size * trg len - 1]\n",
        "            \n",
        "        loss = criterion(output, trg)\n",
        "        \n",
        "        loss.backward()\n",
        "        \n",
        "        torch.nn.utils.clip_grad_norm_(model.parameters(), clip)\n",
        "        \n",
        "        optimizer.step()\n",
        "        \n",
        "        epoch_loss += loss.item()\n",
        "        \n",
        "    return epoch_loss / len(iterator)"
      ],
      "execution_count": 23,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gjs90xwlkSnN"
      },
      "source": [
        "Evaluate Loop"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9e-vPTMukUPW"
      },
      "source": [
        "def evaluate(model, iterator, criterion):\n",
        "    \n",
        "    model.eval()\n",
        "    \n",
        "    epoch_loss = 0\n",
        "    \n",
        "    with torch.no_grad():\n",
        "    \n",
        "        for i, batch in enumerate(iterator):\n",
        "\n",
        "            src = batch.src\n",
        "            trg = batch.trg\n",
        "\n",
        "            output, _ = model(src, trg[:,:-1])\n",
        "            \n",
        "            #output = [batch size, trg len - 1, output dim]\n",
        "            #trg = [batch size, trg len]\n",
        "            \n",
        "            output_dim = output.shape[-1]\n",
        "            \n",
        "            output = output.contiguous().view(-1, output_dim)\n",
        "            trg = trg[:,1:].contiguous().view(-1)\n",
        "            \n",
        "            #output = [batch size * trg len - 1, output dim]\n",
        "            #trg = [batch size * trg len - 1]\n",
        "            \n",
        "            loss = criterion(output, trg)\n",
        "\n",
        "            epoch_loss += loss.item()\n",
        "        \n",
        "    return epoch_loss / len(iterator)"
      ],
      "execution_count": 24,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BfPJIq1fkWE2"
      },
      "source": [
        "Time per epoch"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AgsdAlTwkXP4"
      },
      "source": [
        "def epoch_time(start_time, end_time):\n",
        "    elapsed_time = end_time - start_time\n",
        "    elapsed_mins = int(elapsed_time / 60)\n",
        "    elapsed_secs = int(elapsed_time - (elapsed_mins * 60))\n",
        "    return elapsed_mins, elapsed_secs"
      ],
      "execution_count": 25,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "itILwWt5kY8z"
      },
      "source": [
        "Runner Loop"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Gk8neNJCkbKJ",
        "outputId": "c0aa67b2-9f20-46e8-bce7-40312515008f"
      },
      "source": [
        "N_EPOCHS = 10\n",
        "CLIP = 1\n",
        "\n",
        "best_valid_loss = float('inf')\n",
        "\n",
        "for epoch in range(N_EPOCHS):\n",
        "    \n",
        "    start_time = time.time()\n",
        "    \n",
        "    train_loss = train(model, train_iterator, optimizer, criterion, CLIP)\n",
        "    valid_loss = evaluate(model, valid_iterator, criterion)\n",
        "    \n",
        "    end_time = time.time()\n",
        "    \n",
        "    epoch_mins, epoch_secs = epoch_time(start_time, end_time)\n",
        "    \n",
        "    if valid_loss < best_valid_loss:\n",
        "        best_valid_loss = valid_loss\n",
        "        torch.save(model.state_dict(), 'tut6-model.pt')\n",
        "    \n",
        "    print(f'Epoch: {epoch+1:02} | Time: {epoch_mins}m {epoch_secs}s')\n",
        "    print(f'\\tTrain Loss: {train_loss:.3f} | Train PPL: {math.exp(train_loss):7.3f}')\n",
        "    print(f'\\t Val. Loss: {valid_loss:.3f} |  Val. PPL: {math.exp(valid_loss):7.3f}')"
      ],
      "execution_count": 26,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Epoch: 01 | Time: 0m 18s\n",
            "\tTrain Loss: 4.484 | Train PPL:  88.607\n",
            "\t Val. Loss: 3.339 |  Val. PPL:  28.185\n",
            "Epoch: 02 | Time: 0m 18s\n",
            "\tTrain Loss: 3.150 | Train PPL:  23.338\n",
            "\t Val. Loss: 2.557 |  Val. PPL:  12.901\n",
            "Epoch: 03 | Time: 0m 19s\n",
            "\tTrain Loss: 2.460 | Train PPL:  11.709\n",
            "\t Val. Loss: 2.093 |  Val. PPL:   8.109\n",
            "Epoch: 04 | Time: 0m 19s\n",
            "\tTrain Loss: 2.036 | Train PPL:   7.663\n",
            "\t Val. Loss: 1.833 |  Val. PPL:   6.253\n",
            "Epoch: 05 | Time: 0m 19s\n",
            "\tTrain Loss: 1.754 | Train PPL:   5.778\n",
            "\t Val. Loss: 1.694 |  Val. PPL:   5.443\n",
            "Epoch: 06 | Time: 0m 20s\n",
            "\tTrain Loss: 1.554 | Train PPL:   4.730\n",
            "\t Val. Loss: 1.645 |  Val. PPL:   5.179\n",
            "Epoch: 07 | Time: 0m 20s\n",
            "\tTrain Loss: 1.398 | Train PPL:   4.046\n",
            "\t Val. Loss: 1.591 |  Val. PPL:   4.908\n",
            "Epoch: 08 | Time: 0m 20s\n",
            "\tTrain Loss: 1.271 | Train PPL:   3.565\n",
            "\t Val. Loss: 1.557 |  Val. PPL:   4.746\n",
            "Epoch: 09 | Time: 0m 19s\n",
            "\tTrain Loss: 1.168 | Train PPL:   3.214\n",
            "\t Val. Loss: 1.563 |  Val. PPL:   4.773\n",
            "Epoch: 10 | Time: 0m 19s\n",
            "\tTrain Loss: 1.076 | Train PPL:   2.934\n",
            "\t Val. Loss: 1.548 |  Val. PPL:   4.701\n"
          ],
          "name": "stdout"
        }
      ]
    }
  ]
}