Attention_is_All_You_Need.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "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.5"
    },
    "colab": {
      "name": "Copy of Copy of 6 - Attention is All You Need.ipynb",
      "provenance": [],
      "collapsed_sections": [
        "3WKffF0Gy_x6",
        "q12JMc8by_x7",
        "3XbLg8mty_x7",
        "CCX5MMGky_x8"
      ],
      "toc_visible": true
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sOd68EzV2wcI"
      },
      "source": [
        "# Huggingface Transformers intro"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "p3dmhxc_3nXu",
        "outputId": "f55fe066-0b12-457a-bb1f-d155da34c142"
      },
      "source": [
        "!pip install transformers"
      ],
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Collecting transformers\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/d5/43/cfe4ee779bbd6a678ac6a97c5a5cdeb03c35f9eaebbb9720b036680f9a2d/transformers-4.6.1-py3-none-any.whl (2.2MB)\n",
            "\u001b[K     |████████████████████████████████| 2.3MB 7.4MB/s \n",
            "\u001b[?25hRequirement already satisfied: packaging in /usr/local/lib/python3.7/dist-packages (from transformers) (20.9)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.7/dist-packages (from transformers) (3.0.12)\n",
            "Collecting tokenizers<0.11,>=0.10.1\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/d4/e2/df3543e8ffdab68f5acc73f613de9c2b155ac47f162e725dcac87c521c11/tokenizers-0.10.3-cp37-cp37m-manylinux_2_5_x86_64.manylinux1_x86_64.manylinux_2_12_x86_64.manylinux2010_x86_64.whl (3.3MB)\n",
            "\u001b[K     |████████████████████████████████| 3.3MB 47.5MB/s \n",
            "\u001b[?25hCollecting huggingface-hub==0.0.8\n",
            "  Downloading https://files.pythonhosted.org/packages/a1/88/7b1e45720ecf59c6c6737ff332f41c955963090a18e72acbcbeac6b25e86/huggingface_hub-0.0.8-py3-none-any.whl\n",
            "Requirement already satisfied: tqdm>=4.27 in /usr/local/lib/python3.7/dist-packages (from transformers) (4.41.1)\n",
            "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.7/dist-packages (from transformers) (2019.12.20)\n",
            "Collecting sacremoses\n",
            "\u001b[?25l  Downloading https://files.pythonhosted.org/packages/75/ee/67241dc87f266093c533a2d4d3d69438e57d7a90abb216fa076e7d475d4a/sacremoses-0.0.45-py3-none-any.whl (895kB)\n",
            "\u001b[K     |████████████████████████████████| 901kB 50.8MB/s \n",
            "\u001b[?25hRequirement already satisfied: importlib-metadata; python_version < \"3.8\" in /usr/local/lib/python3.7/dist-packages (from transformers) (4.5.0)\n",
            "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.7/dist-packages (from transformers) (1.19.5)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.7/dist-packages (from transformers) (2.23.0)\n",
            "Requirement already satisfied: pyparsing>=2.0.2 in /usr/local/lib/python3.7/dist-packages (from packaging->transformers) (2.4.7)\n",
            "Requirement already satisfied: click in /usr/local/lib/python3.7/dist-packages (from sacremoses->transformers) (7.1.2)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.7/dist-packages (from sacremoses->transformers) (1.15.0)\n",
            "Requirement already satisfied: joblib in /usr/local/lib/python3.7/dist-packages (from sacremoses->transformers) (1.0.1)\n",
            "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.7/dist-packages (from importlib-metadata; python_version < \"3.8\"->transformers) (3.4.1)\n",
            "Requirement already satisfied: typing-extensions>=3.6.4; python_version < \"3.8\" in /usr/local/lib/python3.7/dist-packages (from importlib-metadata; python_version < \"3.8\"->transformers) (3.7.4.3)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests->transformers) (2.10)\n",
            "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests->transformers) (1.24.3)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests->transformers) (2021.5.30)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests->transformers) (3.0.4)\n",
            "Installing collected packages: tokenizers, huggingface-hub, sacremoses, transformers\n",
            "Successfully installed huggingface-hub-0.0.8 sacremoses-0.0.45 tokenizers-0.10.3 transformers-4.6.1\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Pa2vsNYi2tzK",
        "outputId": "cd0a66ae-0dab-435b-b01c-8481d0482141"
      },
      "source": [
        "from transformers import BertTokenizer, BertForMaskedLM\n",
        "from torch.nn import functional as F\n",
        "import torch\n",
        "tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')\n",
        "model = BertForMaskedLM.from_pretrained('bert-base-uncased',    return_dict = True)\n",
        "text = \"The capital of France, [MASK] contains the Eiffel Tower.\"\n",
        "\n",
        "\n",
        "input = tokenizer.encode_plus(text, return_tensors = \"pt\")\n",
        "mask_index = torch.where(input[\"input_ids\"][0] == tokenizer.mask_token_id)\n",
        "output = model(**input)\n",
        "logits = output.logits\n",
        "softmax = F.softmax(logits, dim = -1)\n",
        "mask_word = softmax[0, mask_index, :]\n",
        "top_10 = torch.topk(mask_word, 10, dim = 1)[1][0]\n",
        "\n",
        "print('\\n\\nSuggestions:')\n",
        "for token in top_10:\n",
        "   word = tokenizer.decode([token])\n",
        "   new_sentence = text.replace(tokenizer.mask_token, word)\n",
        "   print(new_sentence)"
      ],
      "execution_count": 19,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Some weights of the model checkpoint at bert-base-uncased were not used when initializing BertForMaskedLM: ['cls.seq_relationship.bias', 'cls.seq_relationship.weight']\n",
            "- This IS expected if you are initializing BertForMaskedLM from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
            "- This IS NOT expected if you are initializing BertForMaskedLM from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "\n",
            "\n",
            "Suggestions:\n",
            "The capital of France, paris contains the Eiffel Tower.\n",
            "The capital of France, it contains the Eiffel Tower.\n",
            "The capital of France, which contains the Eiffel Tower.\n",
            "The capital of France, versailles contains the Eiffel Tower.\n",
            "The capital of France, brussels contains the Eiffel Tower.\n",
            "The capital of France, monaco contains the Eiffel Tower.\n",
            "The capital of France, and contains the Eiffel Tower.\n",
            "The capital of France, orleans contains the Eiffel Tower.\n",
            "The capital of France, lyon contains the Eiffel Tower.\n",
            "The capital of France, france contains the Eiffel Tower.\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DmjYmuge3fKB"
      },
      "source": [
        ""
      ],
      "execution_count": 10,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dHAsY11Q3fVr"
      },
      "source": [
        ""
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yxdChIEky_xy"
      },
      "source": [
        "## Preparing the Data"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dCawn1l9y_xz"
      },
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "\n",
        "import torchtext\n",
        "from torchtext.legacy.datasets import Multi30k\n",
        "from torchtext.legacy.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": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4BsXoOXTy_xz"
      },
      "source": [
        "SEED = 1234\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": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "V-zSIxEWy_x0",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "9eb3f52a-3be8-4495-d638-b871d37ac750"
      },
      "source": [
        "!python -m spacy download de_core_news_sm\n",
        "!python -m spacy download en_core_web_sm\n",
        "\n",
        "import de_core_news_sm\n",
        "spacy_de = de_core_news_sm.load()\n",
        "spacy_en = spacy.load('en_core_web_sm')"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Collecting de_core_news_sm==2.2.5\n",
            "\u001b[?25l  Downloading https://github.com/explosion/spacy-models/releases/download/de_core_news_sm-2.2.5/de_core_news_sm-2.2.5.tar.gz (14.9MB)\n",
            "\u001b[K     |████████████████████████████████| 14.9MB 5.8MB/s \n",
            "\u001b[?25hRequirement already satisfied: spacy>=2.2.2 in /usr/local/lib/python3.7/dist-packages (from de_core_news_sm==2.2.5) (2.2.4)\n",
            "Requirement already satisfied: preshed<3.1.0,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from spacy>=2.2.2->de_core_news_sm==2.2.5) (3.0.5)\n",
            "Requirement already satisfied: catalogue<1.1.0,>=0.0.7 in /usr/local/lib/python3.7/dist-packages (from spacy>=2.2.2->de_core_news_sm==2.2.5) (1.0.0)\n",
            "Requirement already satisfied: requests<3.0.0,>=2.13.0 in /usr/local/lib/python3.7/dist-packages (from spacy>=2.2.2->de_core_news_sm==2.2.5) (2.23.0)\n",
            "Requirement already satisfied: cymem<2.1.0,>=2.0.2 in /usr/local/lib/python3.7/dist-packages (from spacy>=2.2.2->de_core_news_sm==2.2.5) (2.0.5)\n",
            "Requirement already satisfied: wasabi<1.1.0,>=0.4.0 in /usr/local/lib/python3.7/dist-packages (from spacy>=2.2.2->de_core_news_sm==2.2.5) (0.8.2)\n",
            "Requirement already satisfied: srsly<1.1.0,>=1.0.2 in /usr/local/lib/python3.7/dist-packages (from spacy>=2.2.2->de_core_news_sm==2.2.5) (1.0.5)\n",
            "Requirement already satisfied: plac<1.2.0,>=0.9.6 in /usr/local/lib/python3.7/dist-packages (from spacy>=2.2.2->de_core_news_sm==2.2.5) (1.1.3)\n",
            "Requirement already satisfied: blis<0.5.0,>=0.4.0 in /usr/local/lib/python3.7/dist-packages (from spacy>=2.2.2->de_core_news_sm==2.2.5) (0.4.1)\n",
            "Requirement already satisfied: murmurhash<1.1.0,>=0.28.0 in /usr/local/lib/python3.7/dist-packages (from spacy>=2.2.2->de_core_news_sm==2.2.5) (1.0.5)\n",
            "Requirement already satisfied: numpy>=1.15.0 in /usr/local/lib/python3.7/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.7/dist-packages (from spacy>=2.2.2->de_core_news_sm==2.2.5) (57.0.0)\n",
            "Requirement already satisfied: tqdm<5.0.0,>=4.38.0 in /usr/local/lib/python3.7/dist-packages (from spacy>=2.2.2->de_core_news_sm==2.2.5) (4.41.1)\n",
            "Requirement already satisfied: thinc==7.4.0 in /usr/local/lib/python3.7/dist-packages (from spacy>=2.2.2->de_core_news_sm==2.2.5) (7.4.0)\n",
            "Requirement already satisfied: importlib-metadata>=0.20; python_version < \"3.8\" in /usr/local/lib/python3.7/dist-packages (from catalogue<1.1.0,>=0.0.7->spacy>=2.2.2->de_core_news_sm==2.2.5) (4.5.0)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/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.7/dist-packages (from requests<3.0.0,>=2.13.0->spacy>=2.2.2->de_core_news_sm==2.2.5) (2021.5.30)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/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.7/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: zipp>=0.5 in /usr/local/lib/python3.7/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.1)\n",
            "Requirement already satisfied: typing-extensions>=3.6.4; python_version < \"3.8\" in /usr/local/lib/python3.7/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",
            "Building wheels for collected packages: de-core-news-sm\n",
            "  Building wheel for de-core-news-sm (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for de-core-news-sm: filename=de_core_news_sm-2.2.5-cp37-none-any.whl size=14907055 sha256=cdb065979574a22e5404ade47691dfce21af87caf602f3dc07bd7e9c593d0d56\n",
            "  Stored in directory: /tmp/pip-ephem-wheel-cache-9vxxhv5b/wheels/ba/3f/ed/d4aa8e45e7191b7f32db4bfad565e7da1edbf05c916ca7a1ca\n",
            "Successfully built de-core-news-sm\n",
            "Installing collected packages: de-core-news-sm\n",
            "Successfully installed de-core-news-sm-2.2.5\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",
            "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.7/dist-packages (2.2.5)\n",
            "Requirement already satisfied: spacy>=2.2.2 in /usr/local/lib/python3.7/dist-packages (from en_core_web_sm==2.2.5) (2.2.4)\n",
            "Requirement already satisfied: thinc==7.4.0 in /usr/local/lib/python3.7/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (7.4.0)\n",
            "Requirement already satisfied: blis<0.5.0,>=0.4.0 in /usr/local/lib/python3.7/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (0.4.1)\n",
            "Requirement already satisfied: murmurhash<1.1.0,>=0.28.0 in /usr/local/lib/python3.7/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (1.0.5)\n",
            "Requirement already satisfied: wasabi<1.1.0,>=0.4.0 in /usr/local/lib/python3.7/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.7/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.7/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (2.0.5)\n",
            "Requirement already satisfied: plac<1.2.0,>=0.9.6 in /usr/local/lib/python3.7/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (1.1.3)\n",
            "Requirement already satisfied: setuptools in /usr/local/lib/python3.7/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (57.0.0)\n",
            "Requirement already satisfied: catalogue<1.1.0,>=0.0.7 in /usr/local/lib/python3.7/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.7/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (1.19.5)\n",
            "Requirement already satisfied: srsly<1.1.0,>=1.0.2 in /usr/local/lib/python3.7/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.7/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (2.23.0)\n",
            "Requirement already satisfied: preshed<3.1.0,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from spacy>=2.2.2->en_core_web_sm==2.2.5) (3.0.5)\n",
            "Requirement already satisfied: importlib-metadata>=0.20; python_version < \"3.8\" in /usr/local/lib/python3.7/dist-packages (from catalogue<1.1.0,>=0.0.7->spacy>=2.2.2->en_core_web_sm==2.2.5) (4.5.0)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/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: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.7/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: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests<3.0.0,>=2.13.0->spacy>=2.2.2->en_core_web_sm==2.2.5) (2021.5.30)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.7/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: zipp>=0.5 in /usr/local/lib/python3.7/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.1)\n",
            "Requirement already satisfied: typing-extensions>=3.6.4; python_version < \"3.8\" in /usr/local/lib/python3.7/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",
            "\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"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kLB6GI7ry_x0"
      },
      "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": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D6y3fPJmy_x0"
      },
      "source": [
        "Our fields are the same as the previous notebook. The model expects data to be fed in with the batch dimension first, so we use `batch_first = True`. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "yXCOKl5Fy_x1"
      },
      "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": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QKvxT6bLy_x2"
      },
      "source": [
        "train_data, valid_data, test_data = Multi30k.splits(exts = ('.de', '.en'), \n",
        "                                                    fields = (SRC, TRG))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "T3MFok3Ly_x2"
      },
      "source": [
        "SRC.build_vocab(train_data, min_freq = 2)\n",
        "TRG.build_vocab(train_data, min_freq = 2)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jZ1XMthqy_x3"
      },
      "source": [
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "wUXHZuuCy_x3"
      },
      "source": [
        "BATCH_SIZE = 32\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": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c-bHe9BUy_x3"
      },
      "source": [
        "## Building the Model\n",
        "\n",
        "### Encoder\n",
        "\n",
        "Similar to the ConvSeq2Seq model, the Transformer's encoder does not attempt to compress the entire source sentence, $X = (x_1, ... ,x_n)$, into a single context vector, $z$. Instead it produces a sequence of context vectors, $Z = (z_1, ... , z_n)$. So, if our input sequence was 5 tokens long we would have $Z = (z_1, z_2, z_3, z_4, z_5)$. Why do we call this a sequence of context vectors and not a sequence of hidden states? A hidden state at time $t$ in an RNN has only seen tokens $x_t$ and all the tokens before it. However, each context vector here has seen all tokens at all positions within the input sequence.\n",
        "\n",
        "![](https://github.com/bentrevett/pytorch-seq2seq/blob/master/assets/transformer-encoder.png?raw=1)\n",
        "\n",
        "First, the tokens are passed through a standard embedding layer. Next, as the model has no recurrent it has no idea about the order of the tokens within the sequence. We solve this by using a second embedding layer called a *positional embedding layer*. This is a standard embedding layer where the input is not the token itself but the position of the token within the sequence, starting with the first token, the `<sos>` (start of sequence) token, in position 0. The position embedding has a \"vocabulary\" size of 100, which means our model can accept sentences up to 100 tokens long. This can be increased if we want to handle longer sentences.\n",
        "\n",
        "The original Transformer implementation from the Attention is All You Need paper does not learn positional embeddings. Instead it uses a fixed static embedding. Modern Transformer architectures, like BERT, use positional embeddings instead, hence we have decided to use them in these tutorials. Check out [this](http://nlp.seas.harvard.edu/2018/04/03/attention.html#positional-encoding) section to read more about the positional embeddings used in the original Transformer model.\n",
        "\n",
        "Next, the token and positional embeddings are elementwise summed together to get a vector which contains information about the token and also its position with in the sequence. However, before they are summed, the token embeddings are multiplied by a scaling factor which is $\\sqrt{d_{model}}$, where $d_{model}$ is the hidden dimension size, `hid_dim`. This supposedly reduces variance in the embeddings and the model is difficult to train reliably without this scaling factor. Dropout is then applied to the combined embeddings.\n",
        "\n",
        "The combined embeddings are then passed through $N$ *encoder layers* to get $Z$, which is then output and can be used by the decoder.\n",
        "\n",
        "The source mask, `src_mask`, is simply the same shape as the source sentence but has a value of 1 when the token in the source sentence is not a `<pad>` token and 0 when it is a `<pad>` token. This is used in the encoder layers to mask the multi-head attention mechanisms, which are used to calculate and apply attention over the source sentence, so the model does not pay attention to `<pad>` tokens, which contain no useful information."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "5qMAHun7y_x4"
      },
      "source": [
        "class Encoder(nn.Module):\n",
        "    def __init__(self, \n",
        "                 input_dim, \n",
        "                 hid_dim, \n",
        "                 n_layers, \n",
        "                 n_heads, \n",
        "                 pf_dim,\n",
        "                 dropout, \n",
        "                 device,\n",
        "                 max_length = 100):\n",
        "        super().__init__()\n",
        "\n",
        "        self.device = device\n",
        "        \n",
        "        self.tok_embedding = nn.Embedding(input_dim, hid_dim)\n",
        "        self.pos_embedding = nn.Embedding(max_length, hid_dim)\n",
        "        \n",
        "        self.layers = nn.ModuleList([EncoderLayer(hid_dim, \n",
        "                                                  n_heads, \n",
        "                                                  pf_dim,\n",
        "                                                  dropout, \n",
        "                                                  device) \n",
        "                                     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",
        "        pos = torch.arange(0, src_len).unsqueeze(0).repeat(batch_size, 1).to(self.device)\n",
        "        \n",
        "        #pos = [batch size, src len]\n",
        "        \n",
        "        src = self.dropout((self.tok_embedding(src) * self.scale) + self.pos_embedding(pos))\n",
        "        \n",
        "        #src = [batch size, src len, hid dim]\n",
        "        \n",
        "        for layer in self.layers:\n",
        "            src = layer(src, src_mask)\n",
        "            \n",
        "        #src = [batch size, src len, hid dim]\n",
        "            \n",
        "        return src"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BZ0--Yd2y_x5"
      },
      "source": [
        "### Encoder Layer\n",
        "\n",
        "The encoder layers are where all of the \"meat\" of the encoder is contained. We first pass the source sentence and its mask into the *multi-head attention layer*, then perform dropout on it, apply a residual connection and pass it through a [Layer Normalization](https://arxiv.org/abs/1607.06450) layer. We then pass it through a *position-wise feedforward* layer and then, again, apply dropout, a residual connection and then layer normalization to get the output of this layer which is fed into the next layer. The parameters are not shared between layers. \n",
        "\n",
        "The mutli head attention layer is used by the encoder layer to attend to the source sentence, i.e. it is calculating and applying attention over itself instead of another sequence, hence we call it *self attention*.\n",
        "\n",
        "[This](https://mlexplained.com/2018/01/13/weight-normalization-and-layer-normalization-explained-normalization-in-deep-learning-part-2/) article goes into more detail about layer normalization, but the gist is that it normalizes the values of the features, i.e. across the hidden dimension, so each feature has a mean of 0 and a standard deviation of 1. This allows neural networks with a larger number of layers, like the Transformer, to be trained easier."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HMJlQoNCy_x6"
      },
      "source": [
        "class EncoderLayer(nn.Module):\n",
        "    def __init__(self, \n",
        "                 hid_dim, \n",
        "                 n_heads, \n",
        "                 pf_dim,  \n",
        "                 dropout, \n",
        "                 device):\n",
        "        super().__init__()\n",
        "        \n",
        "        self.self_attn_layer_norm = nn.LayerNorm(hid_dim)\n",
        "        self.ff_layer_norm = nn.LayerNorm(hid_dim)\n",
        "        self.self_attention = MultiHeadAttentionLayer(hid_dim, n_heads, dropout, device)\n",
        "        self.positionwise_feedforward = PositionwiseFeedforwardLayer(hid_dim, \n",
        "                                                                     pf_dim, \n",
        "                                                                     dropout)\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",
        "        #self attention\n",
        "        _src, _ = self.self_attention(src, src, src, src_mask)\n",
        "        \n",
        "        #dropout, residual connection and layer norm\n",
        "        src = self.self_attn_layer_norm(src + self.dropout(_src))\n",
        "        \n",
        "        #src = [batch size, src len, hid dim]\n",
        "        \n",
        "        #positionwise feedforward\n",
        "        _src = self.positionwise_feedforward(src)\n",
        "        \n",
        "        #dropout, residual and layer norm\n",
        "        src = self.ff_layer_norm(src + self.dropout(_src))\n",
        "        \n",
        "        #src = [batch size, src len, hid dim]\n",
        "        \n",
        "        return src"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3WKffF0Gy_x6"
      },
      "source": [
        "### Mutli Head Attention Layer\n",
        "\n",
        "One of the key, novel concepts introduced by the Transformer paper is the *multi-head attention layer*. \n",
        "\n",
        "![](https://github.com/bentrevett/pytorch-seq2seq/blob/master/assets/transformer-attention.png?raw=1)\n",
        "\n",
        "Attention can be though of as *queries*, *keys* and *values* - where the query is used with the key to get an attention vector (usually the output of a *softmax* operation and has all values between 0 and 1 which sum to 1) which is then used to get a weighted sum of the values.\n",
        "\n",
        "The Transformer uses *scaled dot-product attention*, where the query and key are combined by taking the dot product between them, then applying the softmax operation and scaling by $d_k$ before finally then multiplying by the value. $d_k$ is the *head dimension*, `head_dim`, which we will shortly explain further.\n",
        "\n",
        "$$ \\text{Attention}(Q, K, V) = \\text{Softmax} \\big( \\frac{QK^T}{\\sqrt{d_k}} \\big)V $$ \n",
        "\n",
        "This is similar to standard *dot product attention* but is scaled by $d_k$, which the paper states is used to stop the results of the dot products growing large, causing gradients to become too small.\n",
        "\n",
        "However, the scaled dot-product attention isn't simply applied to the queries, keys and values. Instead of doing a single attention application the queries, keys and values have their `hid_dim` split into $h$ *heads* and the scaled dot-product attention is calculated over all heads in parallel. This means instead of paying attention to one concept per attention application, we pay attention to $h$. We then re-combine the heads into their `hid_dim` shape, thus each `hid_dim` is potentially paying attention to $h$ different concepts.\n",
        "\n",
        "$$ \\text{MultiHead}(Q, K, V) = \\text{Concat}(\\text{head}_1,...,\\text{head}_h)W^O $$\n",
        "\n",
        "$$\\text{head}_i = \\text{Attention}(QW_i^Q, KW_i^K, VW_i^V) $$\n",
        "\n",
        "$W^O$ is the linear layer applied at the end of the multi-head attention layer, `fc`. $W^Q, W^K, W^V$ are the linear layers `fc_q`, `fc_k` and `fc_v`.\n",
        "\n",
        "Walking through the module, first we calculate $QW^Q$, $KW^K$ and $VW^V$ with the linear layers, `fc_q`, `fc_k` and `fc_v`, to give us `Q`, `K` and `V`. Next, we split the `hid_dim` of the query, key and value into `n_heads` using `.view` and correctly permute them so they can be multiplied together. We then calculate the `energy` (the un-normalized attention) by multiplying `Q` and `K` together and scaling it by the square root of `head_dim`, which is calulated as `hid_dim // n_heads`. We then mask the energy so we do not pay attention over any elements of the sequeuence we shouldn't, then apply the softmax and dropout. We then apply the attention to the value heads, `V`, before combining the `n_heads` together. Finally, we multiply this $W^O$, represented by `fc_o`. \n",
        "\n",
        "Note that in our implementation the lengths of the keys and values are always the same, thus when matrix multiplying the output of the softmax, `attention`, with `V` we will always have valid dimension sizes for matrix multiplication. This multiplication is carried out using `torch.matmul` which, when both tensors are >2-dimensional, does a batched matrix multiplication over the last two dimensions of each tensor. This will be a **[query len, key len] x [value len, head dim]** batched matrix multiplication over the batch size and each head which provides the **[batch size, n heads, query len, head dim]** result.\n",
        "\n",
        "One thing that looks strange at first is that dropout is applied directly to the attention. This means that our attention vector will most probably not sum to 1 and we may pay full attention to a token but the attention over that token is set to 0 by dropout. This is never explained, or even mentioned, in the paper however is used by the [official implementation](https://github.com/tensorflow/tensor2tensor/) and every Transformer implementation since, [including BERT](https://github.com/google-research/bert/)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "U0VFTLGRy_x6"
      },
      "source": [
        "class MultiHeadAttentionLayer(nn.Module):\n",
        "    def __init__(self, hid_dim, n_heads, dropout, device):\n",
        "        super().__init__()\n",
        "        \n",
        "        assert hid_dim % n_heads == 0\n",
        "        \n",
        "        self.hid_dim = hid_dim\n",
        "        self.n_heads = n_heads\n",
        "        self.head_dim = hid_dim // n_heads\n",
        "        \n",
        "        self.fc_q = nn.Linear(hid_dim, hid_dim)\n",
        "        self.fc_k = nn.Linear(hid_dim, hid_dim)\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",
        "        self.scale = torch.sqrt(torch.FloatTensor([self.head_dim])).to(device)\n",
        "        \n",
        "    def forward(self, query, key, value, mask = None):\n",
        "        \n",
        "        batch_size = query.shape[0]\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",
        "        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",
        "        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",
        "        energy = torch.matmul(Q, K.permute(0, 1, 3, 2)) / self.scale\n",
        "        \n",
        "        #energy = [batch size, n heads, query len, key len]\n",
        "        \n",
        "        if mask is not None:\n",
        "            energy = energy.masked_fill(mask == 0, -1e10)\n",
        "        \n",
        "        attention = torch.softmax(energy, dim = -1)\n",
        "                \n",
        "        #attention = [batch size, n heads, query len, key len]\n",
        "                \n",
        "        x = torch.matmul(self.dropout(attention), V)\n",
        "        \n",
        "        #x = [batch size, n heads, query len, head dim]\n",
        "        \n",
        "        x = x.permute(0, 2, 1, 3).contiguous()\n",
        "        \n",
        "        #x = [batch size, query len, n heads, head dim]\n",
        "        \n",
        "        x = x.view(batch_size, -1, self.hid_dim)\n",
        "        \n",
        "        #x = [batch size, query len, hid dim]\n",
        "        \n",
        "        x = self.fc_o(x)\n",
        "        \n",
        "        #x = [batch size, query len, hid dim]\n",
        "        \n",
        "        return x, attention"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q12JMc8by_x7"
      },
      "source": [
        "### Position-wise Feedforward Layer\n",
        "\n",
        "The other main block inside the encoder layer is the *position-wise feedforward layer* This is relatively simple compared to the multi-head attention layer. The input is transformed from `hid_dim` to `pf_dim`, where `pf_dim` is usually a lot larger than `hid_dim`. The original Transformer used a `hid_dim` of 512 and a `pf_dim` of 2048. The ReLU activation function and dropout are applied before it is transformed back into a `hid_dim` representation. \n",
        "\n",
        "Why is this used? Unfortunately, it is never explained in the paper.\n",
        "\n",
        "BERT uses the [GELU](https://arxiv.org/abs/1606.08415) activation function, which can be used by simply switching `torch.relu` for `F.gelu`. Why did they use GELU? Again, it is never explained."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "tJfnEHOwy_x7"
      },
      "source": [
        "class PositionwiseFeedforwardLayer(nn.Module):\n",
        "    def __init__(self, hid_dim, pf_dim, dropout):\n",
        "        super().__init__()\n",
        "        \n",
        "        self.fc_1 = nn.Linear(hid_dim, pf_dim)\n",
        "        self.fc_2 = nn.Linear(pf_dim, hid_dim)\n",
        "        \n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "        \n",
        "    def forward(self, x):\n",
        "        \n",
        "        #x = [batch size, seq len, hid dim]\n",
        "        \n",
        "        x = self.dropout(torch.relu(self.fc_1(x)))\n",
        "        \n",
        "        #x = [batch size, seq len, pf dim]\n",
        "        \n",
        "        x = self.fc_2(x)\n",
        "        \n",
        "        #x = [batch size, seq len, hid dim]\n",
        "        \n",
        "        return x"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3XbLg8mty_x7"
      },
      "source": [
        "### Decoder\n",
        "\n",
        "The objective of the decoder is to take the encoded representation of the source sentence, $Z$, and convert it into predicted tokens in the target sentence, $\\hat{Y}$. We then compare $\\hat{Y}$ with the actual tokens in the target sentence, $Y$, to calculate our loss, which will be used to calculate the gradients of our parameters and then use our optimizer to update our weights in order to improve our predictions. \n",
        "\n",
        "![](https://github.com/bentrevett/pytorch-seq2seq/blob/master/assets/transformer-decoder.png?raw=1)\n",
        "\n",
        "The decoder is similar to encoder, however it now has two multi-head attention layers. A *masked multi-head attention layer* over the target sequence, and a multi-head attention layer which uses the decoder representation as the query and the encoder representation as the key and value.\n",
        "\n",
        "The decoder uses positional embeddings and combines - via an elementwise sum - them with the scaled embedded target tokens, followed by dropout. Again, our positional encodings have a \"vocabulary\" of 100, which means they can accept sequences up to 100 tokens long. This can be increased if desired.\n",
        "\n",
        "The combined embeddings are then passed through the $N$ decoder layers, along with the encoded source, `enc_src`, and the source and target masks. Note that the number of layers in the encoder does not have to be equal to the number of layers in the decoder, even though they are both denoted by $N$.\n",
        "\n",
        "The decoder representation after the $N^{th}$ layer is then passed through a linear layer, `fc_out`. In PyTorch, the softmax operation is contained within our loss function, so we do not explicitly need to use a softmax layer here.\n",
        "\n",
        "As well as using the source mask, as we did in the encoder to prevent our model attending to `<pad>` tokens, we also use a target mask. This will be explained further in the `Seq2Seq` model which encapsulates both the encoder and decoder, but the gist of it is that it performs a similar operation as the decoder padding in the convolutional sequence-to-sequence model. As we are processing all of the target tokens at once in parallel we need a method of stopping the decoder from \"cheating\" by simply \"looking\" at what the next token in the target sequence is and outputting it. \n",
        "\n",
        "Our decoder layer also outputs the normalized attention values so we can later plot them to see what our model is actually paying attention to."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "szwZ39lSy_x7"
      },
      "source": [
        "class Decoder(nn.Module):\n",
        "    def __init__(self, \n",
        "                 output_dim, \n",
        "                 hid_dim, \n",
        "                 n_layers, \n",
        "                 n_heads, \n",
        "                 pf_dim, \n",
        "                 dropout, \n",
        "                 device,\n",
        "                 max_length = 100):\n",
        "        super().__init__()\n",
        "        \n",
        "        self.device = device\n",
        "        \n",
        "        self.tok_embedding = nn.Embedding(output_dim, hid_dim)\n",
        "        self.pos_embedding = nn.Embedding(max_length, hid_dim)\n",
        "        \n",
        "        self.layers = nn.ModuleList([DecoderLayer(hid_dim, \n",
        "                                                  n_heads, \n",
        "                                                  pf_dim, \n",
        "                                                  dropout, \n",
        "                                                  device)\n",
        "                                     for _ in range(n_layers)])\n",
        "        \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",
        "        pos = torch.arange(0, trg_len).unsqueeze(0).repeat(batch_size, 1).to(self.device)\n",
        "                            \n",
        "        #pos = [batch size, trg len]\n",
        "            \n",
        "        trg = self.dropout((self.tok_embedding(trg) * self.scale) + self.pos_embedding(pos))\n",
        "                \n",
        "        #trg = [batch size, trg len, hid dim]\n",
        "        \n",
        "        for layer in self.layers:\n",
        "            trg, attention = layer(trg, enc_src, trg_mask, src_mask)\n",
        "        \n",
        "        #trg = [batch size, trg len, hid dim]\n",
        "        #attention = [batch size, n heads, trg len, src len]\n",
        "        \n",
        "        output = self.fc_out(trg)\n",
        "        \n",
        "        #output = [batch size, trg len, output dim]\n",
        "            \n",
        "        return output, attention"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CCX5MMGky_x8"
      },
      "source": [
        "### Decoder Layer\n",
        "\n",
        "As mentioned previously, the decoder layer is similar to the encoder layer except that it now has two multi-head attention layers, `self_attention` and `encoder_attention`. \n",
        "\n",
        "The first performs self-attention, as in the encoder, by using the decoder representation so far as the query, key and value. This is followed by dropout, residual connection and layer normalization. This `self_attention` layer uses the target sequence mask, `trg_mask`, in order to prevent the decoder from \"cheating\" by paying attention to tokens that are \"ahead\" of the one it is currently processing as it processes all tokens in the target sentence in parallel.\n",
        "\n",
        "The second is how we actually feed the encoded source sentence, `enc_src`, into our decoder. In this multi-head attention layer the queries are the decoder representations and the keys and values are the encoder representations. Here, the source mask, `src_mask` is used to prevent the multi-head attention layer from attending to `<pad>` tokens within the source sentence. This is then followed by the dropout, residual connection and layer normalization layers. \n",
        "\n",
        "Finally, we pass this through the position-wise feedforward layer and yet another sequence of dropout, residual connection and layer normalization.\n",
        "\n",
        "The decoder layer isn't introducing any new concepts, just using the same set of layers as the encoder in a slightly different way."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "9iviBr13y_x8"
      },
      "source": [
        "class DecoderLayer(nn.Module):\n",
        "    def __init__(self, \n",
        "                 hid_dim, \n",
        "                 n_heads, \n",
        "                 pf_dim, \n",
        "                 dropout, \n",
        "                 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",
        "        self.self_attention = MultiHeadAttentionLayer(hid_dim, n_heads, dropout, device)\n",
        "        self.encoder_attention = MultiHeadAttentionLayer(hid_dim, n_heads, dropout, device)\n",
        "        self.positionwise_feedforward = PositionwiseFeedforwardLayer(hid_dim, \n",
        "                                                                     pf_dim, \n",
        "                                                                     dropout)\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",
        "        #self attention\n",
        "        _trg, _ = self.self_attention(trg, trg, trg, trg_mask)\n",
        "        \n",
        "        #dropout, residual connection and layer norm\n",
        "        trg = self.self_attn_layer_norm(trg + self.dropout(_trg))\n",
        "            \n",
        "        #trg = [batch size, trg len, hid dim]\n",
        "            \n",
        "        #encoder attention\n",
        "        _trg, attention = self.encoder_attention(trg, enc_src, enc_src, src_mask)\n",
        "        \n",
        "        #dropout, residual connection and layer norm\n",
        "        trg = self.enc_attn_layer_norm(trg + self.dropout(_trg))\n",
        "                    \n",
        "        #trg = [batch size, trg len, hid dim]\n",
        "        \n",
        "        #positionwise feedforward\n",
        "        _trg = self.positionwise_feedforward(trg)\n",
        "        \n",
        "        #dropout, residual and layer norm\n",
        "        trg = self.ff_layer_norm(trg + self.dropout(_trg))\n",
        "        \n",
        "        #trg = [batch size, trg len, hid dim]\n",
        "        #attention = [batch size, n heads, trg len, src len]\n",
        "        \n",
        "        return trg, attention"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "_OBeq20Ky_x9"
      },
      "source": [
        "class Seq2Seq(nn.Module):\n",
        "    def __init__(self, \n",
        "                 encoder, \n",
        "                 decoder, \n",
        "                 src_pad_idx, \n",
        "                 trg_pad_idx, \n",
        "                 device):\n",
        "        super().__init__()\n",
        "        \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",
        "        \n",
        "        #src = [batch size, src len]\n",
        "        \n",
        "        src_mask = (src != self.src_pad_idx).unsqueeze(1).unsqueeze(2)\n",
        "\n",
        "        #src_mask = [batch size, 1, 1, src len]\n",
        "\n",
        "        return src_mask\n",
        "    \n",
        "    def make_trg_mask(self, trg):\n",
        "        \n",
        "        #trg = [batch size, trg len]\n",
        "        \n",
        "        trg_pad_mask = (trg != self.trg_pad_idx).unsqueeze(1).unsqueeze(2)\n",
        "        \n",
        "        #trg_pad_mask = [batch size, 1, 1, trg len]\n",
        "        \n",
        "        trg_len = trg.shape[1]\n",
        "        \n",
        "        trg_sub_mask = torch.tril(torch.ones((trg_len, trg_len), device = self.device)).bool()\n",
        "        \n",
        "        #trg_sub_mask = [trg len, trg len]\n",
        "            \n",
        "        trg_mask = trg_pad_mask & trg_sub_mask\n",
        "        \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",
        "        \n",
        "        #enc_src = [batch size, src len, hid dim]\n",
        "                \n",
        "        output, attention = self.decoder(trg, enc_src, trg_mask, src_mask)\n",
        "        \n",
        "        #output = [batch size, trg len, output dim]\n",
        "        #attention = [batch size, n heads, trg len, src len]\n",
        "        \n",
        "        return output, attention"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rJVTCJBTy_x9"
      },
      "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)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8yBBerzFy_x9"
      },
      "source": [
        "Then, use them to define our whole sequence-to-sequence encapsulating model."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "JYJ39bNAy_x9"
      },
      "source": [
        "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": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "J1zCByqIy_x-",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "32500c1c-cfc4-4a2e-ec8b-f9fd13cbcff2"
      },
      "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": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "The model has 9,038,853 trainable parameters\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7WtaK3KZy_x-"
      },
      "source": [
        "def initialize_weights(m):\n",
        "    if hasattr(m, 'weight') and m.weight.dim() > 1:\n",
        "        nn.init.xavier_uniform_(m.weight.data)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "chTWTXGDy_x-"
      },
      "source": [
        "model.apply(initialize_weights);"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "C9FFqqbNy_x_"
      },
      "source": [
        "LEARNING_RATE = 0.0005\n",
        "\n",
        "optimizer = torch.optim.Adam(model.parameters(), lr = LEARNING_RATE)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7D1VEFeoy_x_"
      },
      "source": [
        "criterion = nn.CrossEntropyLoss(ignore_index = TRG_PAD_IDX)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UQYJNk2_y_x_"
      },
      "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": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EWG1iBkJy_yA"
      },
      "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": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AFmtwEy1y_yA"
      },
      "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": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "suKcAoNby_yB",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "8797a47f-4627-40e8-845f-932e12cfb95a"
      },
      "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": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Epoch: 01 | Time: 0m 59s\n",
            "\tTrain Loss: 3.636 | Train PPL:  37.923\n",
            "\t Val. Loss: 2.623 |  Val. PPL:  13.772\n",
            "Epoch: 02 | Time: 0m 59s\n",
            "\tTrain Loss: 2.463 | Train PPL:  11.740\n",
            "\t Val. Loss: 2.061 |  Val. PPL:   7.853\n",
            "Epoch: 03 | Time: 0m 59s\n",
            "\tTrain Loss: 1.979 | Train PPL:   7.232\n",
            "\t Val. Loss: 1.822 |  Val. PPL:   6.182\n",
            "Epoch: 04 | Time: 1m 0s\n",
            "\tTrain Loss: 1.689 | Train PPL:   5.415\n",
            "\t Val. Loss: 1.705 |  Val. PPL:   5.500\n",
            "Epoch: 05 | Time: 0m 59s\n",
            "\tTrain Loss: 1.485 | Train PPL:   4.415\n",
            "\t Val. Loss: 1.664 |  Val. PPL:   5.282\n",
            "Epoch: 06 | Time: 0m 59s\n",
            "\tTrain Loss: 1.326 | Train PPL:   3.767\n",
            "\t Val. Loss: 1.647 |  Val. PPL:   5.192\n",
            "Epoch: 07 | Time: 0m 59s\n",
            "\tTrain Loss: 1.200 | Train PPL:   3.321\n",
            "\t Val. Loss: 1.639 |  Val. PPL:   5.148\n",
            "Epoch: 08 | Time: 0m 59s\n",
            "\tTrain Loss: 1.095 | Train PPL:   2.988\n",
            "\t Val. Loss: 1.651 |  Val. PPL:   5.212\n",
            "Epoch: 09 | Time: 0m 59s\n",
            "\tTrain Loss: 1.003 | Train PPL:   2.725\n",
            "\t Val. Loss: 1.659 |  Val. PPL:   5.253\n",
            "Epoch: 10 | Time: 0m 59s\n",
            "\tTrain Loss: 0.927 | Train PPL:   2.526\n",
            "\t Val. Loss: 1.710 |  Val. PPL:   5.527\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rO_sS7Qby_yB"
      },
      "source": [
        "We load our \"best\" parameters and manage to achieve a better test perplexity than all previous models."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "p7QirlGyy_yB",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "da4ba35b-d36f-4395-e3b3-7de394965d38"
      },
      "source": [
        "model.load_state_dict(torch.load('tut6-model.pt'))\n",
        "\n",
        "test_loss = evaluate(model, test_iterator, criterion)\n",
        "\n",
        "print(f'| Test Loss: {test_loss:.3f} | Test PPL: {math.exp(test_loss):7.3f} |')"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "| Test Loss: 1.693 | Test PPL:   5.435 |\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "skoJQ43Gy_yC"
      },
      "source": [
        "def translate_sentence(sentence, src_field, trg_field, model, device, max_len = 50):\n",
        "    \n",
        "    model.eval()\n",
        "        \n",
        "    if isinstance(sentence, str):\n",
        "        nlp = spacy.load('de_core_news_sm')\n",
        "        tokens = [token.text.lower() for token in nlp(sentence)]\n",
        "    else:\n",
        "        tokens = [token.lower() for token in sentence]\n",
        "\n",
        "    tokens = [src_field.init_token] + tokens + [src_field.eos_token]\n",
        "        \n",
        "    src_indexes = [src_field.vocab.stoi[token] for token in tokens]\n",
        "\n",
        "    src_tensor = torch.LongTensor(src_indexes).unsqueeze(0).to(device)\n",
        "    \n",
        "    src_mask = model.make_src_mask(src_tensor)\n",
        "    \n",
        "    with torch.no_grad():\n",
        "        enc_src = model.encoder(src_tensor, src_mask)\n",
        "\n",
        "    trg_indexes = [trg_field.vocab.stoi[trg_field.init_token]]\n",
        "\n",
        "    for i in range(max_len):\n",
        "\n",
        "        trg_tensor = torch.LongTensor(trg_indexes).unsqueeze(0).to(device)\n",
        "\n",
        "        trg_mask = model.make_trg_mask(trg_tensor)\n",
        "        \n",
        "        with torch.no_grad():\n",
        "            output, attention = model.decoder(trg_tensor, enc_src, trg_mask, src_mask)\n",
        "        \n",
        "        pred_token = output.argmax(2)[:,-1].item()\n",
        "        \n",
        "        trg_indexes.append(pred_token)\n",
        "\n",
        "        if pred_token == trg_field.vocab.stoi[trg_field.eos_token]:\n",
        "            break\n",
        "    \n",
        "    trg_tokens = [trg_field.vocab.itos[i] for i in trg_indexes]\n",
        "    \n",
        "    return trg_tokens[1:], attention"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kPzaC1o5y_yD",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "6f09e97e-2e60-4bb3-ad5d-09fe699373e4"
      },
      "source": [
        "example_idx = 8\n",
        "\n",
        "src = vars(train_data.examples[example_idx])['src']\n",
        "trg = vars(train_data.examples[example_idx])['trg']\n",
        "\n",
        "print(f'src = {src}')\n",
        "print(f'trg = {trg}')"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "src = ['eine', 'frau', 'mit', 'einer', 'großen', 'geldbörse', 'geht', 'an', 'einem', 'tor', 'vorbei', '.']\n",
            "trg = ['a', 'woman', 'with', 'a', 'large', 'purse', 'is', 'walking', 'by', 'a', 'gate', '.']\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IP0dOHYKy_yD",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "d04ad06f-cc79-4e1d-8766-52fec9fbfade"
      },
      "source": [
        "translation, attention = translate_sentence(src, SRC, TRG, model, device)\n",
        "\n",
        "print(f'predicted trg = {translation}')"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "predicted trg = ['a', 'woman', 'with', 'a', 'large', 'purse', 'walks', 'by', 'a', 'gate', '.', '<eos>']\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6Uqgbssvy_yE",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "2a500fa9-71ae-486f-cbad-9dec7a6e6589"
      },
      "source": [
        "example_idx = 6\n",
        "\n",
        "src = vars(valid_data.examples[example_idx])['src']\n",
        "trg = vars(valid_data.examples[example_idx])['trg']\n",
        "\n",
        "print(f'src = {src}')\n",
        "print(f'trg = {trg}')"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "src = ['ein', 'brauner', 'hund', 'rennt', 'dem', 'schwarzen', 'hund', 'hinterher', '.']\n",
            "trg = ['a', 'brown', 'dog', 'is', 'running', 'after', 'the', 'black', 'dog', '.']\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6bEYG7mGy_yE",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "787dd41e-adcc-4aee-d36f-44929849fa08"
      },
      "source": [
        "translation, attention = translate_sentence(src, SRC, TRG, model, device)\n",
        "\n",
        "print(f'predicted trg = {translation}')"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "predicted trg = ['a', 'brown', 'dog', 'chases', 'the', 'black', 'dog', '.', '<eos>']\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "F__v88CPy_yF",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "fdfcbc59-8e35-45f3-88be-3d13badb2849"
      },
      "source": [
        "example_idx = 10\n",
        "\n",
        "src = vars(test_data.examples[example_idx])['src']\n",
        "trg = vars(test_data.examples[example_idx])['trg']\n",
        "\n",
        "print(f'src = {src}')\n",
        "print(f'trg = {trg}')"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "src = ['eine', 'mutter', 'und', 'ihr', 'kleiner', 'sohn', 'genießen', 'einen', 'schönen', 'tag', 'im', 'freien', '.']\n",
            "trg = ['a', 'mother', 'and', 'her', 'young', 'song', 'enjoying', 'a', 'beautiful', 'day', 'outside', '.']\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "AEqV6ykQy_yG",
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "outputId": "abd9ddd5-8455-4aaf-efa4-236624a3eca3"
      },
      "source": [
        "translation, attention = translate_sentence(src, SRC, TRG, model, device)\n",
        "\n",
        "print(f'predicted trg = {translation}')"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "predicted trg = ['a', 'mother', 'and', 'her', 'young', 'son', 'enjoy', 'a', 'beautiful', 'day', '.', '<eos>']\n"
          ],
          "name": "stdout"
        }
      ]
    }
  ]
}