Sequence to Sequence Learning with Neural Networks.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "sXz3mJUr626F",
        "outputId": "4095ab9c-ba7d-4ad1-ef36-3d651931f344",
        "colab": {
          "base_uri": "https://localhost:8080/"
        }
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting datasets\n",
            "  Downloading datasets-2.21.0-py3-none-any.whl.metadata (21 kB)\n",
            "Collecting evaluate\n",
            "  Downloading evaluate-0.4.2-py3-none-any.whl.metadata (9.3 kB)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from datasets) (3.15.4)\n",
            "Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.10/dist-packages (from datasets) (1.26.4)\n",
            "Collecting pyarrow>=15.0.0 (from datasets)\n",
            "  Downloading pyarrow-17.0.0-cp310-cp310-manylinux_2_28_x86_64.whl.metadata (3.3 kB)\n",
            "Collecting dill<0.3.9,>=0.3.0 (from datasets)\n",
            "  Downloading dill-0.3.8-py3-none-any.whl.metadata (10 kB)\n",
            "Requirement already satisfied: pandas in /usr/local/lib/python3.10/dist-packages (from datasets) (2.1.4)\n",
            "Requirement already satisfied: requests>=2.32.2 in /usr/local/lib/python3.10/dist-packages (from datasets) (2.32.3)\n",
            "Requirement already satisfied: tqdm>=4.66.3 in /usr/local/lib/python3.10/dist-packages (from datasets) (4.66.5)\n",
            "Collecting xxhash (from datasets)\n",
            "  Downloading xxhash-3.5.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (12 kB)\n",
            "Collecting multiprocess (from datasets)\n",
            "  Downloading multiprocess-0.70.16-py310-none-any.whl.metadata (7.2 kB)\n",
            "Requirement already satisfied: fsspec<=2024.6.1,>=2023.1.0 in /usr/local/lib/python3.10/dist-packages (from fsspec[http]<=2024.6.1,>=2023.1.0->datasets) (2024.6.1)\n",
            "Requirement already satisfied: aiohttp in /usr/local/lib/python3.10/dist-packages (from datasets) (3.10.5)\n",
            "Requirement already satisfied: huggingface-hub>=0.21.2 in /usr/local/lib/python3.10/dist-packages (from datasets) (0.24.6)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.10/dist-packages (from datasets) (24.1)\n",
            "Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.10/dist-packages (from datasets) (6.0.2)\n",
            "Requirement already satisfied: aiohappyeyeballs>=2.3.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (2.4.0)\n",
            "Requirement already satisfied: aiosignal>=1.1.2 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (1.3.1)\n",
            "Requirement already satisfied: attrs>=17.3.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (24.2.0)\n",
            "Requirement already satisfied: frozenlist>=1.1.1 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (1.4.1)\n",
            "Requirement already satisfied: multidict<7.0,>=4.5 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (6.0.5)\n",
            "Requirement already satisfied: yarl<2.0,>=1.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (1.9.4)\n",
            "Requirement already satisfied: async-timeout<5.0,>=4.0 in /usr/local/lib/python3.10/dist-packages (from aiohttp->datasets) (4.0.3)\n",
            "Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub>=0.21.2->datasets) (4.12.2)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests>=2.32.2->datasets) (3.3.2)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests>=2.32.2->datasets) (3.8)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests>=2.32.2->datasets) (2.0.7)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests>=2.32.2->datasets) (2024.8.30)\n",
            "Requirement already satisfied: python-dateutil>=2.8.2 in /usr/local/lib/python3.10/dist-packages (from pandas->datasets) (2.8.2)\n",
            "Requirement already satisfied: pytz>=2020.1 in /usr/local/lib/python3.10/dist-packages (from pandas->datasets) (2024.1)\n",
            "Requirement already satisfied: tzdata>=2022.1 in /usr/local/lib/python3.10/dist-packages (from pandas->datasets) (2024.1)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.10/dist-packages (from python-dateutil>=2.8.2->pandas->datasets) (1.16.0)\n",
            "Downloading datasets-2.21.0-py3-none-any.whl (527 kB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m527.3/527.3 kB\u001b[0m \u001b[31m9.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading evaluate-0.4.2-py3-none-any.whl (84 kB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m84.1/84.1 kB\u001b[0m \u001b[31m4.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading dill-0.3.8-py3-none-any.whl (116 kB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m116.3/116.3 kB\u001b[0m \u001b[31m5.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading pyarrow-17.0.0-cp310-cp310-manylinux_2_28_x86_64.whl (39.9 MB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m39.9/39.9 MB\u001b[0m \u001b[31m9.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading multiprocess-0.70.16-py310-none-any.whl (134 kB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m134.8/134.8 kB\u001b[0m \u001b[31m6.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hDownloading xxhash-3.5.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (194 kB)\n",
            "\u001b[2K   \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m194.1/194.1 kB\u001b[0m \u001b[31m8.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hInstalling collected packages: xxhash, pyarrow, dill, multiprocess, datasets, evaluate\n",
            "  Attempting uninstall: pyarrow\n",
            "    Found existing installation: pyarrow 14.0.2\n",
            "    Uninstalling pyarrow-14.0.2:\n",
            "      Successfully uninstalled pyarrow-14.0.2\n",
            "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
            "cudf-cu12 24.4.1 requires pyarrow<15.0.0a0,>=14.0.1, but you have pyarrow 17.0.0 which is incompatible.\n",
            "ibis-framework 8.0.0 requires pyarrow<16,>=2, but you have pyarrow 17.0.0 which is incompatible.\u001b[0m\u001b[31m\n",
            "\u001b[0mSuccessfully installed datasets-2.21.0 dill-0.3.8 evaluate-0.4.2 multiprocess-0.70.16 pyarrow-17.0.0 xxhash-3.5.0\n",
            "Collecting de-core-news-sm==3.7.0\n",
            "  Downloading https://github.com/explosion/spacy-models/releases/download/de_core_news_sm-3.7.0/de_core_news_sm-3.7.0-py3-none-any.whl (14.6 MB)\n",
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m14.6/14.6 MB\u001b[0m \u001b[31m85.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25hRequirement already satisfied: spacy<3.8.0,>=3.7.0 in /usr/local/lib/python3.10/dist-packages (from de-core-news-sm==3.7.0) (3.7.6)\n",
            "Requirement already satisfied: spacy-legacy<3.1.0,>=3.0.11 in /usr/local/lib/python3.10/dist-packages (from spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (3.0.12)\n",
            "Requirement already satisfied: spacy-loggers<2.0.0,>=1.0.0 in /usr/local/lib/python3.10/dist-packages (from spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (1.0.5)\n",
            "Requirement already satisfied: murmurhash<1.1.0,>=0.28.0 in /usr/local/lib/python3.10/dist-packages (from spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (1.0.10)\n",
            "Requirement already satisfied: cymem<2.1.0,>=2.0.2 in /usr/local/lib/python3.10/dist-packages (from spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (2.0.8)\n",
            "Requirement already satisfied: preshed<3.1.0,>=3.0.2 in /usr/local/lib/python3.10/dist-packages (from spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (3.0.9)\n",
            "Requirement already satisfied: thinc<8.3.0,>=8.2.2 in /usr/local/lib/python3.10/dist-packages (from spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (8.2.5)\n",
            "Requirement already satisfied: wasabi<1.2.0,>=0.9.1 in /usr/local/lib/python3.10/dist-packages (from spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (1.1.3)\n",
            "Requirement already satisfied: srsly<3.0.0,>=2.4.3 in /usr/local/lib/python3.10/dist-packages (from spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (2.4.8)\n",
            "Requirement already satisfied: catalogue<2.1.0,>=2.0.6 in /usr/local/lib/python3.10/dist-packages (from spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (2.0.10)\n",
            "Requirement already satisfied: weasel<0.5.0,>=0.1.0 in /usr/local/lib/python3.10/dist-packages (from spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (0.4.1)\n",
            "Requirement already satisfied: typer<1.0.0,>=0.3.0 in /usr/local/lib/python3.10/dist-packages (from spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (0.12.5)\n",
            "Requirement already satisfied: tqdm<5.0.0,>=4.38.0 in /usr/local/lib/python3.10/dist-packages (from spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (4.66.5)\n",
            "Requirement already satisfied: requests<3.0.0,>=2.13.0 in /usr/local/lib/python3.10/dist-packages (from spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (2.32.3)\n",
            "Requirement already satisfied: pydantic!=1.8,!=1.8.1,<3.0.0,>=1.7.4 in /usr/local/lib/python3.10/dist-packages (from spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (2.8.2)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (3.1.4)\n",
            "Requirement already satisfied: setuptools in /usr/local/lib/python3.10/dist-packages (from spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (71.0.4)\n",
            "Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.10/dist-packages (from spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (24.1)\n",
            "Requirement already satisfied: langcodes<4.0.0,>=3.2.0 in /usr/local/lib/python3.10/dist-packages (from spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (3.4.0)\n",
            "Requirement already satisfied: numpy>=1.19.0 in /usr/local/lib/python3.10/dist-packages (from spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (1.26.4)\n",
            "Requirement already satisfied: language-data>=1.2 in /usr/local/lib/python3.10/dist-packages (from langcodes<4.0.0,>=3.2.0->spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (1.2.0)\n",
            "Requirement already satisfied: annotated-types>=0.4.0 in /usr/local/lib/python3.10/dist-packages (from pydantic!=1.8,!=1.8.1,<3.0.0,>=1.7.4->spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (0.7.0)\n",
            "Requirement already satisfied: pydantic-core==2.20.1 in /usr/local/lib/python3.10/dist-packages (from pydantic!=1.8,!=1.8.1,<3.0.0,>=1.7.4->spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (2.20.1)\n",
            "Requirement already satisfied: typing-extensions>=4.6.1 in /usr/local/lib/python3.10/dist-packages (from pydantic!=1.8,!=1.8.1,<3.0.0,>=1.7.4->spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (4.12.2)\n",
            "Requirement already satisfied: charset-normalizer<4,>=2 in /usr/local/lib/python3.10/dist-packages (from requests<3.0.0,>=2.13.0->spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (3.3.2)\n",
            "Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests<3.0.0,>=2.13.0->spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (3.8)\n",
            "Requirement already satisfied: urllib3<3,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests<3.0.0,>=2.13.0->spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (2.0.7)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests<3.0.0,>=2.13.0->spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (2024.8.30)\n",
            "Requirement already satisfied: blis<0.8.0,>=0.7.8 in /usr/local/lib/python3.10/dist-packages (from thinc<8.3.0,>=8.2.2->spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (0.7.11)\n",
            "Requirement already satisfied: confection<1.0.0,>=0.0.1 in /usr/local/lib/python3.10/dist-packages (from thinc<8.3.0,>=8.2.2->spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (0.1.5)\n",
            "Requirement already satisfied: click>=8.0.0 in /usr/local/lib/python3.10/dist-packages (from typer<1.0.0,>=0.3.0->spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (8.1.7)\n",
            "Requirement already satisfied: shellingham>=1.3.0 in /usr/local/lib/python3.10/dist-packages (from typer<1.0.0,>=0.3.0->spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (1.5.4)\n",
            "Requirement already satisfied: rich>=10.11.0 in /usr/local/lib/python3.10/dist-packages (from typer<1.0.0,>=0.3.0->spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (13.8.0)\n",
            "Requirement already satisfied: cloudpathlib<1.0.0,>=0.7.0 in /usr/local/lib/python3.10/dist-packages (from weasel<0.5.0,>=0.1.0->spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (0.19.0)\n",
            "Requirement already satisfied: smart-open<8.0.0,>=5.2.1 in /usr/local/lib/python3.10/dist-packages (from weasel<0.5.0,>=0.1.0->spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (7.0.4)\n",
            "Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (2.1.5)\n",
            "Requirement already satisfied: marisa-trie>=0.7.7 in /usr/local/lib/python3.10/dist-packages (from language-data>=1.2->langcodes<4.0.0,>=3.2.0->spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (1.2.0)\n",
            "Requirement already satisfied: markdown-it-py>=2.2.0 in /usr/local/lib/python3.10/dist-packages (from rich>=10.11.0->typer<1.0.0,>=0.3.0->spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (3.0.0)\n",
            "Requirement already satisfied: pygments<3.0.0,>=2.13.0 in /usr/local/lib/python3.10/dist-packages (from rich>=10.11.0->typer<1.0.0,>=0.3.0->spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (2.16.1)\n",
            "Requirement already satisfied: wrapt in /usr/local/lib/python3.10/dist-packages (from smart-open<8.0.0,>=5.2.1->weasel<0.5.0,>=0.1.0->spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (1.16.0)\n",
            "Requirement already satisfied: mdurl~=0.1 in /usr/local/lib/python3.10/dist-packages (from markdown-it-py>=2.2.0->rich>=10.11.0->typer<1.0.0,>=0.3.0->spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (0.1.2)\n",
            "Installing collected packages: de-core-news-sm\n",
            "Successfully installed de-core-news-sm-3.7.0\n",
            "\u001b[38;5;2m✔ Download and installation successful\u001b[0m\n",
            "You can now load the package via spacy.load('de_core_news_sm')\n",
            "\u001b[38;5;3m⚠ Restart to reload dependencies\u001b[0m\n",
            "If you are in a Jupyter or Colab notebook, you may need to restart Python in\n",
            "order to load all the package's dependencies. You can do this by selecting the\n",
            "'Restart kernel' or 'Restart runtime' option.\n"
          ]
        }
      ],
      "source": [
        "# NOTE: If you are running this notebook on Google Colab,\n",
        "#       then uncomment the two lines below and then run this cell!\n",
        "\n",
        "!pip install datasets evaluate --upgrade\n",
        "!python -m spacy download de_core_news_sm"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yFUPv7T_626H"
      },
      "source": [
        "# 1 - Sequence to Sequence Learning with Neural Networks\n",
        "\n",
        "In this series we'll be building a machine learning model to go from one sequence to another, using PyTorch. This will be done on German to English translations, but the models can be applied to any problem that involves going from one sequence to another, such as summarization, i.e. going from a sequence to a shorter sequence in the same language.\n",
        "\n",
        "In this first notebook, we'll start simple to understand the general concepts by implementing the model from the [Sequence to Sequence Learning with Neural Networks](https://arxiv.org/abs/1409.3215) paper.\n",
        "\n",
        "## Introduction\n",
        "\n",
        "The most common sequence-to-sequence (seq2seq) models are _encoder-decoder_ models, which commonly use a _recurrent neural network_ (RNN) to _encode_ the source (input) sentence into a single vector. In this notebook, we'll refer to this single vector as a _context vector_. We can think of the context vector as being an abstract representation of the entire input sentence. This vector is then _decoded_ by a second RNN which learns to output the target (output) sentence by generating it one word at a time.\n",
        "\n",
        "![](https://github.com/bentrevett/pytorch-seq2seq/blob/main/assets/seq2seq1.png?raw=1)\n",
        "\n",
        "The above image shows an example translation. The input/source sentence, \"guten morgen\", is passed through the embedding layer (yellow) and then input into the encoder (green). We also append a _start of sequence_ (`<sos>`) and _end of sequence_ (`<eos>`) token to the start and end of sentence, respectively. At each time-step, the input to the encoder RNN is both the embedding, $e$, of the current word, $e(x_t)$, as well as the hidden state from the previous time-step, $h_{t-1}$, and the encoder RNN outputs a new hidden state $h_t$. We can think of the hidden state as a vector representation of the sentence so far. The RNN can be represented as a function of both of $e(x_t)$ and $h_{t-1}$:\n",
        "\n",
        "$$h_t = \\text{EncoderRNN}(e(x_t), h_{t-1})$$\n",
        "\n",
        "We're using the term RNN generally here, it could be any recurrent architecture, such as an _LSTM_ (Long Short-Term Memory) or a _GRU_ (Gated Recurrent Unit).\n",
        "\n",
        "Here, we have $X = \\{x_1, x_2, ..., x_T\\}$, where $x_1 = \\text{<sos>}, x_2 = \\text{guten}$, etc. The initial hidden state, $h_0$, is usually either initialized to zeros or a learned parameter.\n",
        "\n",
        "Once the final word, $x_T$, has been passed into the RNN via the embedding layer, we use the final hidden state, $h_T$, as the context vector, i.e. $h_T = z$. This is a vector representation of the entire source sentence.\n",
        "\n",
        "Now we have our context vector, $z$, we can start decoding it to get the output/target sentence, \"good morning\". Again, we append start and end of sequence tokens to the target sentence. At each time-step, the input to the decoder RNN (blue) is the embedding, $d$, of current word, $d(y_t)$, as well as the hidden state from the previous time-step, $s_{t-1}$, where the initial decoder hidden state, $s_0$, is the context vector, $s_0 = z = h_T$, i.e. the initial decoder hidden state is the final encoder hidden state. Thus, similar to the encoder, we can represent the decoder as:\n",
        "\n",
        "$$s_t = \\text{DecoderRNN}(d(y_t), s_{t-1})$$\n",
        "\n",
        "Although the input/source embedding layer, $e$, and the output/target embedding layer, $d$, are both shown in yellow in the diagram they are two different embedding layers with their own parameters.\n",
        "\n",
        "In the decoder, we need to go from the hidden state to an actual word, therefore at each time-step we use $s_t$ to predict (by passing it through a `Linear` layer, shown in purple) what we think is the next word in the sequence, $\\hat{y}_t$.\n",
        "\n",
        "$$\\hat{y}_t = f(s_t)$$\n",
        "\n",
        "The words in the decoder are always generated one after another, with one per time-step. We always use `<sos>` for the first input to the decoder, $y_1$, but for subsequent inputs, $y_{t>1}$, we will sometimes use the actual, ground truth next word in the sequence, $y_t$ and sometimes use the word predicted by our decoder, $\\hat{y}_{t-1}$. This is called _teacher forcing_, see a bit more info about it [here](https://machinelearningmastery.com/teacher-forcing-for-recurrent-neural-networks/).\n",
        "\n",
        "When training/testing our model, we always know how many words are in our target sentence, so we stop generating words once we hit that many. During inference it is common to keep generating words until the model outputs an `<eos>` token or after a certain amount of words have been generated.\n",
        "\n",
        "Once we have our predicted target sentence, $\\hat{Y} = \\{ \\hat{y}_1, \\hat{y}_2, ..., \\hat{y}_T \\}$, we compare it against our actual target sentence, $Y = \\{ y_1, y_2, ..., y_T \\}$, to calculate our loss. We then use this loss to update all of the parameters in our model.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qVqgqUsi626L"
      },
      "source": [
        "## Preparing Data\n",
        "\n",
        "First up, importing all the necessary libraries. The main ones we'll be using are:\n",
        "\n",
        "-   [PyTorch](https://pytorch.org/) for creating the models\n",
        "-   [spaCy](https://spacy.io/) to assist in the tokenization of the data\n",
        "-   [torchtext](https://github.com/pytorch/text) to provider some helper functions\n",
        "-   [datasets](https://huggingface.co/docs/datasets/index) to load and manipulate our data\n",
        "-   [evaluate](https://huggingface.co/docs/evaluate/index) for calculating metrics\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XiGVs76e626O",
        "outputId": "883f2b6c-41b8-4297-897f-f8d2c9f28db9"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/home/ben/miniconda3/envs/main/lib/python3.9/site-packages/thinc/compat.py:36: UserWarning: 'has_mps' is deprecated, please use 'torch.backends.mps.is_built()'\n",
            "  hasattr(torch, \"has_mps\")\n",
            "/home/ben/miniconda3/envs/main/lib/python3.9/site-packages/thinc/compat.py:37: UserWarning: 'has_mps' is deprecated, please use 'torch.backends.mps.is_built()'\n",
            "  and torch.has_mps  # type: ignore[attr-defined]\n"
          ]
        }
      ],
      "source": [
        "import torch\n",
        "import torch.nn as nn\n",
        "import torch.optim as optim\n",
        "import random\n",
        "import numpy as np\n",
        "import spacy\n",
        "import datasets\n",
        "import torchtext\n",
        "import tqdm\n",
        "import evaluate"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TWUTfx-P626S"
      },
      "source": [
        "We'll set all possible random seeds for deterministic results.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2_B0Qaep626T"
      },
      "outputs": [],
      "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"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q2R-haEV626U"
      },
      "source": [
        "### Dataset\n",
        "\n",
        "Next, we'll load our dataset using the `datasets` library. When using the `load_dataset` function we pass the name of the dataset, `bentrevett/multi30k`.\n",
        "\n",
        "The dataset we'll be using is a subset of the [Multi30k dataset](https://github.com/multi30k/dataset), which is hosted [here](https://huggingface.co/datasets/bentrevett/multi30k) on the HuggingFace dataset hub. This subset has ~30,000 parallel English and German sentences obtained using the task 1 raw data from [here](https://github.com/multi30k/dataset/tree/master/data/task1/raw). We use the \"2016\" versions of the test sets.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dNO6kFK5626U"
      },
      "outputs": [],
      "source": [
        "dataset = datasets.load_dataset(\"bentrevett/multi30k\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eUiBGYp2626V"
      },
      "source": [
        "We can see the `dataset` object (a `DatasetDict`) contains the train, validation and test splits, the amount of examples in each split, and the features in each split (\"en\" and \"de\").\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jEDcUhFo626V",
        "outputId": "90da4cd8-5f21-4f4c-aa41-02d4e998d9c8"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "DatasetDict({\n",
              "    train: Dataset({\n",
              "        features: ['en', 'de'],\n",
              "        num_rows: 29000\n",
              "    })\n",
              "    validation: Dataset({\n",
              "        features: ['en', 'de'],\n",
              "        num_rows: 1014\n",
              "    })\n",
              "    test: Dataset({\n",
              "        features: ['en', 'de'],\n",
              "        num_rows: 1000\n",
              "    })\n",
              "})"
            ]
          },
          "execution_count": 5,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "dataset"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0lDUlYQz626W"
      },
      "source": [
        "For convenience, we create a variable for each split. Each being a `Dataset` object.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZfKK-7Bs626W"
      },
      "outputs": [],
      "source": [
        "train_data, valid_data, test_data = (\n",
        "    dataset[\"train\"],\n",
        "    dataset[\"validation\"],\n",
        "    dataset[\"test\"],\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "e4spRQgz626X"
      },
      "source": [
        "We can index into each `Dataset` to view an individual example. Each example has two features: \"en\" and \"de\", which are the parallel English and German sentences.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kHCAgmNd626X",
        "outputId": "2b98a07f-c353-4d19-be61-e2ca36bda37b"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{'en': 'Two young, White males are outside near many bushes.',\n",
              " 'de': 'Zwei junge weiße Männer sind im Freien in der Nähe vieler Büsche.'}"
            ]
          },
          "execution_count": 7,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "train_data[0]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mwBCo2v2626X"
      },
      "source": [
        "### Tokenizers\n",
        "\n",
        "Next, we'll load the spaCy models that contain the tokenizers.\n",
        "\n",
        "A tokenizer is used to turn a string into a list of tokens that make up that string, e.g. \"good morning!\" becomes [\"good\", \"morning\", \"!\"]. We'll start talking about the sentences being a sequence of tokens from now, instead of saying they're a sequence of words. What's the difference? Well, \"good\" and \"morning\" are both words and tokens, but \"!\" is not a word. We could say \"!\" is punctuation, but the term token is more general and covers: words, punctuation, numbers and any special symbols.\n",
        "\n",
        "spaCy has model for each language (\"de_core_news_sm\" for German and \"en_core_web_sm\" for English) which need to be loaded so we can access the tokenizer of each model.\n",
        "\n",
        "**Note**: the models must first be downloaded using the following on the command line:\n",
        "\n",
        "```\n",
        "python -m spacy download en_core_web_sm\n",
        "python -m spacy download de_core_news_sm\n",
        "```\n",
        "\n",
        "We load the models as such:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0G7f8Anc626X",
        "outputId": "8e6f74d2-3dca-4db7-d991-4ebefb0f0627"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/home/ben/miniconda3/envs/main/lib/python3.9/site-packages/spacy/util.py:910: UserWarning: [W095] Model 'en_core_web_sm' (3.5.0) was trained with spaCy v3.5.0 and may not be 100% compatible with the current version (3.7.2). If you see errors or degraded performance, download a newer compatible model or retrain your custom model with the current spaCy version. For more details and available updates, run: python -m spacy validate\n",
            "  warnings.warn(warn_msg)\n",
            "/home/ben/miniconda3/envs/main/lib/python3.9/site-packages/spacy/util.py:910: UserWarning: [W095] Model 'de_core_news_sm' (3.5.0) was trained with spaCy v3.5.0 and may not be 100% compatible with the current version (3.7.2). If you see errors or degraded performance, download a newer compatible model or retrain your custom model with the current spaCy version. For more details and available updates, run: python -m spacy validate\n",
            "  warnings.warn(warn_msg)\n"
          ]
        }
      ],
      "source": [
        "en_nlp = spacy.load(\"en_core_web_sm\")\n",
        "de_nlp = spacy.load(\"de_core_news_sm\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z9xbIo_4626Y"
      },
      "source": [
        "We can call the tokenizer for each spaCy model using the `.tokenizer` method, which accepts a string and returns a sequence of `Token` objects. We can get the string from the token object using the `text` attribute.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kVDmFY05626Y",
        "outputId": "92eb5922-e856-44a5-b528-d1701ae2309a"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "['What', 'a', 'lovely', 'day', 'it', 'is', 'today', '!']"
            ]
          },
          "execution_count": 9,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "string = \"What a lovely day it is today!\"\n",
        "\n",
        "[token.text for token in en_nlp.tokenizer(string)]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QXd-bvsy626Z"
      },
      "source": [
        "Next, we'll write a function used to apply the tokenizer to all of the examples in each data split, as well as apply some other processing.\n",
        "\n",
        "This function takes in an example from the `Dataset` object, applies the tokenizers English and German spaCy models, trims the list of tokens to a maximum length, optionally converts each token to lowercase, and then appends the start of sequence and end of sequence tokens to the beginning and end of the list of tokens.\n",
        "\n",
        "This function will be used with the `map` method from a `Dataset`, which needs to return a dictionary containing the names of the features in each example where the outputs are stored. As the output feature names \"en_tokens\" and \"de_tokens\" are not already contained in the example (where we only have \"en\" and \"de\" features), this will create two new features in each example.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SmA2DcsH626Z"
      },
      "outputs": [],
      "source": [
        "def tokenize_example(example, en_nlp, de_nlp, max_length, lower, sos_token, eos_token):\n",
        "    en_tokens = [token.text for token in en_nlp.tokenizer(example[\"en\"])][:max_length]\n",
        "    de_tokens = [token.text for token in de_nlp.tokenizer(example[\"de\"])][:max_length]\n",
        "    if lower:\n",
        "        en_tokens = [token.lower() for token in en_tokens]\n",
        "        de_tokens = [token.lower() for token in de_tokens]\n",
        "    en_tokens = [sos_token] + en_tokens + [eos_token]\n",
        "    de_tokens = [sos_token] + de_tokens + [eos_token]\n",
        "    return {\"en_tokens\": en_tokens, \"de_tokens\": de_tokens}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "igGQm1qE626Z"
      },
      "source": [
        "We apply the `tokenize_example` function using the `map` method as below.\n",
        "\n",
        "The `example` argument is implied, however all additional arguments to the `tokenize_example` function need to be stored in a dictionary and passed to the `fn_kwargs` argument of `map`.\n",
        "\n",
        "Here, we're trimming all sequences to a maximum length of 1000 tokens, converting each token to lower case, and using `<sos>` and `<eos>` as the start and end of sequence tokens, respectively.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "I-NG_D9Q626Z",
        "outputId": "14286c7c-ef30-4492-e5bd-a286c6f6ad26",
        "colab": {
          "referenced_widgets": [
            "c271b513205646c6b87242841f4e79a7",
            "8273d3afab2c43709b24bcd497243a8b"
          ]
        }
      },
      "outputs": [
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "c271b513205646c6b87242841f4e79a7",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "Map:   0%|          | 0/1014 [00:00<?, ? examples/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "8273d3afab2c43709b24bcd497243a8b",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "Map:   0%|          | 0/1000 [00:00<?, ? examples/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "max_length = 1_000\n",
        "lower = True\n",
        "sos_token = \"<sos>\"\n",
        "eos_token = \"<eos>\"\n",
        "\n",
        "fn_kwargs = {\n",
        "    \"en_nlp\": en_nlp,\n",
        "    \"de_nlp\": de_nlp,\n",
        "    \"max_length\": max_length,\n",
        "    \"lower\": lower,\n",
        "    \"sos_token\": sos_token,\n",
        "    \"eos_token\": eos_token,\n",
        "}\n",
        "\n",
        "train_data = train_data.map(tokenize_example, fn_kwargs=fn_kwargs)\n",
        "valid_data = valid_data.map(tokenize_example, fn_kwargs=fn_kwargs)\n",
        "test_data = test_data.map(tokenize_example, fn_kwargs=fn_kwargs)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JdvAj8NE626a"
      },
      "source": [
        "We can now look at an example, confirming the two new features have been added; both of which are lowercased list of strings with the start/end of sequence tokens appended.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9My-ezNA626a",
        "outputId": "369564bb-34ca-48b4-ac94-8a6fcd9dd47c"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{'en': 'Two young, White males are outside near many bushes.',\n",
              " 'de': 'Zwei junge weiße Männer sind im Freien in der Nähe vieler Büsche.',\n",
              " 'en_tokens': ['<sos>',\n",
              "  'two',\n",
              "  'young',\n",
              "  ',',\n",
              "  'white',\n",
              "  'males',\n",
              "  'are',\n",
              "  'outside',\n",
              "  'near',\n",
              "  'many',\n",
              "  'bushes',\n",
              "  '.',\n",
              "  '<eos>'],\n",
              " 'de_tokens': ['<sos>',\n",
              "  'zwei',\n",
              "  'junge',\n",
              "  'weiße',\n",
              "  'männer',\n",
              "  'sind',\n",
              "  'im',\n",
              "  'freien',\n",
              "  'in',\n",
              "  'der',\n",
              "  'nähe',\n",
              "  'vieler',\n",
              "  'büsche',\n",
              "  '.',\n",
              "  '<eos>']}"
            ]
          },
          "execution_count": 12,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "train_data[0]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L8XYw871626a"
      },
      "source": [
        "### Vocabularies\n",
        "\n",
        "Next, we'll build the _vocabulary_ for the source and target languages. The vocabulary is used to associate each unique token in our dataset with an index (an integer), e.g. \"hello\" = 1, \"world\" = 2, \"bye\" = 3, \"hates\" = 4, etc. When feeding text data to our model, we convert the string into tokens and then the tokens into numbers using the vocabulary as a look up table, e.g. \"hello world\" becomes `[\"hello\", \"world\"]` which becomes `[1, 2]` using the example indices given. We do this as neural networks cannot operate on strings, only numerical values.\n",
        "\n",
        "We create the vocabulary (one for each language) from our datasets using the `build_vocab_from_iterator` function, provided by `torchtext`, which accepts an iterator where each item is a list of tokens. It then counts up the number of unique tokens and assigns each a numerical value.\n",
        "\n",
        "In theory, our vocabulary can be large enough to have an index for every unique token in our dataset. However, what happens if a token exists in our validation and test set, but not in our training set? In that case, we replace the token with an \"unknown token\", denoted by `<unk>`, which is given its own index (usually index zero). All unknown tokens are replaced by `<unk>`, even if the tokens are different, i.e. if the tokens \"gilgamesh\" and \"enkidu\" were both not within our vocabulary, then the string \"gilgamesh hates enkidu\" gets tokenized to `[\"gilgamesh\", \"hates\", \"enkidu\"]` and then becomes `[0, 24, 0]` (where \"hates\" has the index 24).\n",
        "\n",
        "Ideally, we want our model to be able to handle unknown tokens by learning to use the context around them to make translations. The only way it can learn that is if we also have unknown tokens in the training set. Hence, when creating our vocabularies with `build_vocab_from_iterator`, we use the `min_freq` argument to not create an index for tokens which appear less than `min_freq` times in our training set. In other words, when using the vocabulary, any token which does not appear at least twice in our training set will get replaced by the unknown token index when converting tokens to indices.\n",
        "\n",
        "It is important to note that a vocabulary should only be built from the training set and never the validation or test set. This prevents \"information leakage\" into our model, giving us artifically inflated validation/test scores.\n",
        "\n",
        "We also use the `specials` argument of `build_vocab_from_iterator` to pass _special tokens_. These are tokens which we want to add to the vocabulary but do not necessarily appear in our tokenized examples. These special tokens will appear first in the vocabulary. We've already discussed the `unk_token`, `sos_token`, and `eos_token`. The final special token is the `pad_token`, denoted by `<pad>`.\n",
        "\n",
        "When inputting sentences into our model, it is more efficient to pass multiple sentences at once (known as a batch), instead of one at a time. The requirement for sentences to be batched together is that they all have to be the same length (in terms of the number of tokens). The majority of our sentences are not the same length, but we can solve this by \"padding\" (adding `<pad>` tokens) the tokenized version of each sentence in a batch until they all have equal tokens to the longest sentence in the batch. For example, if we had two sentences: \"I love pizza\" and \"I hate music videos\". They would be tokenized to something like: `[\"i\", \"love\", \"pizza\"]` and `[\"i\", \"hate\", \"music\", \"videos\"]`. The first sequence of tokens would then be padded to `[\"i\", \"love\", \"pizza\", \"<pad>\"]`. Both sequences could then be converted to indexes using the vocabulary.\n",
        "\n",
        "That was a lot of information, but luckily `torchtext` handles all the fuss of building the vocabulary. We'll handle the padding and batching later.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "JRHh_ILT626a"
      },
      "outputs": [],
      "source": [
        "min_freq = 2\n",
        "unk_token = \"<unk>\"\n",
        "pad_token = \"<pad>\"\n",
        "\n",
        "special_tokens = [\n",
        "    unk_token,\n",
        "    pad_token,\n",
        "    sos_token,\n",
        "    eos_token,\n",
        "]\n",
        "\n",
        "en_vocab = torchtext.vocab.build_vocab_from_iterator(\n",
        "    train_data[\"en_tokens\"],\n",
        "    min_freq=min_freq,\n",
        "    specials=special_tokens,\n",
        ")\n",
        "\n",
        "de_vocab = torchtext.vocab.build_vocab_from_iterator(\n",
        "    train_data[\"de_tokens\"],\n",
        "    min_freq=min_freq,\n",
        "    specials=special_tokens,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EcjL4hBT626b"
      },
      "source": [
        "Now we have our vocabularies, we can check what's actually in them.\n",
        "\n",
        "We can get the first ten tokens in our vocabulary (indices 0 to 9) using the `get_itos` method, where itos = \"**i**nt **to** **s**tring\", which returns a list of tokens.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Mr78ZRUw626b",
        "outputId": "202010ae-4359-4784-97dc-69ed725c6adb"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "['<unk>', '<pad>', '<sos>', '<eos>', 'a', '.', 'in', 'the', 'on', 'man']"
            ]
          },
          "execution_count": 14,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "en_vocab.get_itos()[:10]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UDsJGIIS626c",
        "outputId": "63a9ceb3-3320-40c6-a41e-8f06927175bc"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'man'"
            ]
          },
          "execution_count": 15,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "en_vocab.get_itos()[9]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2A_Ggkvy626e"
      },
      "source": [
        "We can do the same for the German vocabulary. Notice how the special tokens are the same and in the same order (indices 0 to 3), however the rest are different. This is because the vocabularies were effectively created from different data (one English and one German) even though they were from the same examples. The indices given to tokens that are not special tokens are ordered from most frequent to least frequent (though still appearing at least `min_freq` times).\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SNoGnIB-626g",
        "outputId": "69fd61fe-ae75-4eef-fd0b-5e19046296c1"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "['<unk>', '<pad>', '<sos>', '<eos>', '.', 'ein', 'einem', 'in', 'eine', ',']"
            ]
          },
          "execution_count": 16,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "de_vocab.get_itos()[:10]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ucxq92iW626h"
      },
      "source": [
        "We can get the index from a given token using the `get_stoi` (stoi = \" **s**tring **to** **i**nt) method.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wpLsfhtn626j",
        "outputId": "0af42029-cee6-430d-9910-98a6b5717d25"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "7"
            ]
          },
          "execution_count": 17,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "en_vocab.get_stoi()[\"the\"]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oWDtj4g-626k"
      },
      "source": [
        "As a shorthand, we can just use the vocabulary as a dictionary and pass the token to get the index. Note that this doesn't work the other way around, i.e. `en_vocab[7]` does not work.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9zl05yT-626k",
        "outputId": "d39d09fa-288d-4ec7-f462-abe870bc6903"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "7"
            ]
          },
          "execution_count": 18,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "en_vocab[\"the\"]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Eb6Ppg7j626l"
      },
      "source": [
        "The `len` of each vocabulary gives us the number of unique tokens in each one. We can see that our training data had around 2000 more German tokens (that appeared at least twice) than the English data.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AJd15D9K626l",
        "outputId": "25a7ba61-7069-496c-ef0b-35dcf4ebd61e"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(5893, 7853)"
            ]
          },
          "execution_count": 19,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "len(en_vocab), len(de_vocab)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xTMV8stZ626l"
      },
      "source": [
        "We can also use the `in` keyword to get a boolean indicating if a token is in the vocabulary.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UE0c1t3z626l",
        "outputId": "8a3646c3-03c1-4d30-c111-ed531e9420f0"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "True"
            ]
          },
          "execution_count": 20,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "\"the\" in en_vocab"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dlmqqDm0626s"
      },
      "source": [
        "Remember how we converted all of our tokens to lowercase? This means that no tokens containing any uppercase characters appear in our vocabulary.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cYi2MpkC626v",
        "outputId": "e789ea43-08f7-4445-ba88-16446b3bd615"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "False"
            ]
          },
          "execution_count": 21,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "\"The\" in en_vocab"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w4ej5S5t626y"
      },
      "source": [
        "What happens if you try and get the index of a token that isn't in the vocabulary? You get index zero for the `<unk>` (unknown) token, right?\n",
        "\n",
        "Well, no. One quirk of the `torchtext` vocabulary class is that you have to manually set what value you want your vocabulary to return when you try and get the index of an out-of-vocabulary token. If you have not set this value, then you will receive an error! This is so you can set your vocabulary to return any value when trying to get the index of a token not in the vocabulary, even something like `-100`.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6xbNKTnB626z"
      },
      "outputs": [],
      "source": [
        "# en_vocab[\"The\"]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rg1WSPbT626z"
      },
      "source": [
        "We already know the index of our `<unk>` token is zero as it's the first element in our `special_tokens` list, and we've also manually inspected it using `get_itos`.\n",
        "\n",
        "However, here we'll programmatically get it and also check that both our vocabularies have the same index for the unknown and padding tokens as this simplifies some code later on.\n",
        "\n",
        "We also get the index of our `<pad>` token, as we'll use it later\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VoRfOhsf626z"
      },
      "outputs": [],
      "source": [
        "assert en_vocab[unk_token] == de_vocab[unk_token]\n",
        "assert en_vocab[pad_token] == de_vocab[pad_token]\n",
        "\n",
        "unk_index = en_vocab[unk_token]\n",
        "pad_index = en_vocab[pad_token]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "y0V8rrcY626z"
      },
      "source": [
        "Using the `set_default_index` method we can set what value is returned when we try and get the index of a token outside of our vocabulary. In this case, the index of the unknown token, `<unk>`.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qfO-fNA96260"
      },
      "outputs": [],
      "source": [
        "en_vocab.set_default_index(unk_index)\n",
        "de_vocab.set_default_index(unk_index)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mNSMXtMD6260"
      },
      "source": [
        "Now, we can happily get indexes of out of vocabulary tokens until our heart is content!\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7P73oRcM6260",
        "outputId": "abf6202b-8766-4514-b327-fcaa60aabb37"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "0"
            ]
          },
          "execution_count": 25,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "en_vocab[\"The\"]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ouQz_MR96260"
      },
      "source": [
        "And we can get the token corresponding to that index to prove it's the `<unk>` token.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5WGAFQJd6260",
        "outputId": "73f081ba-fdd4-4f7d-b7cd-bff99669d38f"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'<unk>'"
            ]
          },
          "execution_count": 26,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "en_vocab.get_itos()[0]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3m_QNFDo6261"
      },
      "source": [
        "Another useful feature of the vocabulary is the `lookup_indices` method. This takes in a list of tokens and returns a list of indices. In the below example we can see the token \"crime\" does not exist in our vocabulary so is coverted to the index of the `<unk>` token, zero, which we passed to the `set_default_index` method.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lJvF2D146261"
      },
      "outputs": [],
      "source": [
        "tokens = [\"i\", \"love\", \"watching\", \"crime\", \"shows\"]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dgRefjZk6261",
        "outputId": "84b07363-043a-4c07-f65f-d8fdb141387f"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "[956, 2169, 173, 0, 821]"
            ]
          },
          "execution_count": 28,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "en_vocab.lookup_indices(tokens)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DM9TLEKf6262"
      },
      "source": [
        "Conversely, we can use the `lookup_tokens` method to convert a list of indices back into tokens using the vocabulary. Notice how the original \"crime\" token is now an `<unk>` token. There is no way to tell what the original sequence of tokens was.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "On5pEM_V6262",
        "outputId": "9cfdf929-a5a7-4613-87e8-edaea68e8f2c"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "['i', 'love', 'watching', '<unk>', 'shows']"
            ]
          },
          "execution_count": 29,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "en_vocab.lookup_tokens(en_vocab.lookup_indices(tokens))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yCSYqz0D6262"
      },
      "source": [
        "Hopefully we've now got the gist of how the `torchtext.Vocab` class works. Time to put it into action!\n",
        "\n",
        "Just like our `tokenize_example`, we create a `numericalize_example` function which we'll use with the `map` method of our dataset. This will \"numericalize\" (a fancy way of saying convert tokens to indices) our tokens in each example using the vocabularies and return the result into new \"en_ids\" and \"de_ids\" features.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "N4qs3O676262"
      },
      "outputs": [],
      "source": [
        "def numericalize_example(example, en_vocab, de_vocab):\n",
        "    en_ids = en_vocab.lookup_indices(example[\"en_tokens\"])\n",
        "    de_ids = de_vocab.lookup_indices(example[\"de_tokens\"])\n",
        "    return {\"en_ids\": en_ids, \"de_ids\": de_ids}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JOo6OHg56263"
      },
      "source": [
        "We apply the `numericalize_example` function, passing our vocabularies in the `fn_kwargs` dictionary to the `fn_kwargs` argument.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9ZFfZhMp6263",
        "outputId": "d6520e50-2146-4335-e3e2-4810f612e99f",
        "colab": {
          "referenced_widgets": [
            "42a8f88aace442cd92bf627917a59e30",
            "1b97e7bd3a404d1fb9f13385f60aa692"
          ]
        }
      },
      "outputs": [
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "42a8f88aace442cd92bf627917a59e30",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "Map:   0%|          | 0/1014 [00:00<?, ? examples/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "1b97e7bd3a404d1fb9f13385f60aa692",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "Map:   0%|          | 0/1000 [00:00<?, ? examples/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "fn_kwargs = {\"en_vocab\": en_vocab, \"de_vocab\": de_vocab}\n",
        "\n",
        "train_data = train_data.map(numericalize_example, fn_kwargs=fn_kwargs)\n",
        "valid_data = valid_data.map(numericalize_example, fn_kwargs=fn_kwargs)\n",
        "test_data = test_data.map(numericalize_example, fn_kwargs=fn_kwargs)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t8J6vulR6263"
      },
      "source": [
        "Checking an example, we can see that it has the two new features: \"en_ids\" and \"de_ids\", both a list of integers representing their indices in the respective vocabulary.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7GI4Fbqb6263",
        "outputId": "81324b12-3664-4b8b-eff0-6e9bca7f1ff2"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{'en': 'Two young, White males are outside near many bushes.',\n",
              " 'de': 'Zwei junge weiße Männer sind im Freien in der Nähe vieler Büsche.',\n",
              " 'en_tokens': ['<sos>',\n",
              "  'two',\n",
              "  'young',\n",
              "  ',',\n",
              "  'white',\n",
              "  'males',\n",
              "  'are',\n",
              "  'outside',\n",
              "  'near',\n",
              "  'many',\n",
              "  'bushes',\n",
              "  '.',\n",
              "  '<eos>'],\n",
              " 'de_tokens': ['<sos>',\n",
              "  'zwei',\n",
              "  'junge',\n",
              "  'weiße',\n",
              "  'männer',\n",
              "  'sind',\n",
              "  'im',\n",
              "  'freien',\n",
              "  'in',\n",
              "  'der',\n",
              "  'nähe',\n",
              "  'vieler',\n",
              "  'büsche',\n",
              "  '.',\n",
              "  '<eos>'],\n",
              " 'en_ids': [2, 16, 24, 15, 25, 778, 17, 57, 80, 202, 1312, 5, 3],\n",
              " 'de_ids': [2, 18, 26, 253, 30, 84, 20, 88, 7, 15, 110, 7647, 3171, 4, 3]}"
            ]
          },
          "execution_count": 32,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "train_data[0]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RpaEcbYe6264"
      },
      "source": [
        "We can confirm the indices are correct by using the `lookup_tokens` method with the corresponding vocabulary on the list of indices.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3da6Hfo26264",
        "outputId": "3e5ca162-73b4-41a2-c3a1-67082b83a48a"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "['<sos>',\n",
              " 'two',\n",
              " 'young',\n",
              " ',',\n",
              " 'white',\n",
              " 'males',\n",
              " 'are',\n",
              " 'outside',\n",
              " 'near',\n",
              " 'many',\n",
              " 'bushes',\n",
              " '.',\n",
              " '<eos>']"
            ]
          },
          "execution_count": 33,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "en_vocab.lookup_tokens(train_data[0][\"en_ids\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fJFUPbyl6264"
      },
      "source": [
        "One other thing that the `datasets` library handles for us with the `Dataset` class is converting features to the correct type. Our indices in each example are currently basic Python integers. However, they need to be converted to PyTorch tensors in order to use them with PyTorch. We could convert them just before we pass them into the model, however it is more convenient to do it now.\n",
        "\n",
        "The `with_format` method converts features indicated by the `columns` argument to a given `type`. Here, we specify the type as \"torch\" (for PyTorch) and the columns to be \"en_ids\" and \"de_ids\" (the features which we want to convert to PyTorch tensors). By default, `with_format` will remove any features not in the list of features passed to `columns`. We want to keep those features, which we can do with `output_all_columns=True`.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kt-_knCi6264"
      },
      "outputs": [],
      "source": [
        "data_type = \"torch\"\n",
        "format_columns = [\"en_ids\", \"de_ids\"]\n",
        "\n",
        "train_data = train_data.with_format(\n",
        "    type=data_type, columns=format_columns, output_all_columns=True\n",
        ")\n",
        "\n",
        "valid_data = valid_data.with_format(\n",
        "    type=data_type,\n",
        "    columns=format_columns,\n",
        "    output_all_columns=True,\n",
        ")\n",
        "\n",
        "test_data = test_data.with_format(\n",
        "    type=data_type,\n",
        "    columns=format_columns,\n",
        "    output_all_columns=True,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hN9-DeNy6264"
      },
      "source": [
        "We can confirm this worked by checking an example and seeing the \"en_ids\" and \"de_ids\" features are listed as `tensor([...])`.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HY9LXEgt6265",
        "outputId": "845a528a-0320-4b41-b385-f987a1276140"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{'en_ids': tensor([   2,   16,   24,   15,   25,  778,   17,   57,   80,  202, 1312,    5,\n",
              "            3]),\n",
              " 'de_ids': tensor([   2,   18,   26,  253,   30,   84,   20,   88,    7,   15,  110, 7647,\n",
              "         3171,    4,    3]),\n",
              " 'en': 'Two young, White males are outside near many bushes.',\n",
              " 'de': 'Zwei junge weiße Männer sind im Freien in der Nähe vieler Büsche.',\n",
              " 'en_tokens': ['<sos>',\n",
              "  'two',\n",
              "  'young',\n",
              "  ',',\n",
              "  'white',\n",
              "  'males',\n",
              "  'are',\n",
              "  'outside',\n",
              "  'near',\n",
              "  'many',\n",
              "  'bushes',\n",
              "  '.',\n",
              "  '<eos>'],\n",
              " 'de_tokens': ['<sos>',\n",
              "  'zwei',\n",
              "  'junge',\n",
              "  'weiße',\n",
              "  'männer',\n",
              "  'sind',\n",
              "  'im',\n",
              "  'freien',\n",
              "  'in',\n",
              "  'der',\n",
              "  'nähe',\n",
              "  'vieler',\n",
              "  'büsche',\n",
              "  '.',\n",
              "  '<eos>']}"
            ]
          },
          "execution_count": 35,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "train_data[0]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ewa7tvwv6265"
      },
      "source": [
        "We can also check this using the `type` built-in function on one of the features.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aeDjsOeB6265",
        "outputId": "08b0e548-e385-467d-9998-9d00f5bae905"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "torch.Tensor"
            ]
          },
          "execution_count": 36,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "type(train_data[0][\"en_ids\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N50vEbMv6265"
      },
      "source": [
        "## Data Loaders\n",
        "\n",
        "The final step of preparing the data is to create the data loaders. These can be iterated upon to return a batch of data, each batch being a dictionary containing the numericalized English and German sentences (which have also been padded) as PyTorch tensors.\n",
        "\n",
        "First, we need to create a function that collates, i.e. combines, a batch of examples into a batch. The `collate_fn` below takes a \"batch\" as input (a list of examples), we then separate out the English and German indices for each example in the batch, and pass each one to the `pad_sequence` function. `pad_sequence` takes a list of tensors, pads each one to the length of the longest tensor using the `padding_value` (which we set to `pad_index`, the index of our `<pad>` token) and then returns a `[max length, batch size]` shaped tensor, where `batch size` is the number of examples in the batch and `max length` is the length of the longest tensor in the batch. We put each tensor into a dictionary and then return it.\n",
        "\n",
        "The `get_collate_fn` takes in the padding token index and returns the `collate_fn` defined inside it. This technique, of defining a function inside another and returning it, is known as a [closure](<https://en.wikipedia.org/wiki/Closure_(computer_programming)>). It allows the `collate_fn` to continually use the value of `pad_index` it was created with without creating a class or using global variables.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NFimdXQj6266"
      },
      "outputs": [],
      "source": [
        "def get_collate_fn(pad_index):\n",
        "    def collate_fn(batch):\n",
        "        batch_en_ids = [example[\"en_ids\"] for example in batch]\n",
        "        batch_de_ids = [example[\"de_ids\"] for example in batch]\n",
        "        batch_en_ids = nn.utils.rnn.pad_sequence(batch_en_ids, padding_value=pad_index)\n",
        "        batch_de_ids = nn.utils.rnn.pad_sequence(batch_de_ids, padding_value=pad_index)\n",
        "        batch = {\n",
        "            \"en_ids\": batch_en_ids,\n",
        "            \"de_ids\": batch_de_ids,\n",
        "        }\n",
        "        return batch\n",
        "\n",
        "    return collate_fn"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bYNlgOBU6266"
      },
      "source": [
        "Next, we write the functions which give us our data loaders creating using PyTorch's `DataLoader` class.\n",
        "\n",
        "`get_data_loader` is created using a `Dataset`, the batch size, the padding token index (which is used for creating the batches in the `collate_fn`, and a boolean deciding if the examples should be shuffled at the time the data loader is iterated over.\n",
        "\n",
        "The batch size defines the maximum amount of examples within a batch. If the length of the dataset is not evenly divisible by the batch size then the last batch will be smaller.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5calMmOX6266"
      },
      "outputs": [],
      "source": [
        "def get_data_loader(dataset, batch_size, pad_index, shuffle=False):\n",
        "    collate_fn = get_collate_fn(pad_index)\n",
        "    data_loader = torch.utils.data.DataLoader(\n",
        "        dataset=dataset,\n",
        "        batch_size=batch_size,\n",
        "        collate_fn=collate_fn,\n",
        "        shuffle=shuffle,\n",
        "    )\n",
        "    return data_loader"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hIHh23Sy6266"
      },
      "source": [
        "Finally, we create our data loaders.\n",
        "\n",
        "To reduce the training time, we generally want to us the largest batch size possible. When using a GPU, this means using the largest batch size that will fit in GPU memory.\n",
        "\n",
        "Shuffling of data makes training more stable and potentially improves the final performance of the model, however only needs to be done on the training set. The metrics calculated for the validation and test set will be the same no matter what order the data is in.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_NP_7Ilp6266"
      },
      "outputs": [],
      "source": [
        "batch_size = 128\n",
        "\n",
        "train_data_loader = get_data_loader(train_data, batch_size, pad_index, shuffle=True)\n",
        "valid_data_loader = get_data_loader(valid_data, batch_size, pad_index)\n",
        "test_data_loader = get_data_loader(test_data, batch_size, pad_index)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hkyXRWnl6267"
      },
      "source": [
        "## Building the Model\n",
        "\n",
        "We'll be building our model in three parts. The encoder, the decoder and a seq2seq model that encapsulates the encoder and decoder and will provide a way to interface with each.\n",
        "\n",
        "### Encoder\n",
        "\n",
        "First, the encoder, a 2 layer LSTM. The paper we are implementing uses a 4-layer LSTM, but in the interest of training time we cut this down to 2-layers. The concept of multi-layer RNNs is easy to expand from 2 to 4 layers.\n",
        "\n",
        "For a multi-layer RNN, the input sentence, $X$, after being embedded goes into the first (bottom) layer of the RNN and hidden states, $H=\\{h_1, h_2, ..., h_T\\}$, output by this layer are used as inputs to the RNN in the layer above. Thus, representing each layer with a superscript, the hidden states in the first layer are given by:\n",
        "\n",
        "$$h_t^1 = \\text{EncoderRNN}^1(e(x_t), h_{t-1}^1)$$\n",
        "\n",
        "The hidden states in the second layer are given by:\n",
        "\n",
        "$$h_t^2 = \\text{EncoderRNN}^2(h_t^1, h_{t-1}^2)$$\n",
        "\n",
        "Using a multi-layer RNN also means we'll also need an initial hidden state as input per layer, $h_0^l$, and we will also output a context vector per layer, $z^l$.\n",
        "\n",
        "Without going into too much detail about LSTMs (see [this](https://colah.github.io/posts/2015-08-Understanding-LSTMs/) blog post to learn more about them), all we need to know is that they're a type of RNN which instead of just taking in a hidden state and returning a new hidden state per time-step, also take in and return a _cell state_, $c_t$, per time-step.\n",
        "\n",
        "$$\n",
        "\\begin{align*}\n",
        "h_t &= \\text{RNN}(e(x_t), h_{t-1})\\\\\n",
        "(h_t, c_t) &= \\text{LSTM}(e(x_t), h_{t-1}, c_{t-1})\n",
        "\\end{align*}\n",
        "$$\n",
        "\n",
        "We can just think of $c_t$ as another type of hidden state. Similar to $h_0^l$, $c_0^l$ will be initialized to a tensor of all zeros. Also, our context vector will now be both the final hidden state and the final cell state, i.e. $z^l = (h_T^l, c_T^l)$.\n",
        "\n",
        "Extending our multi-layer equations to LSTMs, we get:\n",
        "\n",
        "$$\n",
        "\\begin{align*}\n",
        "(h_t^1, c_t^1) &= \\text{EncoderLSTM}^1(e(x_t), (h_{t-1}^1, c_{t-1}^1))\\\\\n",
        "(h_t^2, c_t^2) &= \\text{EncoderLSTM}^2(h_t^1, (h_{t-1}^2, c_{t-1}^2))\n",
        "\\end{align*}\n",
        "$$\n",
        "\n",
        "Note how only our hidden state from the first layer is passed as input to the second layer, and not the cell state.\n",
        "\n",
        "So our encoder looks something like this:\n",
        "\n",
        "![](https://github.com/bentrevett/pytorch-seq2seq/blob/main/assets/seq2seq2.png?raw=1)\n",
        "\n",
        "We create this in code by making an `Encoder` module, which requires we inherit from `torch.nn.Module` and use the `super().__init__()` as some boilerplate code. The encoder takes the following arguments:\n",
        "\n",
        "-   `input_dim` is the size/dimensionality of the one-hot vectors that will be input to the encoder. This is equal to the input (source) vocabulary size.\n",
        "-   `embedding_dim` is the dimensionality of the embedding layer. This layer converts the one-hot vectors into dense vectors with `embedding_dim` dimensions.\n",
        "-   `hidden_dim` is the dimensionality of the hidden and cell states.\n",
        "-   `n_layers` is the number of layers in the RNN.\n",
        "-   `dropout` is the amount of dropout to use. This is a regularization parameter to prevent overfitting. Check out [this](https://www.coursera.org/lecture/deep-neural-network/understanding-dropout-YaGbR) for more details about dropout.\n",
        "\n",
        "We aren't going to discuss the embedding layer in detail during these tutorials. All we need to know is that there is a step before the words - technically, the indexes of the words - are passed into the RNN, where the words are transformed into vectors. To read more about word embeddings, check these articles: [1](https://monkeylearn.com/blog/word-embeddings-transform-text-numbers/), [2](http://p.migdal.pl/2017/01/06/king-man-woman-queen-why.html), [3](http://mccormickml.com/2016/04/19/word2vec-tutorial-the-skip-gram-model/), [4](http://mccormickml.com/2017/01/11/word2vec-tutorial-part-2-negative-sampling/).\n",
        "\n",
        "The embedding layer is created using `nn.Embedding`, the LSTM with `nn.LSTM` and a dropout layer with `nn.Dropout`. Check the PyTorch [documentation](https://pytorch.org/docs/stable/nn.html) for more about these.\n",
        "\n",
        "One thing to note is that the `dropout` argument to the LSTM is how much dropout to apply between the layers of a multi-layer RNN, i.e. between the hidden states output from layer $l$ and those same hidden states being used for the input of layer $l+1$.\n",
        "\n",
        "In the `forward` method, we pass in the source sentence, $X$, which is converted into dense vectors using the `embedding` layer, and then dropout is applied. These embeddings are then passed into the RNN. As we pass a whole sequence to the RNN, it will automatically do the recurrent calculation of the hidden states over the whole sequence for us! Notice that we do not pass an initial hidden or cell state to the RNN. This is because, as noted in the [documentation](https://pytorch.org/docs/stable/nn.html#torch.nn.LSTM), that if no hidden/cell state is passed to the RNN, it will automatically create an initial hidden/cell state as a tensor of all zeros.\n",
        "\n",
        "The RNN returns: `outputs` (the top-layer hidden state for each time-step), `hidden` (the final hidden state for each layer, $h_T$, stacked on top of each other) and `cell` (the final cell state for each layer, $c_T$, stacked on top of each other).\n",
        "\n",
        "As we only need the final hidden and cell states (to make our context vector), `forward` only returns `hidden` and `cell`.\n",
        "\n",
        "The sizes of each of the tensors is left as comments in the code. In this implementation `n_directions` will always be 1, however note that bidirectional RNNs (covered in tutorial 3) will have `n_directions` as 2.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "j_IoViGL6267"
      },
      "outputs": [],
      "source": [
        "class Encoder(nn.Module):\n",
        "    def __init__(self, input_dim, embedding_dim, hidden_dim, n_layers, dropout):\n",
        "        super().__init__()\n",
        "        self.hidden_dim = hidden_dim\n",
        "        self.n_layers = n_layers\n",
        "        self.embedding = nn.Embedding(input_dim, embedding_dim)\n",
        "        self.rnn = nn.LSTM(embedding_dim, hidden_dim, n_layers, dropout=dropout)\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "    def forward(self, src):\n",
        "        # src = [src length, batch size]\n",
        "        embedded = self.dropout(self.embedding(src))\n",
        "        # embedded = [src length, batch size, embedding dim]\n",
        "        outputs, (hidden, cell) = self.rnn(embedded)\n",
        "        # outputs = [src length, batch size, hidden dim * n directions]\n",
        "        # hidden = [n layers * n directions, batch size, hidden dim]\n",
        "        # cell = [n layers * n directions, batch size, hidden dim]\n",
        "        # outputs are always from the top hidden layer\n",
        "        return hidden, cell"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Kyk9zIkk6267"
      },
      "source": [
        "### Decoder\n",
        "\n",
        "Next, we'll build our decoder, which will also be a 2-layer (4 in the paper) LSTM.\n",
        "\n",
        "![](https://github.com/bentrevett/pytorch-seq2seq/blob/main/assets/seq2seq3.png?raw=1)\n",
        "\n",
        "The `Decoder` class does a single step of decoding, i.e. it ouputs single token per time-step. The first layer will receive a hidden and cell state from the previous time-step, $(s_{t-1}^1, c_{t-1}^1)$, and feeds it through the LSTM with the current embedded token, $y_t$, to produce a new hidden and cell state, $(s_t^1, c_t^1)$. The subsequent layers will use the hidden state from the layer below, $s_t^{l-1}$, and the previous hidden and cell states from their layer, $(s_{t-1}^l, c_{t-1}^l)$. This provides equations very similar to those in the encoder.\n",
        "\n",
        "$$\n",
        "\\begin{align*}\n",
        "(s_t^1, c_t^1) = \\text{DecoderLSTM}^1(d(y_t), (s_{t-1}^1, c_{t-1}^1))\\\\\n",
        "(s_t^2, c_t^2) = \\text{DecoderLSTM}^2(s_t^1, (s_{t-1}^2, c_{t-1}^2))\n",
        "\\end{align*}\n",
        "$$\n",
        "\n",
        "Remember that the initial hidden and cell states to our decoder are our context vectors, which are the final hidden and cell states of our encoder from the same layer, i.e. $(s_0^l,c_0^l)=z^l=(h_T^l,c_T^l)$.\n",
        "\n",
        "We then pass the hidden state from the top layer of the RNN, $s_t^L$, through a linear layer, $f$, to make a prediction of what the next token in the target (output) sequence should be, $\\hat{y}_{t+1}$.\n",
        "\n",
        "$$\\hat{y}_{t+1} = f(s_t^L)$$\n",
        "\n",
        "The arguments and initialization are similar to the `Encoder` class, except we now have an `output_dim` which is the size of the vocabulary for the output/target language. There is also the addition of the `Linear` layer, used to make the predictions from the top layer hidden state.\n",
        "\n",
        "Within the `forward` method, we accept a batch of input tokens, previous hidden states and previous cell states. As we are only decoding one token at a time, the input tokens will always have a sequence length of 1. We `unsqueeze` the input tokens to add a sentence length dimension of 1. Then, similar to the encoder, we pass through an embedding layer and apply dropout. This batch of embedded tokens is then passed into the RNN with the previous hidden and cell states. This produces an `output` (hidden state from the top layer of the RNN), a new `hidden` state (one for each layer, stacked on top of each other) and a new `cell` state (also one per layer, stacked on top of each other). We then pass the `output` (after getting rid of the sentence length dimension) through the linear layer to receive our `prediction`. We then return the `prediction`, the new `hidden` state and the new `cell` state.\n",
        "\n",
        "**Note**: as we always have a sequence length of 1, we could use `nn.LSTMCell`, instead of `nn.LSTM`, as it is designed to handle a batch of inputs that aren't necessarily in a sequence. `nn.LSTMCell` is just a single cell and `nn.LSTM` is a wrapper around potentially multiple cells. Using the `nn.LSTMCell` in this case would mean we don't have to `unsqueeze` to add a fake sequence length dimension, but we would need one `nn.LSTMCell` per layer in the decoder and to ensure each `nn.LSTMCell` receives the correct initial hidden state from the encoder. All of this makes the code less concise -- hence the decision to stick with the regular `nn.LSTM`.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "33PFyjCl6268"
      },
      "outputs": [],
      "source": [
        "class Decoder(nn.Module):\n",
        "    def __init__(self, output_dim, embedding_dim, hidden_dim, n_layers, dropout):\n",
        "        super().__init__()\n",
        "        self.output_dim = output_dim\n",
        "        self.hidden_dim = hidden_dim\n",
        "        self.n_layers = n_layers\n",
        "        self.embedding = nn.Embedding(output_dim, embedding_dim)\n",
        "        self.rnn = nn.LSTM(embedding_dim, hidden_dim, n_layers, dropout=dropout)\n",
        "        self.fc_out = nn.Linear(hidden_dim, output_dim)\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "\n",
        "    def forward(self, input, hidden, cell):\n",
        "        # input = [batch size]\n",
        "        # hidden = [n layers * n directions, batch size, hidden dim]\n",
        "        # cell = [n layers * n directions, batch size, hidden dim]\n",
        "        # n directions in the decoder will both always be 1, therefore:\n",
        "        # hidden = [n layers, batch size, hidden dim]\n",
        "        # context = [n layers, batch size, hidden dim]\n",
        "        input = input.unsqueeze(0)\n",
        "        # input = [1, batch size]\n",
        "        embedded = self.dropout(self.embedding(input))\n",
        "        # embedded = [1, batch size, embedding dim]\n",
        "        output, (hidden, cell) = self.rnn(embedded, (hidden, cell))\n",
        "        # output = [seq length, batch size, hidden dim * n directions]\n",
        "        # hidden = [n layers * n directions, batch size, hidden dim]\n",
        "        # cell = [n layers * n directions, batch size, hidden dim]\n",
        "        # seq length and n directions will always be 1 in this decoder, therefore:\n",
        "        # output = [1, batch size, hidden dim]\n",
        "        # hidden = [n layers, batch size, hidden dim]\n",
        "        # cell = [n layers, batch size, hidden dim]\n",
        "        prediction = self.fc_out(output.squeeze(0))\n",
        "        # prediction = [batch size, output dim]\n",
        "        return prediction, hidden, cell"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mEJndB9E6269"
      },
      "source": [
        "### Seq2Seq\n",
        "\n",
        "For the final part of the implemenetation, we'll implement the seq2seq model. This will handle:\n",
        "\n",
        "-   receiving the input/source sentence\n",
        "-   using the encoder to produce the context vectors\n",
        "-   using the decoder to produce the predicted output/target sentence\n",
        "\n",
        "Our full model will look like this:\n",
        "\n",
        "![](https://github.com/bentrevett/pytorch-seq2seq/blob/main/assets/seq2seq4.png?raw=1)\n",
        "\n",
        "The `Seq2Seq` model takes in an `Encoder`, `Decoder`, and a `device` (used to place tensors on the GPU, if it exists).\n",
        "\n",
        "For this implementation, we have to ensure that the number of layers and the hidden (and cell) dimensions are equal in the `Encoder` and `Decoder`. This is not always the case, we do not necessarily need the same number of layers or the same hidden dimension sizes in a sequence-to-sequence model. However, if we did something like having a different number of layers then we would need to make decisions about how this is handled. For example, if our encoder has 2 layers and our decoder only has 1, how is this handled? Do we average the two context vectors output by the decoder? Do we pass both through a linear layer? Do we only use the context vector from the highest layer? Etc.\n",
        "\n",
        "Our `forward` method takes the source sentence, target sentence and a teacher-forcing ratio. The teacher forcing ratio is used when training our model. When decoding, at each time-step we will predict what the next token in the target sequence will be from the previous tokens decoded, $\\hat{y}_{t+1}=f(s_t^L)$. With probability equal to the teaching forcing ratio (`teacher_forcing_ratio`) we will use the actual ground-truth next token in the sequence as the input to the decoder during the next time-step. However, with probability `1 - teacher_forcing_ratio`, we will use the token that the model predicted as the next input to the model, even if it doesn't match the actual next token in the sequence.\n",
        "\n",
        "The first thing we do in the `forward` method is to create an `outputs` tensor that will store all of our predictions, $\\hat{Y}$.\n",
        "\n",
        "We then feed the input/source sentence, `src`, into the encoder and receive out final hidden and cell states.\n",
        "\n",
        "The first input to the decoder is the start of sequence (`<sos>`) token. As our `trg` tensor already has the `<sos>` token appended (all the way back when we tokenized our English sentences) we get our $y_1$ by slicing into it. We know how long our target sentences should be (`trg_length`), so we loop that many times. The last token input into the decoder is the one **before** the `<eos>` token -- the `<eos>` token is never input into the decoder.\n",
        "\n",
        "During each iteration of the loop, we:\n",
        "\n",
        "-   pass the input, previous hidden and previous cell states ($y_t, s_{t-1}, c_{t-1}$) into the decoder\n",
        "-   receive a prediction, next hidden state and next cell state ($\\hat{y}_{t+1}, s_{t}, c_{t}$) from the decoder\n",
        "-   place our prediction, $\\hat{y}_{t+1}$/`output` in our tensor of predictions, $\\hat{Y}$/`outputs`\n",
        "-   decide if we are going to \"teacher force\" or not\n",
        "    -   if we do, the next `input` is the ground-truth next token in the sequence, $y_{t+1}$/`trg[t]`\n",
        "    -   if we don't, the next `input` is the predicted next token in the sequence, $\\hat{y}_{t+1}$/`top1`, which we get by doing an `argmax` over the output tensor\n",
        "\n",
        "Once we've made all of our predictions, we return our tensor full of predictions, $\\hat{Y}$/`outputs`.\n",
        "\n",
        "**Note**: our decoder loop starts at 1, not 0. This means the 0th element of our `outputs` tensor remains all zeros. So our `trg` and `outputs` look something like:\n",
        "\n",
        "$$\n",
        "\\begin{align*}\n",
        "\\text{trg} = [<sos>, &y_1, y_2, y_3, <eos>]\\\\\n",
        "\\text{outputs} = [0, &\\hat{y}_1, \\hat{y}_2, \\hat{y}_3, <eos>]\n",
        "\\end{align*}\n",
        "$$\n",
        "\n",
        "Later on when we calculate the loss, we cut off the first element of each tensor to get:\n",
        "\n",
        "$$\n",
        "\\begin{align*}\n",
        "\\text{trg} = [&y_1, y_2, y_3, <eos>]\\\\\n",
        "\\text{outputs} = [&\\hat{y}_1, \\hat{y}_2, \\hat{y}_3, <eos>]\n",
        "\\end{align*}\n",
        "$$\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ESqoaRvD626-"
      },
      "outputs": [],
      "source": [
        "class Seq2Seq(nn.Module):\n",
        "    def __init__(self, encoder, decoder, device):\n",
        "        super().__init__()\n",
        "        self.encoder = encoder\n",
        "        self.decoder = decoder\n",
        "        self.device = device\n",
        "        assert (\n",
        "            encoder.hidden_dim == decoder.hidden_dim\n",
        "        ), \"Hidden dimensions of encoder and decoder must be equal!\"\n",
        "        assert (\n",
        "            encoder.n_layers == decoder.n_layers\n",
        "        ), \"Encoder and decoder must have equal number of layers!\"\n",
        "\n",
        "    def forward(self, src, trg, teacher_forcing_ratio):\n",
        "        # src = [src length, batch size]\n",
        "        # trg = [trg length, batch size]\n",
        "        # teacher_forcing_ratio is probability to use teacher forcing\n",
        "        # e.g. if teacher_forcing_ratio is 0.75 we use ground-truth inputs 75% of the time\n",
        "        batch_size = trg.shape[1]\n",
        "        trg_length = trg.shape[0]\n",
        "        trg_vocab_size = self.decoder.output_dim\n",
        "        # tensor to store decoder outputs\n",
        "        outputs = torch.zeros(trg_length, batch_size, trg_vocab_size).to(self.device)\n",
        "        # last hidden state of the encoder is used as the initial hidden state of the decoder\n",
        "        hidden, cell = self.encoder(src)\n",
        "        # hidden = [n layers * n directions, batch size, hidden dim]\n",
        "        # cell = [n layers * n directions, batch size, hidden dim]\n",
        "        # first input to the decoder is the <sos> tokens\n",
        "        input = trg[0, :]\n",
        "        # input = [batch size]\n",
        "        for t in range(1, trg_length):\n",
        "            # insert input token embedding, previous hidden and previous cell states\n",
        "            # receive output tensor (predictions) and new hidden and cell states\n",
        "            output, hidden, cell = self.decoder(input, hidden, cell)\n",
        "            # output = [batch size, output dim]\n",
        "            # hidden = [n layers, batch size, hidden dim]\n",
        "            # cell = [n layers, batch size, hidden dim]\n",
        "            # place predictions in a tensor holding predictions for each token\n",
        "            outputs[t] = output\n",
        "            # decide if we are going to use teacher forcing or not\n",
        "            teacher_force = random.random() < teacher_forcing_ratio\n",
        "            # get the highest predicted token from our predictions\n",
        "            top1 = output.argmax(1)\n",
        "            # if teacher forcing, use actual next token as next input\n",
        "            # if not, use predicted token\n",
        "            input = trg[t] if teacher_force else top1\n",
        "            # input = [batch size]\n",
        "        return outputs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-OCacVLN626_"
      },
      "source": [
        "## Training the Model\n",
        "\n",
        "Now we have our model implemented, we can begin training it.\n",
        "\n",
        "### Model Initialization\n",
        "\n",
        "First, we'll initialize our model. As mentioned before, the input and output dimensions are defined by the size of the vocabulary. The embedding dimesions and dropout for the encoder and decoder can be different, but the number of layers and the size of the hidden/cell states must be the same.\n",
        "\n",
        "We then define the encoder, decoder and then our Seq2Seq model, which we place on the `device`. The `device` is used to tell PyTorch whether a model or a tensor should be processed on a GPU or CPU. The `torch.cuda.is_available()` function returns `True` if a GPU is detected on our machine. Thus, our model will be placed on the GPU, if we have one.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ayWk0_dr627B"
      },
      "outputs": [],
      "source": [
        "input_dim = len(de_vocab)\n",
        "output_dim = len(en_vocab)\n",
        "encoder_embedding_dim = 256\n",
        "decoder_embedding_dim = 256\n",
        "hidden_dim = 512\n",
        "n_layers = 2\n",
        "encoder_dropout = 0.5\n",
        "decoder_dropout = 0.5\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "\n",
        "encoder = Encoder(\n",
        "    input_dim,\n",
        "    encoder_embedding_dim,\n",
        "    hidden_dim,\n",
        "    n_layers,\n",
        "    encoder_dropout,\n",
        ")\n",
        "\n",
        "decoder = Decoder(\n",
        "    output_dim,\n",
        "    decoder_embedding_dim,\n",
        "    hidden_dim,\n",
        "    n_layers,\n",
        "    decoder_dropout,\n",
        ")\n",
        "\n",
        "model = Seq2Seq(encoder, decoder, device).to(device)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7Z3k9DYn627C"
      },
      "source": [
        "### Weight Initialization\n",
        "\n",
        "Next up is initializing the weights of our model. In the paper they state they initialize all weights from a uniform distribution between -0.08 and +0.08, i.e. $\\mathcal{U}(-0.08, 0.08)$.\n",
        "\n",
        "We initialize weights in PyTorch by creating a function which we `apply` to our model. When using `apply`, the `init_weights` function will be called on every module and sub-module within our model. For each module we loop through all of the parameters and sample them from a uniform distribution with `nn.init.uniform_`.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "u6kzu-Lk627D",
        "outputId": "05e6c721-9112-4d82-cdea-52e950d351cf"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "Seq2Seq(\n",
              "  (encoder): Encoder(\n",
              "    (embedding): Embedding(7853, 256)\n",
              "    (rnn): LSTM(256, 512, num_layers=2, dropout=0.5)\n",
              "    (dropout): Dropout(p=0.5, inplace=False)\n",
              "  )\n",
              "  (decoder): Decoder(\n",
              "    (embedding): Embedding(5893, 256)\n",
              "    (rnn): LSTM(256, 512, num_layers=2, dropout=0.5)\n",
              "    (fc_out): Linear(in_features=512, out_features=5893, bias=True)\n",
              "    (dropout): Dropout(p=0.5, inplace=False)\n",
              "  )\n",
              ")"
            ]
          },
          "execution_count": 44,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "def init_weights(m):\n",
        "    for name, param in m.named_parameters():\n",
        "        nn.init.uniform_(param.data, -0.08, 0.08)\n",
        "\n",
        "\n",
        "model.apply(init_weights)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2JQD8UkU627E"
      },
      "source": [
        "We can also count the number of parameters in our model.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5h3ZcD0-627E",
        "outputId": "2bf78d11-3d04-415a-dfa4-d1c124af92fd"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "The model has 13,898,501 trainable parameters\n"
          ]
        }
      ],
      "source": [
        "def count_parameters(model):\n",
        "    return sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
        "\n",
        "\n",
        "print(f\"The model has {count_parameters(model):,} trainable parameters\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "un15XkUb627E"
      },
      "source": [
        "### Optimizer\n",
        "\n",
        "We define our optimizer, which we use to update our parameters in the training loop. Check out [this](http://ruder.io/optimizing-gradient-descent/) post for information about different optimizers. Here, we'll use Adam.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "reNDDUoD627E"
      },
      "outputs": [],
      "source": [
        "optimizer = optim.Adam(model.parameters())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ENXPUvyd627E"
      },
      "source": [
        "### Loss Function\n",
        "\n",
        "Next, we define our loss function. The `CrossEntropyLoss` function calculates both the log softmax as well as the negative log-likelihood of our predictions.\n",
        "\n",
        "Our loss function calculates the average loss per token, however by passing the index of the `<pad>` token as the `ignore_index` argument we ignore the loss whenever the target token is a padding token.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0PRMkei5627F"
      },
      "outputs": [],
      "source": [
        "criterion = nn.CrossEntropyLoss(ignore_index=pad_index)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9FKsC1sb627F"
      },
      "source": [
        "### Training Loop\n",
        "\n",
        "Next, we'll define our training loop.\n",
        "\n",
        "First, we'll set the model into \"training mode\" with `model.train()`. This will turn on dropout (and batch normalization, which we aren't using) and then iterate through our data iterator.\n",
        "\n",
        "As stated before, our decoder loop starts at 1, not 0. This means the 0th element of our `outputs` tensor remains all zeros. So our `trg` and `outputs` look something like:\n",
        "\n",
        "$$\n",
        "\\begin{align*}\n",
        "\\text{trg} = [<sos>, &y_1, y_2, y_3, <eos>]\\\\\n",
        "\\text{outputs} = [0, &\\hat{y}_1, \\hat{y}_2, \\hat{y}_3, <eos>]\n",
        "\\end{align*}\n",
        "$$\n",
        "\n",
        "Here, when we calculate the loss, we cut off the first element of each tensor to get:\n",
        "\n",
        "$$\n",
        "\\begin{align*}\n",
        "\\text{trg} = [&y_1, y_2, y_3, <eos>]\\\\\n",
        "\\text{outputs} = [&\\hat{y}_1, \\hat{y}_2, \\hat{y}_3, <eos>]\n",
        "\\end{align*}\n",
        "$$\n",
        "\n",
        "At each iteration:\n",
        "\n",
        "-   get the source and target sentences from the batch, $X$ and $Y$\n",
        "-   zero the gradients calculated from the last batch\n",
        "-   feed the source and target into the model to get the output, $\\hat{Y}$\n",
        "-   as the loss function only works on 2d inputs with 1d targets we need to flatten each of them with `.view`\n",
        "    -   we slice off the first column of the output and target tensors as mentioned above\n",
        "-   calculate the gradients with `loss.backward()`\n",
        "-   clip the gradients to prevent them from exploding (a common issue in RNNs)\n",
        "-   update the parameters of our model by doing an optimizer step\n",
        "-   sum the loss value to a running total\n",
        "\n",
        "Finally, we return the loss that is averaged over all batches.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IHRIzgc7627F"
      },
      "outputs": [],
      "source": [
        "def train_fn(\n",
        "    model, data_loader, optimizer, criterion, clip, teacher_forcing_ratio, device\n",
        "):\n",
        "    model.train()\n",
        "    epoch_loss = 0\n",
        "    for i, batch in enumerate(data_loader):\n",
        "        src = batch[\"de_ids\"].to(device)\n",
        "        trg = batch[\"en_ids\"].to(device)\n",
        "        # src = [src length, batch size]\n",
        "        # trg = [trg length, batch size]\n",
        "        optimizer.zero_grad()\n",
        "        output = model(src, trg, teacher_forcing_ratio)\n",
        "        # output = [trg length, batch size, trg vocab size]\n",
        "        output_dim = output.shape[-1]\n",
        "        output = output[1:].view(-1, output_dim)\n",
        "        # output = [(trg length - 1) * batch size, trg vocab size]\n",
        "        trg = trg[1:].view(-1)\n",
        "        # trg = [(trg length - 1) * batch size]\n",
        "        loss = criterion(output, trg)\n",
        "        loss.backward()\n",
        "        torch.nn.utils.clip_grad_norm_(model.parameters(), clip)\n",
        "        optimizer.step()\n",
        "        epoch_loss += loss.item()\n",
        "    return epoch_loss / len(data_loader)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-efXPuP-627G"
      },
      "source": [
        "### Evaluation Loop\n",
        "\n",
        "Our evaluation loop is similar to our training loop, however as we aren't updating any parameters we don't need to pass an optimizer or a clip value.\n",
        "\n",
        "We must remember to set the model to evaluation mode with `model.eval()`. This will turn off dropout (and batch normalization, if used).\n",
        "\n",
        "We use the `with torch.no_grad()` block to ensure no gradients are calculated within the block. This reduces memory consumption and speeds things up.\n",
        "\n",
        "The iteration loop is similar (without the parameter updates), however we must ensure we turn teacher forcing off for evaluation. This will cause the model to only use it's own predictions to make further predictions within a sentence, which mirrors how it would be used in deployment.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "h22RxrgJ627G"
      },
      "outputs": [],
      "source": [
        "def evaluate_fn(model, data_loader, criterion, device):\n",
        "    model.eval()\n",
        "    epoch_loss = 0\n",
        "    with torch.no_grad():\n",
        "        for i, batch in enumerate(data_loader):\n",
        "            src = batch[\"de_ids\"].to(device)\n",
        "            trg = batch[\"en_ids\"].to(device)\n",
        "            # src = [src length, batch size]\n",
        "            # trg = [trg length, batch size]\n",
        "            output = model(src, trg, 0)  # turn off teacher forcing\n",
        "            # output = [trg length, batch size, trg vocab size]\n",
        "            output_dim = output.shape[-1]\n",
        "            output = output[1:].view(-1, output_dim)\n",
        "            # output = [(trg length - 1) * batch size, trg vocab size]\n",
        "            trg = trg[1:].view(-1)\n",
        "            # trg = [(trg length - 1) * batch size]\n",
        "            loss = criterion(output, trg)\n",
        "            epoch_loss += loss.item()\n",
        "    return epoch_loss / len(data_loader)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uieqPRTt627G"
      },
      "source": [
        "### Model Training\n",
        "\n",
        "We can finally start training our model!\n",
        "\n",
        "At each epoch, we'll be checking if our model has achieved the best validation loss so far. If it has, we'll update our best validation loss and save the parameters of our model (called `state_dict` in PyTorch). Then, when we come to test our model, we'll use the saved parameters used to achieve the best validation loss.\n",
        "\n",
        "We'll be printing out both the loss and the perplexity at each epoch. It is easier to see a change in perplexity than a change in loss as the numbers are much bigger.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EAc0I0hI627G",
        "outputId": "13a6c838-9605-46f5-a46e-264365ce8194"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            " 10%|██████████████▋                                                                                                                                    | 1/10 [00:15<02:21, 15.76s/it]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\tTrain Loss:   5.035 | Train PPL: 153.632\n",
            "\tValid Loss:   4.991 | Valid PPL: 147.028\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            " 20%|█████████████████████████████▍                                                                                                                     | 2/10 [00:31<02:05, 15.71s/it]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\tTrain Loss:   4.399 | Train PPL:  81.330\n",
            "\tValid Loss:   4.661 | Valid PPL: 105.741\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            " 30%|████████████████████████████████████████████                                                                                                       | 3/10 [00:47<01:50, 15.73s/it]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\tTrain Loss:   4.075 | Train PPL:  58.860\n",
            "\tValid Loss:   4.487 | Valid PPL:  88.884\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            " 40%|██████████████████████████████████████████████████████████▊                                                                                        | 4/10 [01:02<01:33, 15.62s/it]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\tTrain Loss:   3.873 | Train PPL:  48.071\n",
            "\tValid Loss:   4.274 | Valid PPL:  71.819\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            " 50%|█████████████████████████████████████████████████████████████████████████▌                                                                         | 5/10 [01:18<01:18, 15.63s/it]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\tTrain Loss:   3.688 | Train PPL:  39.972\n",
            "\tValid Loss:   4.209 | Valid PPL:  67.278\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            " 60%|████████████████████████████████████████████████████████████████████████████████████████▏                                                          | 6/10 [01:33<01:02, 15.64s/it]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\tTrain Loss:   3.541 | Train PPL:  34.492\n",
            "\tValid Loss:   4.075 | Valid PPL:  58.835\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            " 70%|██████████████████████████████████████████████████████████████████████████████████████████████████████▉                                            | 7/10 [01:49<00:46, 15.65s/it]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\tTrain Loss:   3.404 | Train PPL:  30.074\n",
            "\tValid Loss:   4.034 | Valid PPL:  56.458\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            " 80%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▌                             | 8/10 [02:05<00:31, 15.68s/it]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\tTrain Loss:   3.248 | Train PPL:  25.751\n",
            "\tValid Loss:   3.987 | Valid PPL:  53.888\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            " 90%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████▎              | 9/10 [02:20<00:15, 15.65s/it]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\tTrain Loss:   3.109 | Train PPL:  22.396\n",
            "\tValid Loss:   3.894 | Valid PPL:  49.093\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 10/10 [02:36<00:00, 15.66s/it]"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\tTrain Loss:   2.986 | Train PPL:  19.802\n",
            "\tValid Loss:   3.787 | Valid PPL:  44.137\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "\n"
          ]
        }
      ],
      "source": [
        "n_epochs = 10\n",
        "clip = 1.0\n",
        "teacher_forcing_ratio = 0.5\n",
        "\n",
        "best_valid_loss = float(\"inf\")\n",
        "\n",
        "for epoch in tqdm.tqdm(range(n_epochs)):\n",
        "    train_loss = train_fn(\n",
        "        model,\n",
        "        train_data_loader,\n",
        "        optimizer,\n",
        "        criterion,\n",
        "        clip,\n",
        "        teacher_forcing_ratio,\n",
        "        device,\n",
        "    )\n",
        "    valid_loss = evaluate_fn(\n",
        "        model,\n",
        "        valid_data_loader,\n",
        "        criterion,\n",
        "        device,\n",
        "    )\n",
        "    if valid_loss < best_valid_loss:\n",
        "        best_valid_loss = valid_loss\n",
        "        torch.save(model.state_dict(), \"tut1-model.pt\")\n",
        "    print(f\"\\tTrain Loss: {train_loss:7.3f} | Train PPL: {np.exp(train_loss):7.3f}\")\n",
        "    print(f\"\\tValid Loss: {valid_loss:7.3f} | Valid PPL: {np.exp(valid_loss):7.3f}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qUCp92B9627H"
      },
      "source": [
        "We've now successfully trained a model that translates German into English! But how well does it perform?\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CtULTx3W627H"
      },
      "source": [
        "## Evaluating the Model\n",
        "\n",
        "The first thing to do is to test the model's performance on the test set.\n",
        "\n",
        "We'll load the parameters (`state_dict`) that gave our model the best validation loss and run it on the test set to get our test loss and perplexity.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sWACYOZX627H",
        "outputId": "e3b56169-4052-4f56-dbb5-abb77c2ee381"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "| Test Loss: 3.780 | Test PPL:  43.833 |\n"
          ]
        }
      ],
      "source": [
        "model.load_state_dict(torch.load(\"tut1-model.pt\"))\n",
        "\n",
        "test_loss = evaluate_fn(model, test_data_loader, criterion, device)\n",
        "\n",
        "print(f\"| Test Loss: {test_loss:.3f} | Test PPL: {np.exp(test_loss):7.3f} |\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "isj3jr0V627H"
      },
      "source": [
        "Pretty similar to the validation performance, which is a good sign. It means we aren't overfitting on the validation set.\n",
        "\n",
        "You might think it's impossible to overfit on the validation set, but it's not. Every time you tweak your hyperparameters (e.g. optimizer, learning rate, model architecture, weight initialization, etc.) in order to get better results on the validation set, you are slowly overfitting those hyperparameters to your validation set. You can also do this on the test set too! Hence, you should evaluate your model on your test set as few times as possible.\n",
        "\n",
        "Most papers using neural networks for translation don't give their results in terms of loss and perplexity on the test set, they usually give the [BLEU](https://en.wikipedia.org/wiki/BLEU) score. Unlike loss/perplexity, BLEU is a value between zero and one, where higher is better, and [according to the original BLEU paper](https://aclanthology.org/P02-1040.pdf) it has a high correlation with human judgement.\n",
        "\n",
        "To get our model's BLEU score on the test set, we first need to use our model to translate every example from our test set, which we do with the `translate_sentence` function below. The function first converts the input `sentence` into `tokens`, optionally lowercases each `token`, and then appends the start and end of sequence tokens, `sos_token` and `eos_token`. It then uses the vocabulary to numericalize the tokens into `ids` and converts these `ids` into a `tensor`, adding a \"fake\" batch dimension, and then passes the `tensor` through the `encoder` to get the `hidden` and `cell` states. We then perform the decoding, starting with the `sos_token`, converting it into a tensor, passing it through the `decoder`, getting the `predicted_token` our model thinks is most likely to be next in the sequence, which we append to our list of `inputs` to the decoder. If the `predicted_token` is the end of sequence token then we stop decoding, if not we continue the loop, using the `predicted_token` as the next input to the decoder. We keep decoding until the decoder outputs the `eos_token` or we hit `max_output_length` (which we use to avoid the decoder just generating tokens forever). Once we've stopped decoding, we convert out inputs into `tokens` using our vocabulary and return them.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "i11WCcpb627I"
      },
      "outputs": [],
      "source": [
        "def translate_sentence(\n",
        "    sentence,\n",
        "    model,\n",
        "    en_nlp,\n",
        "    de_nlp,\n",
        "    en_vocab,\n",
        "    de_vocab,\n",
        "    lower,\n",
        "    sos_token,\n",
        "    eos_token,\n",
        "    device,\n",
        "    max_output_length=25,\n",
        "):\n",
        "    model.eval()\n",
        "    with torch.no_grad():\n",
        "        if isinstance(sentence, str):\n",
        "            tokens = [token.text for token in de_nlp.tokenizer(sentence)]\n",
        "        else:\n",
        "            tokens = [token for token in sentence]\n",
        "        if lower:\n",
        "            tokens = [token.lower() for token in tokens]\n",
        "        tokens = [sos_token] + tokens + [eos_token]\n",
        "        ids = de_vocab.lookup_indices(tokens)\n",
        "        tensor = torch.LongTensor(ids).unsqueeze(-1).to(device)\n",
        "        hidden, cell = model.encoder(tensor)\n",
        "        inputs = en_vocab.lookup_indices([sos_token])\n",
        "        for _ in range(max_output_length):\n",
        "            inputs_tensor = torch.LongTensor([inputs[-1]]).to(device)\n",
        "            output, hidden, cell = model.decoder(inputs_tensor, hidden, cell)\n",
        "            predicted_token = output.argmax(-1).item()\n",
        "            inputs.append(predicted_token)\n",
        "            if predicted_token == en_vocab[eos_token]:\n",
        "                break\n",
        "        tokens = en_vocab.lookup_tokens(inputs)\n",
        "    return tokens"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4DChroPC627I"
      },
      "source": [
        "We'll pass a test example (something the model hasn't been trained on) to use as a sentence to test our `translate_sentence` function, passing in the German sentence and expecting to get something that looks like the English sentence.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "PCH4pfsY627I",
        "outputId": "8a2f7976-df00-4a58-b419-b16cd6073feb"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "('Ein Mann mit einem orangefarbenen Hut, der etwas anstarrt.',\n",
              " 'A man in an orange hat starring at something.')"
            ]
          },
          "execution_count": 53,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "sentence = test_data[0][\"de\"]\n",
        "expected_translation = test_data[0][\"en\"]\n",
        "\n",
        "sentence, expected_translation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dAbeOMU8627J"
      },
      "outputs": [],
      "source": [
        "translation = translate_sentence(\n",
        "    sentence,\n",
        "    model,\n",
        "    en_nlp,\n",
        "    de_nlp,\n",
        "    en_vocab,\n",
        "    de_vocab,\n",
        "    lower,\n",
        "    sos_token,\n",
        "    eos_token,\n",
        "    device,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NsGDgwGW627J"
      },
      "source": [
        "Our model has seemed to have figured out that the input sentence mentions a man wearing an item of clothing (though gets both the color and the item wrong), but it can't seem to figure out what the man is doing.\n",
        "\n",
        "We shouldn't be expecting amazing results, our model is relatively small to what is used in the paper we're implementing (they use four layers with embedding and hidden dimensions of 1000) and is miniscule compared to modern translation models (which have billions of parameters).\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "tdarpk-c627J",
        "outputId": "aaaeed1d-3393-4593-f77d-d66e628e0c59"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "['<sos>',\n",
              " 'a',\n",
              " 'man',\n",
              " 'in',\n",
              " 'a',\n",
              " 'white',\n",
              " 'shirt',\n",
              " 'is',\n",
              " 'cutting',\n",
              " 'something',\n",
              " '.',\n",
              " '<eos>']"
            ]
          },
          "execution_count": 55,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "translation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HrLkhsip627J"
      },
      "source": [
        "The model doesn't just translate examples in the training, validation and test sets. We can use it to translate arbitrary sentences by passing any string to the `translate_sentence`.\n",
        "\n",
        "Note that the multi30k dataset consists of image captions that have been translated from English to German, and our model has been trained to translate German to English. Therefore, the model will only output reasonable translations if the sentences are German sentences that could potentially be image captions. (It's also important to re-iterate that the model trained here is relatively small and the translation performance will generally be poor.)\n",
        "\n",
        "Below, we input the German translation of **\"A man is watching a film.\"**\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "J_oAPQMC627J"
      },
      "outputs": [],
      "source": [
        "sentence = \"Ein Mann sitzt auf einer Bank.\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ShhlSiPz627J"
      },
      "outputs": [],
      "source": [
        "translation = translate_sentence(\n",
        "    sentence,\n",
        "    model,\n",
        "    en_nlp,\n",
        "    de_nlp,\n",
        "    en_vocab,\n",
        "    de_vocab,\n",
        "    lower,\n",
        "    sos_token,\n",
        "    eos_token,\n",
        "    device,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Kqjog-at627K"
      },
      "source": [
        "And we receive our translation, which is reasonably close.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sjtqhsI-627K",
        "outputId": "21b91469-89a8-4bbc-b503-7d61e87be1ea"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "['<sos>', 'man', 'sitting', 'on', 'a', 'bench', '.', '<eos>']"
            ]
          },
          "execution_count": 58,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "translation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ih4FM2NP627K"
      },
      "source": [
        "We can now loop over our `test_data`, getting our model's translation of each test sentence.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XPoq2CPs627K",
        "outputId": "712f6fc5-7cb1-48a3-a84f-cbf83aec0a34"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:04<00:00, 235.22it/s]\n"
          ]
        }
      ],
      "source": [
        "translations = [\n",
        "    translate_sentence(\n",
        "        example[\"de\"],\n",
        "        model,\n",
        "        en_nlp,\n",
        "        de_nlp,\n",
        "        en_vocab,\n",
        "        de_vocab,\n",
        "        lower,\n",
        "        sos_token,\n",
        "        eos_token,\n",
        "        device,\n",
        "    )\n",
        "    for example in tqdm.tqdm(test_data)\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ILnHaA57627L"
      },
      "source": [
        "To calculate BLEU, we'll use the `evaluate` library. It's recommended to use libraries for measuring metrics to ensure there are now bugs in your metric calculations and giving you potentially incorrect results.\n",
        "\n",
        "The BLEU metric can be loaded from the `evaluate` library like so:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jGyPRD4M627L"
      },
      "outputs": [],
      "source": [
        "bleu = evaluate.load(\"bleu\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Nj8KYoNf627L"
      },
      "source": [
        "One quirk one the BLEU metric is that it expects the predictions (predicted translations) to be strings and the references (actual English sentences) to be a list of sentences. This is because BLEU works if you have multiple correct sentences per prediction as there may be potentially be multiple ways to translate a sentence. In our case, we only have a single reference sentence so we just wrap our target sentence in a list. We also convert our translations from a list of tokens into a string by joining them with whitespace inbetween and getting rid of the `<sos>` and `<eos>` tokens (as they will never appear in our reference sentences).\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jPkdsGR9627L"
      },
      "outputs": [],
      "source": [
        "predictions = [\" \".join(translation[1:-1]) for translation in translations]\n",
        "\n",
        "references = [[example[\"en\"]] for example in test_data]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9fdfu3k3627L",
        "outputId": "50f51997-cb55-4b60-8e8c-135bd5d9965f"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "('a man in a white shirt is cutting something .',\n",
              " ['A man in an orange hat starring at something.'])"
            ]
          },
          "execution_count": 62,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "predictions[0], references[0]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KrgYRPpB627L"
      },
      "source": [
        "We also need to define a function which tokenizes an input string. This will be used to calculate the BLEU score by comparing our predicted tokens against the reference tokens.\n",
        "\n",
        "It seems a bit odd that we joined our translated tokens together into a string only to just tokenize them again, and also used the English string from our test data instead of the existing tokens (`en_tokens`), however this is another quirk of the BLEU metric provided by the `evaluate` library; the `predictions` and `references` must be strings and not tokens, and that we must tell the metric how these strings should be tokenized.\n",
        "\n",
        "The `get_tokenize_fn` returns our `tokenizer_fn`, which uses our `spaCy` tokenizer and lowercases tokens if necessary.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "P5LXYxKk627M"
      },
      "outputs": [],
      "source": [
        "def get_tokenizer_fn(nlp, lower):\n",
        "    def tokenizer_fn(s):\n",
        "        tokens = [token.text for token in nlp.tokenizer(s)]\n",
        "        if lower:\n",
        "            tokens = [token.lower() for token in tokens]\n",
        "        return tokens\n",
        "\n",
        "    return tokenizer_fn"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CSGlXzv1627M"
      },
      "outputs": [],
      "source": [
        "tokenizer_fn = get_tokenizer_fn(en_nlp, lower)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OnqTO6ZO627M",
        "outputId": "12e23411-b59a-4f8c-cdc5-067b79a2ee4c"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "(['a', 'man', 'in', 'a', 'white', 'shirt', 'is', 'cutting', 'something', '.'],\n",
              " ['a', 'man', 'in', 'an', 'orange', 'hat', 'starring', 'at', 'something', '.'])"
            ]
          },
          "execution_count": 65,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "tokenizer_fn(predictions[0]), tokenizer_fn(references[0][0])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-b6xklqM627M"
      },
      "source": [
        "Finally, we calculate the BLEU metric across our test set!\n",
        "\n",
        "We pass our `predictions`, `references` and our `tokenizer_fn` to the `compute` method of the BLEU metric to get our results.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MyLZXCGz627M"
      },
      "outputs": [],
      "source": [
        "results = bleu.compute(\n",
        "    predictions=predictions, references=references, tokenizer=tokenizer_fn\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sZTNH_SZ627N",
        "outputId": "1fc1bd8f-2471-43a0-d68f-8788ee781791"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "{'bleu': 0.13756994726803526,\n",
              " 'precisions': [0.4681576952236543,\n",
              "  0.18843314191960622,\n",
              "  0.09133154602323502,\n",
              "  0.04445534838076546],\n",
              " 'brevity_penalty': 1.0,\n",
              " 'length_ratio': 1.0101087455965692,\n",
              " 'translation_length': 13190,\n",
              " 'reference_length': 13058}"
            ]
          },
          "execution_count": 67,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "results"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4BT5LqnC627N"
      },
      "source": [
        "We get a BLEU score of 0.14! Nothing to write home about, but not bad for our first translation model.\n",
        "\n",
        "In the subsequent notebooks we'll be implementing more translation papers and slowly increasing the BLEU score achieved.\n"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3 (ipykernel)",
      "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.9.12"
    },
    "colab": {
      "provenance": []
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}