notebooks/training-nw-nb.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f2aeea51",
   "metadata": {
    "_cell_guid": "b1076dfc-b9ad-4769-8c92-a6c4dae69d19",
    "_uuid": "8f2839f25d086af736a60e9eeb907d3b93b6e0e5",
    "papermill": {
     "duration": 0.018113,
     "end_time": "2024-01-23T10:24:53.640885",
     "exception": false,
     "start_time": "2024-01-23T10:24:53.622772",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Training model for Machine Translation\n",
    "\n",
    "Library Imports for the jupyter notebook. We are refering to this [blog](https://medium.com/@hunter-j-phillips/putting-it-all-together-the-implemented-transformer-bfb11ac1ddfehttps://medium.com/@hunter-j-phillips/putting-it-all-together-the-implemented-transformer-bfb11ac1ddfe) to understand attention network in depth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "ec700df6",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:24:53.677228Z",
     "iopub.status.busy": "2024-01-23T10:24:53.676883Z",
     "iopub.status.idle": "2024-01-23T10:25:14.521412Z",
     "shell.execute_reply": "2024-01-23T10:25:14.520418Z"
    },
    "papermill": {
     "duration": 20.86528,
     "end_time": "2024-01-23T10:25:14.523614",
     "exception": false,
     "start_time": "2024-01-23T10:24:53.658334",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.10/site-packages/scipy/__init__.py:146: UserWarning: A NumPy version >=1.16.5 and <1.23.0 is required for this version of SciPy (detected version 1.24.3\n",
      "  warnings.warn(f\"A NumPy version >={np_minversion} and <{np_maxversion}\"\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "device(type='cuda')"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "!pip install -q portalocker\n",
    "\n",
    "# importing required libraries\n",
    "import math\n",
    "import copy\n",
    "import time\n",
    "import random\n",
    "import spacy\n",
    "import numpy as np\n",
    "import os\n",
    "\n",
    "# torch packages\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch import Tensor\n",
    "import torch.optim as optim\n",
    "\n",
    "# load and build datasets\n",
    "import torchtext\n",
    "from torchtext.data.functional import to_map_style_dataset\n",
    "from torch.nn.functional import pad\n",
    "from torch.utils.data import DataLoader\n",
    "from torchtext.vocab import build_vocab_from_iterator\n",
    "import torchtext.datasets as datasets\n",
    "import portalocker\n",
    "\n",
    "# visualization packages\n",
    "from mpl_toolkits import mplot3d\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.ticker as ticker\n",
    "from torchinfo import summary\n",
    "\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "device"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "076fe546",
   "metadata": {
    "papermill": {
     "duration": 0.017701,
     "end_time": "2024-01-23T10:25:14.560539",
     "exception": false,
     "start_time": "2024-01-23T10:25:14.542838",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Reading the dataframe and converting it into iterable for consuming in pytorch dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9c006dab",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:25:14.598136Z",
     "iopub.status.busy": "2024-01-23T10:25:14.597246Z",
     "iopub.status.idle": "2024-01-23T10:25:14.601704Z",
     "shell.execute_reply": "2024-01-23T10:25:14.600817Z"
    },
    "papermill": {
     "duration": 0.025355,
     "end_time": "2024-01-23T10:25:14.603594",
     "exception": false,
     "start_time": "2024-01-23T10:25:14.578239",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d4400ac0",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:25:14.640538Z",
     "iopub.status.busy": "2024-01-23T10:25:14.640209Z",
     "iopub.status.idle": "2024-01-23T10:25:14.808002Z",
     "shell.execute_reply": "2024-01-23T10:25:14.806999Z"
    },
    "papermill": {
     "duration": 0.188658,
     "end_time": "2024-01-23T10:25:14.810032",
     "exception": false,
     "start_time": "2024-01-23T10:25:14.621374",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>english</th>\n",
       "      <th>german</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Two young, White males are outside near many b...</td>\n",
       "      <td>Zwei junge weiße Männer sind im Freien in der ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Several men in hard hats are operating a giant...</td>\n",
       "      <td>Mehrere Männer mit Schutzhelmen bedienen ein A...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>A little girl climbing into a wooden playhouse.</td>\n",
       "      <td>Ein kleines Mädchen klettert in ein Spielhaus ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>A man in a blue shirt is standing on a ladder ...</td>\n",
       "      <td>Ein Mann in einem blauen Hemd steht auf einer ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Two men are at the stove preparing food.</td>\n",
       "      <td>Zwei Männer stehen am Herd und bereiten Essen zu.</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                             english  \\\n",
       "0  Two young, White males are outside near many b...   \n",
       "1  Several men in hard hats are operating a giant...   \n",
       "2    A little girl climbing into a wooden playhouse.   \n",
       "3  A man in a blue shirt is standing on a ladder ...   \n",
       "4           Two men are at the stove preparing food.   \n",
       "\n",
       "                                              german  \n",
       "0  Zwei junge weiße Männer sind im Freien in der ...  \n",
       "1  Mehrere Männer mit Schutzhelmen bedienen ein A...  \n",
       "2  Ein kleines Mädchen klettert in ein Spielhaus ...  \n",
       "3  Ein Mann in einem blauen Hemd steht auf einer ...  \n",
       "4  Zwei Männer stehen am Herd und bereiten Essen zu.  "
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = pd.read_csv(\n",
    "    \"/kaggle/input/machine-translation-dataset-de-en/translation_train.csv\"\n",
    ")\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "2bf5dbae",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:25:14.847688Z",
     "iopub.status.busy": "2024-01-23T10:25:14.847355Z",
     "iopub.status.idle": "2024-01-23T10:25:14.852814Z",
     "shell.execute_reply": "2024-01-23T10:25:14.851917Z"
    },
    "papermill": {
     "duration": 0.026333,
     "end_time": "2024-01-23T10:25:14.854631",
     "exception": false,
     "start_time": "2024-01-23T10:25:14.828298",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(29000, 2)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "9206d5ef",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:25:14.891779Z",
     "iopub.status.busy": "2024-01-23T10:25:14.891502Z",
     "iopub.status.idle": "2024-01-23T10:25:14.903827Z",
     "shell.execute_reply": "2024-01-23T10:25:14.902944Z"
    },
    "papermill": {
     "duration": 0.033206,
     "end_time": "2024-01-23T10:25:14.905734",
     "exception": false,
     "start_time": "2024-01-23T10:25:14.872528",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Perform train - val split\n",
    "train_df = df.sample(frac=0.95, random_state=200)\n",
    "val_df = df.drop(train_df.index)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "acd57d86",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:25:14.943589Z",
     "iopub.status.busy": "2024-01-23T10:25:14.943307Z",
     "iopub.status.idle": "2024-01-23T10:25:14.948573Z",
     "shell.execute_reply": "2024-01-23T10:25:14.947746Z"
    },
    "papermill": {
     "duration": 0.025894,
     "end_time": "2024-01-23T10:25:14.950450",
     "exception": false,
     "start_time": "2024-01-23T10:25:14.924556",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def create_iterable(sample_df):\n",
    "    \"\"\"\n",
    "    This converts pandas dataframe into list of tuples\n",
    "    consisting of (german sentences, english sentences).\n",
    "\n",
    "    This iterable in used in our data preparation\n",
    "    \"\"\"\n",
    "    sample_iter = sample_df.to_dict(orient=\"records\")\n",
    "    out_iter = [(dict1[\"german\"], dict1[\"english\"]) for dict1 in sample_iter]\n",
    "    print(f\"length of iterable: {len(out_iter)}\")\n",
    "    return out_iter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "6bbf8401",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:25:15.028313Z",
     "iopub.status.busy": "2024-01-23T10:25:15.028019Z",
     "iopub.status.idle": "2024-01-23T10:25:15.341229Z",
     "shell.execute_reply": "2024-01-23T10:25:15.340295Z"
    },
    "papermill": {
     "duration": 0.375161,
     "end_time": "2024-01-23T10:25:15.343415",
     "exception": false,
     "start_time": "2024-01-23T10:25:14.968254",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "length of iterable: 27550\n",
      "length of iterable: 1450\n"
     ]
    }
   ],
   "source": [
    "train_iter = create_iterable(train_df)\n",
    "val_iter = create_iterable(val_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "3055d58c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:25:15.381926Z",
     "iopub.status.busy": "2024-01-23T10:25:15.381296Z",
     "iopub.status.idle": "2024-01-23T10:25:15.398534Z",
     "shell.execute_reply": "2024-01-23T10:25:15.397632Z"
    },
    "papermill": {
     "duration": 0.038555,
     "end_time": "2024-01-23T10:25:15.400557",
     "exception": false,
     "start_time": "2024-01-23T10:25:15.362002",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>english</th>\n",
       "      <th>german</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>A man in an orange hat starring at something.</td>\n",
       "      <td>Ein Mann mit einem orangefarbenen Hut, der etw...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>A Boston Terrier is running on lush green gras...</td>\n",
       "      <td>Ein Boston Terrier läuft über saftig-grünes Gr...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>A girl in karate uniform breaking a stick with...</td>\n",
       "      <td>Ein Mädchen in einem Karateanzug bricht einen ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Five people wearing winter jackets and helmets...</td>\n",
       "      <td>Fünf Leute in Winterjacken und mit Helmen steh...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>People are fixing the roof of a house.</td>\n",
       "      <td>Leute Reparieren das Dach eines Hauses.</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                             english  \\\n",
       "0      A man in an orange hat starring at something.   \n",
       "1  A Boston Terrier is running on lush green gras...   \n",
       "2  A girl in karate uniform breaking a stick with...   \n",
       "3  Five people wearing winter jackets and helmets...   \n",
       "4             People are fixing the roof of a house.   \n",
       "\n",
       "                                              german  \n",
       "0  Ein Mann mit einem orangefarbenen Hut, der etw...  \n",
       "1  Ein Boston Terrier läuft über saftig-grünes Gr...  \n",
       "2  Ein Mädchen in einem Karateanzug bricht einen ...  \n",
       "3  Fünf Leute in Winterjacken und mit Helmen steh...  \n",
       "4            Leute Reparieren das Dach eines Hauses.  "
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_df = pd.read_csv(\n",
    "    \"/kaggle/input/machine-translation-dataset-de-en/translation_test.csv\"\n",
    ")\n",
    "test_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "35b85e1d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:25:15.439797Z",
     "iopub.status.busy": "2024-01-23T10:25:15.438995Z",
     "iopub.status.idle": "2024-01-23T10:25:15.448712Z",
     "shell.execute_reply": "2024-01-23T10:25:15.447703Z"
    },
    "papermill": {
     "duration": 0.031113,
     "end_time": "2024-01-23T10:25:15.450640",
     "exception": false,
     "start_time": "2024-01-23T10:25:15.419527",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "length of iterable: 1000\n"
     ]
    }
   ],
   "source": [
    "test_iter = create_iterable(test_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d46bbb5a",
   "metadata": {
    "papermill": {
     "duration": 0.018436,
     "end_time": "2024-01-23T10:25:15.487772",
     "exception": false,
     "start_time": "2024-01-23T10:25:15.469336",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Create Vocab from data sources"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "3c2fa2ae",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:25:15.526619Z",
     "iopub.status.busy": "2024-01-23T10:25:15.525963Z",
     "iopub.status.idle": "2024-01-23T10:25:15.531892Z",
     "shell.execute_reply": "2024-01-23T10:25:15.531044Z"
    },
    "papermill": {
     "duration": 0.027563,
     "end_time": "2024-01-23T10:25:15.533780",
     "exception": false,
     "start_time": "2024-01-23T10:25:15.506217",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def load_tokenizers():\n",
    "    \"\"\"\n",
    "    Load the German and English tokenizers provided by spaCy.\n",
    "\n",
    "    Returns:\n",
    "        spacy_de:     German tokenizer\n",
    "        spacy_en:     English tokenizer\n",
    "    \"\"\"\n",
    "    try:\n",
    "        spacy_de = spacy.load(\"de_core_news_sm\")\n",
    "    except OSError:\n",
    "        os.system(\"python -m spacy download de_core_news_sm\")\n",
    "        spacy_de = spacy.load(\"de_core_news_sm\")\n",
    "\n",
    "    try:\n",
    "        spacy_en = spacy.load(\"en_core_web_sm\")\n",
    "    except OSError:\n",
    "        os.system(\"python -m spacy download en_core_web_sm\")\n",
    "        spacy_en = spacy.load(\"en_core_web_sm\")\n",
    "\n",
    "    print(\"Loaded English and German tokenizers.\")\n",
    "    return spacy_de, spacy_en"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "b90ea8a4",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:25:15.573112Z",
     "iopub.status.busy": "2024-01-23T10:25:15.572808Z",
     "iopub.status.idle": "2024-01-23T10:25:15.577717Z",
     "shell.execute_reply": "2024-01-23T10:25:15.576810Z"
    },
    "papermill": {
     "duration": 0.027322,
     "end_time": "2024-01-23T10:25:15.579719",
     "exception": false,
     "start_time": "2024-01-23T10:25:15.552397",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def tokenize(text: str, tokenizer):\n",
    "    \"\"\"\n",
    "    Split a string into its tokens using the provided tokenizer.\n",
    "\n",
    "    Args:\n",
    "        text:         string\n",
    "        tokenizer:    tokenizer for the language\n",
    "\n",
    "    Returns:\n",
    "        tokenized list of strings\n",
    "    \"\"\"\n",
    "    return [tok.text.lower() for tok in tokenizer.tokenizer(text)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "ebbe948c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:25:15.618116Z",
     "iopub.status.busy": "2024-01-23T10:25:15.617797Z",
     "iopub.status.idle": "2024-01-23T10:25:15.622673Z",
     "shell.execute_reply": "2024-01-23T10:25:15.621819Z"
    },
    "papermill": {
     "duration": 0.026243,
     "end_time": "2024-01-23T10:25:15.624529",
     "exception": false,
     "start_time": "2024-01-23T10:25:15.598286",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def yield_tokens(data_iter, tokenizer, index: int):\n",
    "    \"\"\"\n",
    "    Return the tokens for the appropriate language.\n",
    "\n",
    "    Args:\n",
    "        data_iter:    text here\n",
    "        tokenizer:    tokenizer for the language\n",
    "        index:        index of the language in the tuple | (de=0, en=1)\n",
    "\n",
    "    Yields:\n",
    "        sequences based on index\n",
    "    \"\"\"\n",
    "    for from_tuple in data_iter:\n",
    "        yield tokenizer(from_tuple[index])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "04340e4c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:25:15.663117Z",
     "iopub.status.busy": "2024-01-23T10:25:15.662178Z",
     "iopub.status.idle": "2024-01-23T10:25:15.670949Z",
     "shell.execute_reply": "2024-01-23T10:25:15.670102Z"
    },
    "papermill": {
     "duration": 0.029938,
     "end_time": "2024-01-23T10:25:15.672929",
     "exception": false,
     "start_time": "2024-01-23T10:25:15.642991",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def build_vocabulary(\n",
    "    spacy_de, spacy_en, train_iter, val_iter, test_iter, min_freq: int = 2\n",
    "):\n",
    "    def tokenize_de(text: str):\n",
    "        \"\"\"\n",
    "        Call the German tokenizer.\n",
    "\n",
    "        Args:\n",
    "            text:         string\n",
    "            min_freq:     minimum frequency needed to include a word in the vocabulary\n",
    "\n",
    "        Returns:\n",
    "            tokenized list of strings\n",
    "        \"\"\"\n",
    "        return tokenize(text, spacy_de)\n",
    "\n",
    "    def tokenize_en(text: str):\n",
    "        \"\"\"\n",
    "        Call the English tokenizer.\n",
    "\n",
    "        Args:\n",
    "            text:         string\n",
    "\n",
    "        Returns:\n",
    "            tokenized list of strings\n",
    "        \"\"\"\n",
    "        return tokenize(text, spacy_en)\n",
    "\n",
    "    print(\"Building German Vocabulary...\")\n",
    "\n",
    "    #     # load train, val, and test data pipelines\n",
    "    #     train, val, test = datasets.IWSLT2016(language_pair=(\"de\", \"en\"))\n",
    "    train = train_iter\n",
    "    val = val_iter\n",
    "    test = test_iter\n",
    "\n",
    "    # generate source vocabulary\n",
    "    vocab_src = build_vocab_from_iterator(\n",
    "        yield_tokens(\n",
    "            train + val + test, tokenize_de, index=0\n",
    "        ),  # tokens for each German sentence (index 0)\n",
    "        min_freq=min_freq,\n",
    "        specials=[\"<bos>\", \"<eos>\", \"<pad>\", \"<unk>\"],\n",
    "    )\n",
    "\n",
    "    print(\"Building English Vocabulary...\")\n",
    "\n",
    "    # generate target vocabulary\n",
    "    vocab_trg = build_vocab_from_iterator(\n",
    "        yield_tokens(\n",
    "            train + val + test, tokenize_en, index=1\n",
    "        ),  # tokens for each English sentence (index 1)\n",
    "        min_freq=2,  #\n",
    "        specials=[\"<bos>\", \"<eos>\", \"<pad>\", \"<unk>\"],\n",
    "    )\n",
    "\n",
    "    # set default token for out-of-vocabulary words (OOV)\n",
    "    vocab_src.set_default_index(vocab_src[\"<unk>\"])\n",
    "    vocab_trg.set_default_index(vocab_trg[\"<unk>\"])\n",
    "\n",
    "    return vocab_src, vocab_trg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "26d825a4",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:25:15.711242Z",
     "iopub.status.busy": "2024-01-23T10:25:15.710680Z",
     "iopub.status.idle": "2024-01-23T10:25:15.716978Z",
     "shell.execute_reply": "2024-01-23T10:25:15.716165Z"
    },
    "papermill": {
     "duration": 0.027498,
     "end_time": "2024-01-23T10:25:15.718828",
     "exception": false,
     "start_time": "2024-01-23T10:25:15.691330",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def load_vocab(spacy_de, spacy_en, train_iter, val_iter, test_iter, min_freq: int = 2):\n",
    "    \"\"\"\n",
    "    Args:\n",
    "        spacy_de:     German tokenizer\n",
    "        spacy_en:     English tokenizer\n",
    "        min_freq:     minimum frequency needed to include a word in the vocabulary\n",
    "\n",
    "    Returns:\n",
    "        vocab_src:    German vocabulary\n",
    "        vocab_trg:     English vocabulary\n",
    "    \"\"\"\n",
    "    if not os.path.exists(\"vocab.pt\"):\n",
    "        # build the German/English vocabulary if it does not exist\n",
    "        vocab_src, vocab_trg = build_vocabulary(\n",
    "            spacy_de, spacy_en, train_iter, val_iter, test_iter, min_freq\n",
    "        )\n",
    "        # save it to a file\n",
    "        torch.save((vocab_src, vocab_trg), \"vocab.pt\")\n",
    "    else:\n",
    "        # load the vocab if it exists\n",
    "        vocab_src, vocab_trg = torch.load(\"vocab.pt\")\n",
    "\n",
    "    print(\"Finished.\\nVocabulary sizes:\")\n",
    "    print(\"\\tSource:\", len(vocab_src))\n",
    "    print(\"\\tTarget:\", len(vocab_trg))\n",
    "    return vocab_src, vocab_trg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "e839c2f6",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:25:15.756582Z",
     "iopub.status.busy": "2024-01-23T10:25:15.756323Z",
     "iopub.status.idle": "2024-01-23T10:25:34.786846Z",
     "shell.execute_reply": "2024-01-23T10:25:34.785974Z"
    },
    "papermill": {
     "duration": 19.051673,
     "end_time": "2024-01-23T10:25:34.788809",
     "exception": false,
     "start_time": "2024-01-23T10:25:15.737136",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.10/site-packages/scipy/__init__.py:146: UserWarning: A NumPy version >=1.16.5 and <1.23.0 is required for this version of SciPy (detected version 1.24.3\n",
      "  warnings.warn(f\"A NumPy version >={np_minversion} and <{np_maxversion}\"\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "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[31m64.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
      "\u001b[?25hRequirement already satisfied: spacy<3.8.0,>=3.7.0 in /opt/conda/lib/python3.10/site-packages (from de-core-news-sm==3.7.0) (3.7.2)\n",
      "Requirement already satisfied: spacy-legacy<3.1.0,>=3.0.11 in /opt/conda/lib/python3.10/site-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 /opt/conda/lib/python3.10/site-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 /opt/conda/lib/python3.10/site-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 /opt/conda/lib/python3.10/site-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 /opt/conda/lib/python3.10/site-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.1.8 in /opt/conda/lib/python3.10/site-packages (from spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (8.2.2)\n",
      "Requirement already satisfied: wasabi<1.2.0,>=0.9.1 in /opt/conda/lib/python3.10/site-packages (from spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (1.1.2)\n",
      "Requirement already satisfied: srsly<3.0.0,>=2.4.3 in /opt/conda/lib/python3.10/site-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 /opt/conda/lib/python3.10/site-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.4.0,>=0.1.0 in /opt/conda/lib/python3.10/site-packages (from spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (0.3.4)\n",
      "Requirement already satisfied: typer<0.10.0,>=0.3.0 in /opt/conda/lib/python3.10/site-packages (from spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (0.9.0)\n",
      "Requirement already satisfied: smart-open<7.0.0,>=5.2.1 in /opt/conda/lib/python3.10/site-packages (from spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (6.3.0)\n",
      "Requirement already satisfied: tqdm<5.0.0,>=4.38.0 in /opt/conda/lib/python3.10/site-packages (from spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (4.66.1)\n",
      "Requirement already satisfied: requests<3.0.0,>=2.13.0 in /opt/conda/lib/python3.10/site-packages (from spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (2.31.0)\n",
      "Requirement already satisfied: pydantic!=1.8,!=1.8.1,<3.0.0,>=1.7.4 in /opt/conda/lib/python3.10/site-packages (from spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (1.10.12)\n",
      "Requirement already satisfied: jinja2 in /opt/conda/lib/python3.10/site-packages (from spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (3.1.2)\n",
      "Requirement already satisfied: setuptools in /opt/conda/lib/python3.10/site-packages (from spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (68.1.2)\n",
      "Requirement already satisfied: packaging>=20.0 in /opt/conda/lib/python3.10/site-packages (from spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (21.3)\n",
      "Requirement already satisfied: langcodes<4.0.0,>=3.2.0 in /opt/conda/lib/python3.10/site-packages (from spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (3.3.0)\n",
      "Requirement already satisfied: numpy>=1.19.0 in /opt/conda/lib/python3.10/site-packages (from spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (1.24.3)\n",
      "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in /opt/conda/lib/python3.10/site-packages (from packaging>=20.0->spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (3.0.9)\n",
      "Requirement already satisfied: typing-extensions>=4.2.0 in /opt/conda/lib/python3.10/site-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.5.0)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /opt/conda/lib/python3.10/site-packages (from requests<3.0.0,>=2.13.0->spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (3.2.0)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /opt/conda/lib/python3.10/site-packages (from requests<3.0.0,>=2.13.0->spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (3.4)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /opt/conda/lib/python3.10/site-packages (from requests<3.0.0,>=2.13.0->spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (1.26.15)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /opt/conda/lib/python3.10/site-packages (from requests<3.0.0,>=2.13.0->spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (2023.11.17)\n",
      "Requirement already satisfied: blis<0.8.0,>=0.7.8 in /opt/conda/lib/python3.10/site-packages (from thinc<8.3.0,>=8.1.8->spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (0.7.10)\n",
      "Requirement already satisfied: confection<1.0.0,>=0.0.1 in /opt/conda/lib/python3.10/site-packages (from thinc<8.3.0,>=8.1.8->spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (0.1.4)\n",
      "Requirement already satisfied: click<9.0.0,>=7.1.1 in /opt/conda/lib/python3.10/site-packages (from typer<0.10.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: cloudpathlib<0.17.0,>=0.7.0 in /opt/conda/lib/python3.10/site-packages (from weasel<0.4.0,>=0.1.0->spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (0.16.0)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /opt/conda/lib/python3.10/site-packages (from jinja2->spacy<3.8.0,>=3.7.0->de-core-news-sm==3.7.0) (2.1.3)\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",
      "Loaded English and German tokenizers.\n"
     ]
    }
   ],
   "source": [
    "# global variables used later in the script\n",
    "spacy_de, spacy_en = load_tokenizers()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "75aff5df",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:25:34.829168Z",
     "iopub.status.busy": "2024-01-23T10:25:34.828860Z",
     "iopub.status.idle": "2024-01-23T10:25:34.834280Z",
     "shell.execute_reply": "2024-01-23T10:25:34.833465Z"
    },
    "papermill": {
     "duration": 0.027621,
     "end_time": "2024-01-23T10:25:34.836219",
     "exception": false,
     "start_time": "2024-01-23T10:25:34.808598",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<spacy.lang.de.German at 0x7c646fe8a620>"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "spacy_de"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "77d038b1",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:25:34.876233Z",
     "iopub.status.busy": "2024-01-23T10:25:34.875950Z",
     "iopub.status.idle": "2024-01-23T10:25:39.894623Z",
     "shell.execute_reply": "2024-01-23T10:25:39.893309Z"
    },
    "papermill": {
     "duration": 5.040891,
     "end_time": "2024-01-23T10:25:39.896688",
     "exception": false,
     "start_time": "2024-01-23T10:25:34.855797",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Building German Vocabulary...\n",
      "Building English Vocabulary...\n",
      "Finished.\n",
      "Vocabulary sizes:\n",
      "\tSource: 7983\n",
      "\tTarget: 5979\n"
     ]
    }
   ],
   "source": [
    "vocab_src, vocab_trg = load_vocab(spacy_de, spacy_en, train_iter, val_iter, test_iter)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "1ff58caa",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:25:39.938041Z",
     "iopub.status.busy": "2024-01-23T10:25:39.937718Z",
     "iopub.status.idle": "2024-01-23T10:25:39.941960Z",
     "shell.execute_reply": "2024-01-23T10:25:39.941101Z"
    },
    "papermill": {
     "duration": 0.027053,
     "end_time": "2024-01-23T10:25:39.943984",
     "exception": false,
     "start_time": "2024-01-23T10:25:39.916931",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "BOS_IDX = vocab_trg[\"<bos>\"]\n",
    "EOS_IDX = vocab_trg[\"<eos>\"]\n",
    "PAD_IDX = vocab_trg[\"<pad>\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4bf98f2d",
   "metadata": {
    "papermill": {
     "duration": 0.019362,
     "end_time": "2024-01-23T10:25:39.983007",
     "exception": false,
     "start_time": "2024-01-23T10:25:39.963645",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Perform data preprocess"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "e2d66c38",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:25:40.025337Z",
     "iopub.status.busy": "2024-01-23T10:25:40.024715Z",
     "iopub.status.idle": "2024-01-23T10:25:40.031834Z",
     "shell.execute_reply": "2024-01-23T10:25:40.030964Z"
    },
    "papermill": {
     "duration": 0.031075,
     "end_time": "2024-01-23T10:25:40.033771",
     "exception": false,
     "start_time": "2024-01-23T10:25:40.002696",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def data_process(raw_data):\n",
    "    \"\"\"\n",
    "    Process raw sentences by tokenizing and converting to integers based on\n",
    "    the vocabulary.\n",
    "\n",
    "    Args:\n",
    "        raw_data:     German-English sentence pairs\n",
    "    Returns:\n",
    "        data:         tokenized data converted to index based on vocabulary\n",
    "    \"\"\"\n",
    "    data = []\n",
    "    # loop through each sentence pair\n",
    "    for raw_de, raw_en in tqdm(raw_data):\n",
    "        de_tensor_ = []\n",
    "        # tokenize the sentence and convert each word to an integers\n",
    "        for token in spacy_de.tokenizer(raw_de):\n",
    "            de_tensor_.append(vocab_src[token.text.lower()])\n",
    "\n",
    "        en_tensor_ = []\n",
    "        # tokenize the sentence and convert each word to an integers\n",
    "        for token in spacy_en.tokenizer(raw_en):\n",
    "            en_tensor_.append(vocab_trg[token.text.lower()])\n",
    "\n",
    "        de_tensor_ = torch.tensor(de_tensor_, dtype=torch.long)\n",
    "        en_tensor_ = torch.tensor(en_tensor_, dtype=torch.long)\n",
    "        # append tensor representations\n",
    "        data.append((de_tensor_, en_tensor_))\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "11c9233f",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:25:40.074794Z",
     "iopub.status.busy": "2024-01-23T10:25:40.074497Z",
     "iopub.status.idle": "2024-01-23T10:25:44.107407Z",
     "shell.execute_reply": "2024-01-23T10:25:44.106446Z"
    },
    "papermill": {
     "duration": 4.055653,
     "end_time": "2024-01-23T10:25:44.109474",
     "exception": false,
     "start_time": "2024-01-23T10:25:40.053821",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 27550/27550 [00:03<00:00, 7487.18it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train data shape: 27550\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1450/1450 [00:00<00:00, 7258.85it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Val data shape: 1450\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1000/1000 [00:00<00:00, 7254.49it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test data shape: 1000\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# processed data\n",
    "train_data = data_process(train_iter)\n",
    "print(f\"Train data shape: {len(train_data)}\")\n",
    "val_data = data_process(val_iter)\n",
    "print(f\"Val data shape: {len(val_data)}\")\n",
    "test_data = data_process(test_iter)\n",
    "print(f\"Test data shape: {len(test_data)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2476953",
   "metadata": {
    "papermill": {
     "duration": 0.02333,
     "end_time": "2024-01-23T10:25:44.156551",
     "exception": false,
     "start_time": "2024-01-23T10:25:44.133221",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Create collate function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "10eb78dd",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:25:44.204664Z",
     "iopub.status.busy": "2024-01-23T10:25:44.204233Z",
     "iopub.status.idle": "2024-01-23T10:25:44.212088Z",
     "shell.execute_reply": "2024-01-23T10:25:44.211246Z"
    },
    "papermill": {
     "duration": 0.034132,
     "end_time": "2024-01-23T10:25:44.214017",
     "exception": false,
     "start_time": "2024-01-23T10:25:44.179885",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def generate_batch(data_batch):\n",
    "    \"\"\"\n",
    "    Process indexed-sequences by adding <bos>, <eos>, and <pad> tokens.\n",
    "\n",
    "    Args:\n",
    "        data_batch:     German-English indexed-sentence pairs\n",
    "\n",
    "    Returns:\n",
    "        two batches:    one for German and one for English\n",
    "    \"\"\"\n",
    "    de_batch, en_batch = [], []\n",
    "\n",
    "    # for each sentence\n",
    "    for de_item, en_item in data_batch:\n",
    "        # add <bos> and <eos> indices before and after the sentence\n",
    "        de_temp = torch.cat(\n",
    "            [torch.tensor([BOS_IDX]), de_item, torch.tensor([EOS_IDX])], dim=0\n",
    "        ).to(device)\n",
    "        en_temp = torch.cat(\n",
    "            [torch.tensor([BOS_IDX]), en_item, torch.tensor([EOS_IDX])], dim=0\n",
    "        ).to(device)\n",
    "\n",
    "        # add padding\n",
    "        de_batch.append(\n",
    "            pad(\n",
    "                de_temp,\n",
    "                (\n",
    "                    0,  # dimension to pad\n",
    "                    MAX_PADDING - len(de_temp),  # amount of padding to add\n",
    "                ),\n",
    "                value=PAD_IDX,\n",
    "            )\n",
    "        )\n",
    "\n",
    "        # add padding\n",
    "        en_batch.append(\n",
    "            pad(\n",
    "                en_temp,\n",
    "                (\n",
    "                    0,  # dimension to pad\n",
    "                    MAX_PADDING - len(en_temp),  # amount of padding to add\n",
    "                ),\n",
    "                value=PAD_IDX,\n",
    "            )\n",
    "        )\n",
    "\n",
    "    return torch.stack(de_batch), torch.stack(en_batch)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67520041",
   "metadata": {
    "papermill": {
     "duration": 0.023574,
     "end_time": "2024-01-23T10:25:44.261934",
     "exception": false,
     "start_time": "2024-01-23T10:25:44.238360",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Creating DataLoaders for training and testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "385e649e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:25:44.312590Z",
     "iopub.status.busy": "2024-01-23T10:25:44.311879Z",
     "iopub.status.idle": "2024-01-23T10:25:44.322112Z",
     "shell.execute_reply": "2024-01-23T10:25:44.321200Z"
    },
    "papermill": {
     "duration": 0.037067,
     "end_time": "2024-01-23T10:25:44.324044",
     "exception": false,
     "start_time": "2024-01-23T10:25:44.286977",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "MAX_PADDING = 20\n",
    "BATCH_SIZE = 128\n",
    "\n",
    "train_iter = DataLoader(\n",
    "    to_map_style_dataset(train_data),\n",
    "    batch_size=BATCH_SIZE,\n",
    "    shuffle=True,\n",
    "    drop_last=True,\n",
    "    collate_fn=generate_batch,\n",
    ")\n",
    "\n",
    "valid_iter = DataLoader(\n",
    "    to_map_style_dataset(val_data),\n",
    "    batch_size=BATCH_SIZE,\n",
    "    shuffle=True,\n",
    "    drop_last=True,\n",
    "    collate_fn=generate_batch,\n",
    ")\n",
    "\n",
    "test_iter = DataLoader(\n",
    "    to_map_style_dataset(test_data),\n",
    "    batch_size=BATCH_SIZE,\n",
    "    shuffle=True,\n",
    "    drop_last=True,\n",
    "    collate_fn=generate_batch,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1366cde4",
   "metadata": {
    "papermill": {
     "duration": 0.023276,
     "end_time": "2024-01-23T10:25:44.370785",
     "exception": false,
     "start_time": "2024-01-23T10:25:44.347509",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Creating Attention Network Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "2d07811c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:25:44.419283Z",
     "iopub.status.busy": "2024-01-23T10:25:44.418634Z",
     "iopub.status.idle": "2024-01-23T10:25:44.422739Z",
     "shell.execute_reply": "2024-01-23T10:25:44.421833Z"
    },
    "papermill": {
     "duration": 0.03036,
     "end_time": "2024-01-23T10:25:44.424710",
     "exception": false,
     "start_time": "2024-01-23T10:25:44.394350",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import math"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d14270b5",
   "metadata": {
    "papermill": {
     "duration": 0.023248,
     "end_time": "2024-01-23T10:25:44.471163",
     "exception": false,
     "start_time": "2024-01-23T10:25:44.447915",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Create submodules for network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "bb044b6b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:25:44.519487Z",
     "iopub.status.busy": "2024-01-23T10:25:44.518600Z",
     "iopub.status.idle": "2024-01-23T10:25:44.525171Z",
     "shell.execute_reply": "2024-01-23T10:25:44.524484Z"
    },
    "papermill": {
     "duration": 0.032718,
     "end_time": "2024-01-23T10:25:44.527158",
     "exception": false,
     "start_time": "2024-01-23T10:25:44.494440",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class Embedding(nn.Module):\n",
    "    \"\"\"\n",
    "    Embedding lookup table which is used by the positional embedding block.\n",
    "    Embedding lookup table is shared across input and output\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, vocab_size, dmodel):\n",
    "        \"\"\"\n",
    "        Embedding lookup needs a vocab size and model dimension size matrix for\n",
    "        creating lookups\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.embedding_lookup = nn.Embedding(vocab_size, dmodel)\n",
    "        self.vocab_size = vocab_size\n",
    "        self.dmodel = dmodel\n",
    "\n",
    "    def forward(self, token_ids):\n",
    "        \"\"\"\n",
    "        For a given token lookup the embedding vector\n",
    "\n",
    "        As per the paper, we also multiply the embedding vector with sqrt of dmodel\n",
    "        \"\"\"\n",
    "        assert (\n",
    "            token_ids.ndim == 2\n",
    "        ), f\"Expected: (batch size, max token sequence length), got {token_ids.shape}\"\n",
    "\n",
    "        embedding_vector = self.embedding_lookup(token_ids)\n",
    "\n",
    "        return embedding_vector * math.sqrt(self.dmodel)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "faac69f5",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:25:44.575408Z",
     "iopub.status.busy": "2024-01-23T10:25:44.574920Z",
     "iopub.status.idle": "2024-01-23T10:25:44.584292Z",
     "shell.execute_reply": "2024-01-23T10:25:44.583607Z"
    },
    "papermill": {
     "duration": 0.035635,
     "end_time": "2024-01-23T10:25:44.586213",
     "exception": false,
     "start_time": "2024-01-23T10:25:44.550578",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class PositionalEncoding(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        dmodel,\n",
    "        max_seq_length=5000,\n",
    "        pdropout=0.1,\n",
    "    ):\n",
    "        \"\"\"\n",
    "        dmodel(int): model dimensions\n",
    "        max_seq_length(int): Maximum input sequence length\n",
    "        pdropout(float): Dropout probability\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        self.dropout = nn.Dropout(p=pdropout)\n",
    "\n",
    "        # Calculate frequencies\n",
    "        position_ids = torch.arange(0, max_seq_length).unsqueeze(1)\n",
    "        # -ve sign is added because the exponents are inverted when you multiply position and frequencies\n",
    "        frequencies = torch.pow(\n",
    "            10000, -torch.arange(0, dmodel, 2, dtype=torch.float) / dmodel\n",
    "        )\n",
    "\n",
    "        # Create positional encoding table\n",
    "        positional_encoding_table = torch.zeros(max_seq_length, dmodel)\n",
    "        # Fill the table with even entries with sin and odd entries with cosine\n",
    "        positional_encoding_table[:, 0::2] = torch.sin(position_ids * frequencies)\n",
    "        positional_encoding_table[:, 1::2] = torch.cos(position_ids * frequencies)\n",
    "\n",
    "        # Registering the position enconding in state_dict but the its not included\n",
    "        # in named parameter as it is not trainable\n",
    "        self.register_buffer(\"positional_encoding_table\", positional_encoding_table)\n",
    "\n",
    "    def forward(self, embeddings_batch):\n",
    "        \"\"\"\n",
    "        embeddings_batch shape = (batch size, seq_length, dmodel)\n",
    "        positional_encoding_table shape = (max_seq_length, dmodel)\n",
    "        \"\"\"\n",
    "        assert (\n",
    "            embeddings_batch.ndim == 3\n",
    "        ), f\"Embeddings batch should have dimension of 3 but got {embeddings_batch.ndim}\"\n",
    "        assert (\n",
    "            embeddings_batch.size()[-1] == self.positional_encoding_table.size()[-1]\n",
    "        ), f\"Embedding batch shape and positional_encoding_table shape should match, expected Embedding batch shape : {embeddings_batch.shape[-1]} while positional_encoding_table shape : {positional_encoding_table[-1]}\"\n",
    "\n",
    "        # Get encodings for the given input sequence length\n",
    "        pos_encodings = self.positional_encoding_table[\n",
    "            : embeddings_batch.shape[1]\n",
    "        ]  # Choose only seq_length out of max_seq_length\n",
    "\n",
    "        # Final output\n",
    "        out = embeddings_batch + pos_encodings\n",
    "        out = self.dropout(out)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "7b09ec14",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:25:44.634884Z",
     "iopub.status.busy": "2024-01-23T10:25:44.634543Z",
     "iopub.status.idle": "2024-01-23T10:25:44.641220Z",
     "shell.execute_reply": "2024-01-23T10:25:44.640328Z"
    },
    "papermill": {
     "duration": 0.033183,
     "end_time": "2024-01-23T10:25:44.643118",
     "exception": false,
     "start_time": "2024-01-23T10:25:44.609935",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class PositionwiseFeedForward(nn.Module):\n",
    "    def __init__(self, dmodel, dff, pdropout=0.1):\n",
    "        super().__init__()\n",
    "\n",
    "        self.dropout = nn.Dropout(p=pdropout)\n",
    "\n",
    "        self.W1 = nn.Linear(dmodel, dff)  # Intermediate layer\n",
    "        self.W2 = nn.Linear(dff, dmodel)  # Output layer\n",
    "\n",
    "        self.relu = nn.ReLU()\n",
    "\n",
    "    def forward(self, x):\n",
    "        \"\"\"\n",
    "        Perform Feedforward calculation\n",
    "\n",
    "        x shape = (B - batch size, S/T - max token sequence length, D- model dimension).\n",
    "        \"\"\"\n",
    "        out = self.W2(self.relu(self.dropout(self.W1(x))))\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "26ec852d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:25:44.691625Z",
     "iopub.status.busy": "2024-01-23T10:25:44.691273Z",
     "iopub.status.idle": "2024-01-23T10:25:44.705319Z",
     "shell.execute_reply": "2024-01-23T10:25:44.704463Z"
    },
    "papermill": {
     "duration": 0.040578,
     "end_time": "2024-01-23T10:25:44.707225",
     "exception": false,
     "start_time": "2024-01-23T10:25:44.666647",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class MultiHeadAttention(nn.Module):\n",
    "    \"\"\"\n",
    "    We can refer to the following blog to understand in depth about the transformer and MHA\n",
    "    https://medium.com/@hunter-j-phillips/multi-head-attention-7924371d477a\n",
    "\n",
    "    Here we are clubbing all the linear layers together and duplicating the inputs and\n",
    "    then performing matrix multiplications\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, dk, dv, h, pdropout=0.1):\n",
    "        \"\"\"\n",
    "        Input Args:\n",
    "\n",
    "        dk(int): Key dimensions used for generating Key weight matrix\n",
    "        dv(int): Val dimensions used for generating val weight matrix\n",
    "        h(int) : Number of heads in MHA\n",
    "        \"\"\"\n",
    "        super().__init__()\n",
    "        assert dk == dv\n",
    "        self.dk = dk\n",
    "        self.dv = dv\n",
    "        self.h = h\n",
    "        self.dmodel = self.dk * self.h  # model dimension\n",
    "\n",
    "        # Add the params in modulelist as the params in the conv list needs to be tracked\n",
    "        # wq, wk, wv -> multiple linear weights for the number of heads\n",
    "        self.WQ = nn.Linear(self.dmodel, self.dmodel)  # shape -> (dmodel, dmodel)\n",
    "        self.WK = nn.Linear(self.dmodel, self.dmodel)  # shape -> (dmodel, dmodel)\n",
    "        self.WV = nn.Linear(self.dmodel, self.dmodel)  # shape -> (dmodel, dmodel)\n",
    "        # Output Weights\n",
    "        self.WO = nn.Linear(self.h * self.dv, self.dmodel)  # shape -> (dmodel, dmodel)\n",
    "        self.softmax = nn.Softmax(dim=-1)\n",
    "        self.dropout = nn.Dropout(p=pdropout)\n",
    "\n",
    "    def forward(self, query, key, val, mask=None):\n",
    "        \"\"\"\n",
    "        Forward pass for MHA\n",
    "\n",
    "        X has a size of (batch_size, seq_length, d_model)\n",
    "        Wq, Wk, and Wv have a size of (d_model, d_model)\n",
    "\n",
    "        Perform Scaled Dot Product Attention on multi head attention.\n",
    "\n",
    "        Notation: B - batch size, S/T - max src/trg token-sequence length\n",
    "        query shape = (B, S, dmodel)\n",
    "        key shape = (B, S, dmodel)\n",
    "        val shape = (B, S, dmodel)\n",
    "        \"\"\"\n",
    "        # Weight the queries\n",
    "        Q = self.WQ(query)  # shape -> (B, S, dmodel)\n",
    "        K = self.WK(key)  # shape -> (B, S, dmodel)\n",
    "        V = self.WV(val)  # shape -> (B, S, dmodel)\n",
    "\n",
    "        # Separate last dimension to number of head and dk\n",
    "        batch_size = Q.size(0)\n",
    "        Q = Q.view(batch_size, -1, self.h, self.dk)  # shape -> (B, S, h, dk)\n",
    "        K = K.view(batch_size, -1, self.h, self.dk)  # shape -> (B, S, h, dk)\n",
    "        V = V.view(batch_size, -1, self.h, self.dk)  # shape -> (B, S, h, dk)\n",
    "\n",
    "        # each sequence is split across n_heads, with each head receiving seq_length tokens\n",
    "        # with d_key elements in each token instead of d_model.\n",
    "        Q = Q.permute(0, 2, 1, 3)  # shape -> (B, h, S, dk)\n",
    "        K = K.permute(0, 2, 1, 3)  # shape -> (B, h, S, dk)\n",
    "        V = V.permute(0, 2, 1, 3)  # shape -> (B, h, S, dk)\n",
    "\n",
    "        # dot product of Q and K\n",
    "        scaled_dot_product = torch.matmul(Q, K.permute(0, 1, 3, 2)) / math.sqrt(self.dk)\n",
    "\n",
    "        # fill those positions of product as (-1e10) where mask positions are 0\n",
    "        if mask is not None:\n",
    "            scaled_dot_product = scaled_dot_product.masked_fill(mask == 0, -1e10)\n",
    "\n",
    "        attn_probs = self.softmax(scaled_dot_product)\n",
    "\n",
    "        # Create head\n",
    "        head = torch.matmul(\n",
    "            self.dropout(attn_probs), V\n",
    "        )  # shape -> (B, h, S, S) * (B, h, S, dk) = (B, h, S, dk)\n",
    "        # Prepare the head to pass it through output linear layer\n",
    "        head = head.permute(0, 2, 1, 3).contiguous()  # shape -> (B, S, h, dk)\n",
    "        # Concatenate the head together\n",
    "        head = head.view(\n",
    "            batch_size, -1, self.h * self.dk\n",
    "        )  # shape -> (B, S, (h*dk = dmodel))\n",
    "        # Pass through output layer\n",
    "        token_representation = self.WO(head)\n",
    "        return token_representation, attn_probs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28ffab98",
   "metadata": {
    "papermill": {
     "duration": 0.023219,
     "end_time": "2024-01-23T10:25:44.753941",
     "exception": false,
     "start_time": "2024-01-23T10:25:44.730722",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Create Encoder for the Network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "c7cd71a4",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:25:44.802690Z",
     "iopub.status.busy": "2024-01-23T10:25:44.801756Z",
     "iopub.status.idle": "2024-01-23T10:25:44.810801Z",
     "shell.execute_reply": "2024-01-23T10:25:44.809902Z"
    },
    "papermill": {
     "duration": 0.035548,
     "end_time": "2024-01-23T10:25:44.812777",
     "exception": false,
     "start_time": "2024-01-23T10:25:44.777229",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class EncoderLayer(nn.Module):\n",
    "    \"\"\"\n",
    "    This building block in the encoder layer consists of the following\n",
    "    1. MultiHead Attention\n",
    "    2. Sublayer Logic\n",
    "    3. Positional FeedForward Network\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, dk, dv, h, dim_multiplier=4, pdropout=0.1):\n",
    "        super().__init__()\n",
    "        self.attention = MultiHeadAttention(dk, dv, h, pdropout)\n",
    "        # Reference page 5 chapter 3.2.2 Multi-head attention\n",
    "        dmodel = dk * h\n",
    "        # Reference page 5 chapter 3.3 positionwise FeedForward\n",
    "        dff = dmodel * dim_multiplier\n",
    "        self.attn_norm = nn.LayerNorm(dmodel)\n",
    "        self.ff = PositionwiseFeedForward(dmodel, dff, pdropout=pdropout)\n",
    "        self.ff_norm = nn.LayerNorm(dmodel)\n",
    "\n",
    "        self.dropout = nn.Dropout(p=pdropout)\n",
    "\n",
    "    def forward(self, src_inputs, src_mask=None):\n",
    "        \"\"\"\n",
    "        Forward pass as per page 3 chapter 3.1\n",
    "        \"\"\"\n",
    "        mha_out, attention_wts = self.attention(\n",
    "            query=src_inputs, key=src_inputs, val=src_inputs, mask=src_mask\n",
    "        )\n",
    "\n",
    "        # Residual connection between input and sublayer output, details: Page 7, Chapter 5.4 \"Regularization\",\n",
    "        # Actual paper design is the following\n",
    "        intermediate_out = self.attn_norm(src_inputs + self.dropout(mha_out))\n",
    "\n",
    "        pff_out = self.ff(intermediate_out)\n",
    "\n",
    "        # Perform Add Norm again\n",
    "        out = self.ff_norm(intermediate_out + self.dropout(pff_out))\n",
    "        return out, attention_wts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "3b081599",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:25:44.860964Z",
     "iopub.status.busy": "2024-01-23T10:25:44.860657Z",
     "iopub.status.idle": "2024-01-23T10:25:44.867285Z",
     "shell.execute_reply": "2024-01-23T10:25:44.866458Z"
    },
    "papermill": {
     "duration": 0.032991,
     "end_time": "2024-01-23T10:25:44.869100",
     "exception": false,
     "start_time": "2024-01-23T10:25:44.836109",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class Encoder(nn.Module):\n",
    "    def __init__(self, dk, dv, h, num_encoders, dim_multiplier=4, pdropout=0.1):\n",
    "        super().__init__()\n",
    "        self.encoder_layers = nn.ModuleList(\n",
    "            [\n",
    "                EncoderLayer(dk, dv, h, dim_multiplier, pdropout)\n",
    "                for _ in range(num_encoders)\n",
    "            ]\n",
    "        )\n",
    "\n",
    "    def forward(self, src_inputs, src_mask=None):\n",
    "        \"\"\"\n",
    "        Input from the Embedding layer\n",
    "        src_inputs = (B - batch size, S/T - max token sequence length, D- model dimension)\n",
    "        \"\"\"\n",
    "        src_representation = src_inputs\n",
    "\n",
    "        # Forward pass through encoder stack\n",
    "        for enc in self.encoder_layers:\n",
    "            src_representation, attn_probs = enc(src_representation, src_mask)\n",
    "\n",
    "        self.attn_probs = attn_probs\n",
    "        return src_representation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af9f7e86",
   "metadata": {
    "papermill": {
     "duration": 0.023288,
     "end_time": "2024-01-23T10:25:44.915806",
     "exception": false,
     "start_time": "2024-01-23T10:25:44.892518",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Create Decoder Layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "caa87d49",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:25:44.964009Z",
     "iopub.status.busy": "2024-01-23T10:25:44.963711Z",
     "iopub.status.idle": "2024-01-23T10:25:44.974080Z",
     "shell.execute_reply": "2024-01-23T10:25:44.973241Z"
    },
    "papermill": {
     "duration": 0.036928,
     "end_time": "2024-01-23T10:25:44.976014",
     "exception": false,
     "start_time": "2024-01-23T10:25:44.939086",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class DecoderLayer(nn.Module):\n",
    "    def __init__(self, dk, dv, h, dim_multiplier=4, pdropout=0.1):\n",
    "        super().__init__()\n",
    "\n",
    "        # Reference page 5 chapter 3.2.2 Multi-head attention\n",
    "        dmodel = dk * h\n",
    "        # Reference page 5 chapter 3.3 positionwise FeedForward\n",
    "        dff = dmodel * dim_multiplier\n",
    "\n",
    "        # Masked Multi Head Attention\n",
    "        self.masked_attention = MultiHeadAttention(dk, dv, h, pdropout)\n",
    "        self.masked_attn_norm = nn.LayerNorm(dmodel)\n",
    "\n",
    "        # Multi head attention\n",
    "        self.attention = MultiHeadAttention(dk, dv, h, pdropout)\n",
    "        self.attn_norm = nn.LayerNorm(dmodel)\n",
    "\n",
    "        # Add position FeedForward Network\n",
    "        self.ff = PositionwiseFeedForward(dmodel, dff, pdropout=pdropout)\n",
    "        self.ff_norm = nn.LayerNorm(dmodel)\n",
    "\n",
    "        self.dropout = nn.Dropout(p=pdropout)\n",
    "\n",
    "    def forward(self, trg: Tensor, src: Tensor, trg_mask: Tensor, src_mask: Tensor):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            trg:          embedded sequences                (batch_size, trg_seq_length, d_model)\n",
    "            src:          embedded sequences                (batch_size, src_seq_length, d_model)\n",
    "            trg_mask:     mask for the sequences            (batch_size, 1, trg_seq_length, trg_seq_length)\n",
    "            src_mask:     mask for the sequences            (batch_size, 1, 1, src_seq_length)\n",
    "\n",
    "        Returns:\n",
    "            trg:          sequences after self-attention    (batch_size, trg_seq_length, d_model)\n",
    "            attn_probs:   self-attention softmax scores     (batch_size, n_heads, trg_seq_length, src_seq_length)\n",
    "        \"\"\"\n",
    "        _trg, attn_probs = self.masked_attention(\n",
    "            query=trg, key=trg, val=trg, mask=trg_mask\n",
    "        )\n",
    "\n",
    "        # Residual connection between input and sublayer output, details: Page 7, Chapter 5.4 \"Regularization\",\n",
    "        # Actual paper design is the following\n",
    "        trg = self.masked_attn_norm(trg + self.dropout(_trg))\n",
    "\n",
    "        # Inputs to the decoder attention is given as follows\n",
    "        # query = previous decoder layer\n",
    "        # key and val = output of encoder\n",
    "        # mask = src_mask\n",
    "        # Reference : page 5 chapter 3.2.3 point 1\n",
    "        _trg, attn_probs = self.attention(query=trg, key=src, val=src, mask=src_mask)\n",
    "        trg = self.attn_norm(trg + self.dropout(_trg))\n",
    "\n",
    "        # position-wise feed-forward network\n",
    "        _trg = self.ff(trg)\n",
    "        # Perform Add Norm again\n",
    "        trg = self.ff_norm(trg + self.dropout(_trg))\n",
    "        return trg, attn_probs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "2128839e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:25:45.023868Z",
     "iopub.status.busy": "2024-01-23T10:25:45.023188Z",
     "iopub.status.idle": "2024-01-23T10:25:45.030752Z",
     "shell.execute_reply": "2024-01-23T10:25:45.029711Z"
    },
    "papermill": {
     "duration": 0.033899,
     "end_time": "2024-01-23T10:25:45.033031",
     "exception": false,
     "start_time": "2024-01-23T10:25:44.999132",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class Decoder(nn.Module):\n",
    "    def __init__(self, dk, dv, h, num_decoders, dim_multiplier=4, pdropout=0.1):\n",
    "        super().__init__()\n",
    "        self.decoder_layers = nn.ModuleList(\n",
    "            [\n",
    "                DecoderLayer(dk, dv, h, dim_multiplier, pdropout)\n",
    "                for _ in range(num_decoders)\n",
    "            ]\n",
    "        )\n",
    "\n",
    "    def forward(self, target_inputs, src_inputs, target_mask, src_mask):\n",
    "        \"\"\"\n",
    "        Input from the Embedding layer\n",
    "        target_inputs = embedded sequences    (batch_size, trg_seq_length, d_model)\n",
    "        src_inputs = embedded sequences       (batch_size, src_seq_length, d_model)\n",
    "        target_mask = mask for the sequences  (batch_size, 1, trg_seq_length, trg_seq_length)\n",
    "        src_mask = mask for the sequences     (batch_size, 1, 1, src_seq_length)\n",
    "        \"\"\"\n",
    "        target_representation = target_inputs\n",
    "\n",
    "        # Forward pass through decoder stack\n",
    "        for layer in self.decoder_layers:\n",
    "            target_representation, attn_probs = layer(\n",
    "                target_representation, src_inputs, target_mask, src_mask\n",
    "            )\n",
    "        self.attn_probs = attn_probs\n",
    "        return target_representation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a071433",
   "metadata": {
    "papermill": {
     "duration": 0.023094,
     "end_time": "2024-01-23T10:25:45.079566",
     "exception": false,
     "start_time": "2024-01-23T10:25:45.056472",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Adding all up to construct the complete model for language translation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "e61b5ff6",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:25:45.127748Z",
     "iopub.status.busy": "2024-01-23T10:25:45.127442Z",
     "iopub.status.idle": "2024-01-23T10:25:45.144040Z",
     "shell.execute_reply": "2024-01-23T10:25:45.143065Z"
    },
    "papermill": {
     "duration": 0.043124,
     "end_time": "2024-01-23T10:25:45.145982",
     "exception": false,
     "start_time": "2024-01-23T10:25:45.102858",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "class Transformer(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        dk,\n",
    "        dv,\n",
    "        h,\n",
    "        src_vocab_size,\n",
    "        target_vocab_size,\n",
    "        num_encoders,\n",
    "        num_decoders,\n",
    "        src_pad_idx,\n",
    "        target_pad_idx,\n",
    "        dim_multiplier=4,\n",
    "        pdropout=0.1,\n",
    "        device=\"cpu\",\n",
    "    ):\n",
    "        super().__init__()\n",
    "\n",
    "        # Reference page 5 chapter 3.2.2 Multi-head attention\n",
    "        dmodel = dk * h\n",
    "        # Modules required to build Encoder\n",
    "        self.src_embeddings = Embedding(src_vocab_size, dmodel)\n",
    "        self.src_positional_encoding = PositionalEncoding(\n",
    "            dmodel, max_seq_length=src_vocab_size, pdropout=pdropout\n",
    "        )\n",
    "        self.encoder = Encoder(\n",
    "            dk, dv, h, num_encoders, dim_multiplier=dim_multiplier, pdropout=pdropout\n",
    "        )\n",
    "\n",
    "        # Modules required to build Decoder\n",
    "        self.target_embeddings = Embedding(target_vocab_size, dmodel)\n",
    "        self.target_positional_encoding = PositionalEncoding(\n",
    "            dmodel, max_seq_length=target_vocab_size, pdropout=pdropout\n",
    "        )\n",
    "        self.decoder = Decoder(dk, dv, h, num_decoders, dim_multiplier=4, pdropout=0.1)\n",
    "\n",
    "        # Final output\n",
    "        self.linear = nn.Linear(dmodel, target_vocab_size)\n",
    "        #         self.softmax = nn.Softmax(dim=-1)\n",
    "        self.device = device\n",
    "        self.src_pad_idx = src_pad_idx\n",
    "        self.target_pad_idx = target_pad_idx\n",
    "        self.init_params()\n",
    "\n",
    "    # This part wasn't mentioned in the paper, but it's super important!\n",
    "    def init_params(self):\n",
    "        \"\"\"\n",
    "        xavier has tremendous impact! I didn't expect\n",
    "        that the model's perf, with normalization layers,\n",
    "        is so dependent on the choice of weight initialization.\n",
    "        \"\"\"\n",
    "        for name, p in self.named_parameters():\n",
    "            if p.dim() > 1:\n",
    "                nn.init.xavier_uniform_(p)\n",
    "\n",
    "    def make_src_mask(self, src):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            src: raw sequences with padding        (batch_size, seq_length)\n",
    "            src_pad_idx(int): index where the token need not be attended\n",
    "\n",
    "        Returns:\n",
    "            src_mask: mask for each sequence            (batch_size, 1, 1, seq_length)\n",
    "        \"\"\"\n",
    "        batch_size = src.shape[0]\n",
    "        # assign 1 to tokens that need attended to and 0 to padding tokens,\n",
    "        # then add 2 dimensions\n",
    "        src_mask = (src != self.src_pad_idx).view(batch_size, 1, 1, -1)\n",
    "        return src_mask\n",
    "\n",
    "    def make_target_mask(self, target):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            target:  raw sequences with padding        (batch_size, seq_length)\n",
    "            target_pad_idx(int): index where the token need not be attended\n",
    "\n",
    "        Returns:\n",
    "            target_mask: mask for each sequence   (batch_size, 1, seq_length, seq_length)\n",
    "        \"\"\"\n",
    "\n",
    "        seq_length = target.shape[1]\n",
    "        batch_size = target.shape[0]\n",
    "\n",
    "        # assign True to tokens that need attended to and\n",
    "        # False to padding tokens, then add 2 dimensions\n",
    "        target_mask = (target != self.target_pad_idx).view(\n",
    "            batch_size, 1, 1, -1\n",
    "        )  # (batch_size, 1, 1, seq_length)\n",
    "\n",
    "        # generate subsequent mask\n",
    "        trg_sub_mask = torch.tril(\n",
    "            torch.ones((seq_length, seq_length), device=self.device)\n",
    "        ).bool()  # (batch_size, 1, seq_length, seq_length)\n",
    "\n",
    "        # bitwise \"and\" operator | 0 & 0 = 0, 1 & 1 = 1, 1 & 0 = 0\n",
    "        target_mask = target_mask & trg_sub_mask\n",
    "\n",
    "        return target_mask\n",
    "\n",
    "    def forward(self, src_token_ids_batch, target_token_ids_batch):\n",
    "        # create source and target masks\n",
    "        src_mask = self.make_src_mask(\n",
    "            src_token_ids_batch\n",
    "        )  # (batch_size, 1, 1, src_seq_length)\n",
    "        target_mask = self.make_target_mask(\n",
    "            target_token_ids_batch\n",
    "        )  # (batch_size, 1, trg_seq_length, trg_seq_length)\n",
    "\n",
    "        # Create embeddings\n",
    "        src_representations = self.src_embeddings(src_token_ids_batch)\n",
    "        src_representations = self.src_positional_encoding(src_representations)\n",
    "\n",
    "        target_representations = self.target_embeddings(target_token_ids_batch)\n",
    "        target_representations = self.target_positional_encoding(target_representations)\n",
    "\n",
    "        # Encode\n",
    "        encoded_src = self.encoder(src_representations, src_mask)\n",
    "\n",
    "        # Decode\n",
    "        decoded_output = self.decoder(\n",
    "            target_representations, encoded_src, target_mask, src_mask\n",
    "        )\n",
    "\n",
    "        # Post processing\n",
    "        out = self.linear(decoded_output)\n",
    "        # Don't use softmax as we are not comparing against softmaxed output while\n",
    "        # computing loss. We are comparing against linear outputs\n",
    "        #         # Output\n",
    "        #         out = self.softmax(out)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9cbd709",
   "metadata": {
    "papermill": {
     "duration": 0.023086,
     "end_time": "2024-01-23T10:25:45.192389",
     "exception": false,
     "start_time": "2024-01-23T10:25:45.169303",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Create Model Instance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "c3918202",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:25:45.241145Z",
     "iopub.status.busy": "2024-01-23T10:25:45.240652Z",
     "iopub.status.idle": "2024-01-23T10:25:45.259212Z",
     "shell.execute_reply": "2024-01-23T10:25:45.258372Z"
    },
    "papermill": {
     "duration": 0.045157,
     "end_time": "2024-01-23T10:25:45.261155",
     "exception": false,
     "start_time": "2024-01-23T10:25:45.215998",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(2, 2)"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "src_pad_idx = vocab_src.get_stoi()[\"<pad>\"]\n",
    "target_pad_idx = vocab_trg.get_stoi()[\"<pad>\"]\n",
    "src_pad_idx, target_pad_idx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "2205ad16",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:25:45.309311Z",
     "iopub.status.busy": "2024-01-23T10:25:45.309032Z",
     "iopub.status.idle": "2024-01-23T10:25:45.314289Z",
     "shell.execute_reply": "2024-01-23T10:25:45.313487Z"
    },
    "papermill": {
     "duration": 0.03155,
     "end_time": "2024-01-23T10:25:45.316342",
     "exception": false,
     "start_time": "2024-01-23T10:25:45.284792",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(7983, 5979)"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(vocab_src), len(vocab_trg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "60e08f3e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:25:45.364709Z",
     "iopub.status.busy": "2024-01-23T10:25:45.364188Z",
     "iopub.status.idle": "2024-01-23T10:25:45.381233Z",
     "shell.execute_reply": "2024-01-23T10:25:45.380413Z"
    },
    "papermill": {
     "duration": 0.043289,
     "end_time": "2024-01-23T10:25:45.383089",
     "exception": false,
     "start_time": "2024-01-23T10:25:45.339800",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'dk': 32,\n",
       " 'dv': 32,\n",
       " 'h': 8,\n",
       " 'src_vocab_size': 8500,\n",
       " 'target_vocab_size': 6500,\n",
       " 'src_pad_idx': 2,\n",
       " 'target_pad_idx': 2,\n",
       " 'num_encoders': 3,\n",
       " 'num_decoders': 3,\n",
       " 'dim_multiplier': 4,\n",
       " 'pdropout': 0.1,\n",
       " 'lr': 0.0003,\n",
       " 'N_EPOCHS': 50,\n",
       " 'CLIP': 1,\n",
       " 'patience': 5}"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "config = {\n",
    "    \"dk\": 32,\n",
    "    \"dv\": 32,\n",
    "    \"h\": 8,\n",
    "    \"src_vocab_size\": 8500,\n",
    "    \"target_vocab_size\": 6500,\n",
    "    \"src_pad_idx\": vocab_src.get_stoi()[\"<pad>\"],\n",
    "    \"target_pad_idx\": vocab_trg.get_stoi()[\"<pad>\"],\n",
    "    \"num_encoders\": 3,\n",
    "    \"num_decoders\": 3,\n",
    "    \"dim_multiplier\": 4,\n",
    "    \"pdropout\": 0.1,\n",
    "    \"lr\": 0.0003,\n",
    "    \"N_EPOCHS\": 50,\n",
    "    \"CLIP\": 1,\n",
    "    \"patience\": 5,\n",
    "}\n",
    "config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "0f2a2052",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:25:45.432402Z",
     "iopub.status.busy": "2024-01-23T10:25:45.431733Z",
     "iopub.status.idle": "2024-01-23T10:25:46.062912Z",
     "shell.execute_reply": "2024-01-23T10:25:46.061893Z"
    },
    "papermill": {
     "duration": 0.658717,
     "end_time": "2024-01-23T10:25:46.065539",
     "exception": false,
     "start_time": "2024-01-23T10:25:45.406822",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Transformer(\n",
      "  (src_embeddings): Embedding(\n",
      "    (embedding_lookup): Embedding(8500, 256)\n",
      "  )\n",
      "  (src_positional_encoding): PositionalEncoding(\n",
      "    (dropout): Dropout(p=0.1, inplace=False)\n",
      "  )\n",
      "  (encoder): Encoder(\n",
      "    (encoder_layers): ModuleList(\n",
      "      (0-2): 3 x EncoderLayer(\n",
      "        (attention): MultiHeadAttention(\n",
      "          (WQ): Linear(in_features=256, out_features=256, bias=True)\n",
      "          (WK): Linear(in_features=256, out_features=256, bias=True)\n",
      "          (WV): Linear(in_features=256, out_features=256, bias=True)\n",
      "          (WO): Linear(in_features=256, out_features=256, bias=True)\n",
      "          (softmax): Softmax(dim=-1)\n",
      "          (dropout): Dropout(p=0.1, inplace=False)\n",
      "        )\n",
      "        (attn_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
      "        (ff): PositionwiseFeedForward(\n",
      "          (dropout): Dropout(p=0.1, inplace=False)\n",
      "          (W1): Linear(in_features=256, out_features=1024, bias=True)\n",
      "          (W2): Linear(in_features=1024, out_features=256, bias=True)\n",
      "          (relu): ReLU()\n",
      "        )\n",
      "        (ff_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
      "        (dropout): Dropout(p=0.1, inplace=False)\n",
      "      )\n",
      "    )\n",
      "  )\n",
      "  (target_embeddings): Embedding(\n",
      "    (embedding_lookup): Embedding(6500, 256)\n",
      "  )\n",
      "  (target_positional_encoding): PositionalEncoding(\n",
      "    (dropout): Dropout(p=0.1, inplace=False)\n",
      "  )\n",
      "  (decoder): Decoder(\n",
      "    (decoder_layers): ModuleList(\n",
      "      (0-2): 3 x DecoderLayer(\n",
      "        (masked_attention): MultiHeadAttention(\n",
      "          (WQ): Linear(in_features=256, out_features=256, bias=True)\n",
      "          (WK): Linear(in_features=256, out_features=256, bias=True)\n",
      "          (WV): Linear(in_features=256, out_features=256, bias=True)\n",
      "          (WO): Linear(in_features=256, out_features=256, bias=True)\n",
      "          (softmax): Softmax(dim=-1)\n",
      "          (dropout): Dropout(p=0.1, inplace=False)\n",
      "        )\n",
      "        (masked_attn_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
      "        (attention): MultiHeadAttention(\n",
      "          (WQ): Linear(in_features=256, out_features=256, bias=True)\n",
      "          (WK): Linear(in_features=256, out_features=256, bias=True)\n",
      "          (WV): Linear(in_features=256, out_features=256, bias=True)\n",
      "          (WO): Linear(in_features=256, out_features=256, bias=True)\n",
      "          (softmax): Softmax(dim=-1)\n",
      "          (dropout): Dropout(p=0.1, inplace=False)\n",
      "        )\n",
      "        (attn_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
      "        (ff): PositionwiseFeedForward(\n",
      "          (dropout): Dropout(p=0.1, inplace=False)\n",
      "          (W1): Linear(in_features=256, out_features=1024, bias=True)\n",
      "          (W2): Linear(in_features=1024, out_features=256, bias=True)\n",
      "          (relu): ReLU()\n",
      "        )\n",
      "        (ff_norm): LayerNorm((256,), eps=1e-05, elementwise_affine=True)\n",
      "        (dropout): Dropout(p=0.1, inplace=False)\n",
      "      )\n",
      "    )\n",
      "  )\n",
      "  (linear): Linear(in_features=256, out_features=6500, bias=True)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "dk = config[\"dk\"]\n",
    "dv = config[\"dv\"]\n",
    "h = config[\"h\"]\n",
    "src_vocab_size = config[\"src_vocab_size\"]\n",
    "target_vocab_size = config[\"target_vocab_size\"]\n",
    "# src_vocab_size = 8500\n",
    "# target_vocab_size = 6500\n",
    "src_pad_idx = config[\"src_pad_idx\"]\n",
    "target_pad_idx = config[\"target_pad_idx\"]\n",
    "num_encoders = config[\"num_encoders\"]\n",
    "num_decoders = config[\"num_decoders\"]\n",
    "dim_multiplier = config[\"dim_multiplier\"]\n",
    "pdropout = config[\"pdropout\"]\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "model = Transformer(\n",
    "    dk,\n",
    "    dv,\n",
    "    h,\n",
    "    src_vocab_size,\n",
    "    target_vocab_size,\n",
    "    num_encoders,\n",
    "    num_decoders,\n",
    "    dim_multiplier,\n",
    "    pdropout,\n",
    "    device=device,\n",
    ")\n",
    "model.cuda()\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "a69bbb0a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:25:46.115184Z",
     "iopub.status.busy": "2024-01-23T10:25:46.114885Z",
     "iopub.status.idle": "2024-01-23T10:25:46.120730Z",
     "shell.execute_reply": "2024-01-23T10:25:46.119805Z"
    },
    "papermill": {
     "duration": 0.033107,
     "end_time": "2024-01-23T10:25:46.122982",
     "exception": false,
     "start_time": "2024-01-23T10:25:46.089875",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The model has 11,040,100 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",
   "id": "645724b8",
   "metadata": {
    "papermill": {
     "duration": 0.023754,
     "end_time": "2024-01-23T10:25:46.170983",
     "exception": false,
     "start_time": "2024-01-23T10:25:46.147229",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Training part"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "362b8849",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:25:46.220566Z",
     "iopub.status.busy": "2024-01-23T10:25:46.220241Z",
     "iopub.status.idle": "2024-01-23T10:25:46.225799Z",
     "shell.execute_reply": "2024-01-23T10:25:46.224887Z"
    },
    "papermill": {
     "duration": 0.032812,
     "end_time": "2024-01-23T10:25:46.227663",
     "exception": false,
     "start_time": "2024-01-23T10:25:46.194851",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "LEARNING_RATE = config[\"lr\"]\n",
    "\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=LEARNING_RATE)\n",
    "criterion = nn.CrossEntropyLoss(ignore_index=PAD_IDX)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "1eba41af",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:25:46.276700Z",
     "iopub.status.busy": "2024-01-23T10:25:46.276381Z",
     "iopub.status.idle": "2024-01-23T10:25:46.285147Z",
     "shell.execute_reply": "2024-01-23T10:25:46.284431Z"
    },
    "papermill": {
     "duration": 0.035439,
     "end_time": "2024-01-23T10:25:46.287013",
     "exception": false,
     "start_time": "2024-01-23T10:25:46.251574",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def train(model, iterator, optimizer, criterion, clip):\n",
    "    \"\"\"\n",
    "    Train the model on the given data.\n",
    "\n",
    "    Args:\n",
    "        model:        Transformer model to be trained\n",
    "        iterator:     data to be trained on\n",
    "        optimizer:    optimizer for updating parameters\n",
    "        criterion:    loss function for updating parameters\n",
    "        clip:         value to help prevent exploding gradients\n",
    "\n",
    "    Returns:\n",
    "        loss for the epoch\n",
    "    \"\"\"\n",
    "\n",
    "    # set the model to training mode\n",
    "    model.train()\n",
    "\n",
    "    epoch_loss = 0\n",
    "\n",
    "    show_summary = True\n",
    "\n",
    "    # loop through each batch in the iterator\n",
    "    for i, batch in enumerate(iterator):\n",
    "        # set the source and target batches\n",
    "        src, trg = batch\n",
    "\n",
    "        #         print(\"src\", src, src.size())\n",
    "        #         print(\"trg\", trg, trg.size())\n",
    "\n",
    "        #         print(\"trg[:,:-1]\", trg[:,:-1], trg[:,:-1].size())\n",
    "\n",
    "        # expected output\n",
    "        expected_output = trg[:, 1:]\n",
    "        #         print(\"expected_output\", expected_output, expected_output.size())\n",
    "\n",
    "        # zero the gradients\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # Show model summary\n",
    "        if show_summary:\n",
    "            show_summary = False\n",
    "            summary(model, input_data=[src, trg[:, :-1]])\n",
    "\n",
    "        # logits for each output\n",
    "        logits = model(src, trg[:, :-1])\n",
    "\n",
    "        # calculate the loss\n",
    "        loss = criterion(\n",
    "            logits.contiguous().view(-1, logits.shape[-1]),\n",
    "            expected_output.contiguous().view(-1),\n",
    "        )\n",
    "\n",
    "        # backpropagation\n",
    "        loss.backward()\n",
    "\n",
    "        # clip the weights\n",
    "        torch.nn.utils.clip_grad_norm_(model.parameters(), clip)\n",
    "\n",
    "        # update the weights\n",
    "        optimizer.step()\n",
    "\n",
    "        # update the loss\n",
    "        epoch_loss += loss.item()\n",
    "\n",
    "    # return the average loss for the epoch\n",
    "    return epoch_loss / len(iterator)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "78f61790",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:25:46.336134Z",
     "iopub.status.busy": "2024-01-23T10:25:46.335853Z",
     "iopub.status.idle": "2024-01-23T10:25:46.342756Z",
     "shell.execute_reply": "2024-01-23T10:25:46.341874Z"
    },
    "papermill": {
     "duration": 0.033659,
     "end_time": "2024-01-23T10:25:46.344605",
     "exception": false,
     "start_time": "2024-01-23T10:25:46.310946",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def evaluate(model, iterator, criterion):\n",
    "    \"\"\"\n",
    "    Evaluate the model on the given data.\n",
    "\n",
    "    Args:\n",
    "        model:        Transformer model to be trained\n",
    "        iterator:     data to be evaluated\n",
    "        criterion:    loss function for assessing outputs\n",
    "\n",
    "    Returns:\n",
    "        loss for the data\n",
    "    \"\"\"\n",
    "\n",
    "    # set the model to evaluation mode\n",
    "    model.eval()\n",
    "\n",
    "    epoch_loss = 0\n",
    "\n",
    "    # evaluate without updating gradients\n",
    "    with torch.no_grad():\n",
    "        # loop through each batch in the iterator\n",
    "        for i, batch in enumerate(iterator):\n",
    "            # set the source and target batches\n",
    "            src, trg = batch\n",
    "\n",
    "            # expected output\n",
    "            expected_output = trg[:, 1:]\n",
    "\n",
    "            # logits for each output\n",
    "            logits = model(src, trg[:, :-1])\n",
    "\n",
    "            # calculate the loss\n",
    "            loss = criterion(\n",
    "                logits.contiguous().view(-1, logits.shape[-1]),\n",
    "                expected_output.contiguous().view(-1),\n",
    "            )\n",
    "\n",
    "            # update the loss\n",
    "            epoch_loss += loss.item()\n",
    "\n",
    "    # return the average loss for the epoch\n",
    "    return epoch_loss / len(iterator)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "8a59df8e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:25:46.393820Z",
     "iopub.status.busy": "2024-01-23T10:25:46.393539Z",
     "iopub.status.idle": "2024-01-23T10:25:46.397956Z",
     "shell.execute_reply": "2024-01-23T10:25:46.397116Z"
    },
    "papermill": {
     "duration": 0.031096,
     "end_time": "2024-01-23T10:25:46.399911",
     "exception": false,
     "start_time": "2024-01-23T10:25:46.368815",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def epoch_time(start_time, end_time):\n",
    "    elapsed_time = end_time - start_time\n",
    "    elapsed_mins = int(elapsed_time / 60)\n",
    "    elapsed_secs = int(elapsed_time - (elapsed_mins * 60))\n",
    "    return elapsed_mins, elapsed_secs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9becf6c4",
   "metadata": {
    "papermill": {
     "duration": 0.023615,
     "end_time": "2024-01-23T10:25:46.447573",
     "exception": false,
     "start_time": "2024-01-23T10:25:46.423958",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Perform Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "0582b634",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:25:46.496960Z",
     "iopub.status.busy": "2024-01-23T10:25:46.496628Z",
     "iopub.status.idle": "2024-01-23T10:29:26.407666Z",
     "shell.execute_reply": "2024-01-23T10:29:26.406761Z"
    },
    "papermill": {
     "duration": 219.967167,
     "end_time": "2024-01-23T10:29:26.438574",
     "exception": false,
     "start_time": "2024-01-23T10:25:46.471407",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  2%|▏         | 1/50 [00:12<10:25, 12.76s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 01 | Time: 0m 12s\n",
      "\tTrain Loss: 4.850 | Train PPL: 127.757\n",
      "\t Val. Loss: 3.746 |  Val. PPL:  42.362\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  4%|▍         | 2/50 [00:24<09:57, 12.45s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 02 | Time: 0m 12s\n",
      "\tTrain Loss: 3.431 | Train PPL:  30.896\n",
      "\t Val. Loss: 2.979 |  Val. PPL:  19.663\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  6%|▌         | 3/50 [00:37<09:40, 12.36s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 03 | Time: 0m 12s\n",
      "\tTrain Loss: 2.782 | Train PPL:  16.158\n",
      "\t Val. Loss: 2.506 |  Val. PPL:  12.251\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  8%|▊         | 4/50 [00:49<09:29, 12.37s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 04 | Time: 0m 12s\n",
      "\tTrain Loss: 2.362 | Train PPL:  10.611\n",
      "\t Val. Loss: 2.245 |  Val. PPL:   9.440\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 10%|█         | 5/50 [01:01<09:14, 12.32s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 05 | Time: 0m 12s\n",
      "\tTrain Loss: 2.071 | Train PPL:   7.933\n",
      "\t Val. Loss: 2.070 |  Val. PPL:   7.922\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 12%|█▏        | 6/50 [01:14<09:00, 12.28s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 06 | Time: 0m 12s\n",
      "\tTrain Loss: 1.852 | Train PPL:   6.372\n",
      "\t Val. Loss: 1.955 |  Val. PPL:   7.061\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 14%|█▍        | 7/50 [01:26<08:47, 12.26s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 07 | Time: 0m 12s\n",
      "\tTrain Loss: 1.680 | Train PPL:   5.365\n",
      "\t Val. Loss: 1.861 |  Val. PPL:   6.432\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 16%|█▌        | 8/50 [01:38<08:33, 12.23s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 08 | Time: 0m 12s\n",
      "\tTrain Loss: 1.538 | Train PPL:   4.654\n",
      "\t Val. Loss: 1.798 |  Val. PPL:   6.038\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 18%|█▊        | 9/50 [01:50<08:21, 12.23s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 09 | Time: 0m 12s\n",
      "\tTrain Loss: 1.417 | Train PPL:   4.125\n",
      "\t Val. Loss: 1.769 |  Val. PPL:   5.864\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 20%|██        | 10/50 [02:02<08:08, 12.21s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 10 | Time: 0m 12s\n",
      "\tTrain Loss: 1.313 | Train PPL:   3.716\n",
      "\t Val. Loss: 1.745 |  Val. PPL:   5.724\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 22%|██▏       | 11/50 [02:14<07:55, 12.20s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 11 | Time: 0m 12s\n",
      "\tTrain Loss: 1.222 | Train PPL:   3.392\n",
      "\t Val. Loss: 1.730 |  Val. PPL:   5.638\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 24%|██▍       | 12/50 [02:27<07:42, 12.17s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 12 | Time: 0m 12s\n",
      "\tTrain Loss: 1.141 | Train PPL:   3.130\n",
      "\t Val. Loss: 1.746 |  Val. PPL:   5.731\n",
      "Loss did not improve\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 26%|██▌       | 13/50 [02:39<07:30, 12.17s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 13 | Time: 0m 12s\n",
      "\tTrain Loss: 1.067 | Train PPL:   2.906\n",
      "\t Val. Loss: 1.721 |  Val. PPL:   5.591\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 28%|██▊       | 14/50 [02:51<07:17, 12.16s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 14 | Time: 0m 12s\n",
      "\tTrain Loss: 1.000 | Train PPL:   2.718\n",
      "\t Val. Loss: 1.754 |  Val. PPL:   5.778\n",
      "Loss did not improve\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 30%|███       | 15/50 [03:03<07:04, 12.14s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 15 | Time: 0m 12s\n",
      "\tTrain Loss: 0.937 | Train PPL:   2.552\n",
      "\t Val. Loss: 1.747 |  Val. PPL:   5.735\n",
      "Loss did not improve\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 32%|███▏      | 16/50 [03:15<06:52, 12.14s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 16 | Time: 0m 12s\n",
      "\tTrain Loss: 0.883 | Train PPL:   2.417\n",
      "\t Val. Loss: 1.754 |  Val. PPL:   5.777\n",
      "Loss did not improve\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 34%|███▍      | 17/50 [03:27<06:40, 12.14s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 17 | Time: 0m 12s\n",
      "\tTrain Loss: 0.830 | Train PPL:   2.293\n",
      "\t Val. Loss: 1.786 |  Val. PPL:   5.964\n",
      "Loss did not improve\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 34%|███▍      | 17/50 [03:39<07:06, 12.94s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 18 | Time: 0m 12s\n",
      "\tTrain Loss: 0.779 | Train PPL:   2.180\n",
      "\t Val. Loss: 1.809 |  Val. PPL:   6.104\n",
      "Loss did not improve\n",
      "Early Stopping since validation count did not decrease after 5 patience!\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "N_EPOCHS = config[\"N_EPOCHS\"]\n",
    "CLIP = config[\"CLIP\"]\n",
    "patience = config[\"patience\"]\n",
    "\n",
    "best_valid_loss = float(\"inf\")\n",
    "\n",
    "patience_count = 0\n",
    "\n",
    "# loop through each epoch\n",
    "for epoch in tqdm(range(N_EPOCHS)):\n",
    "    start_time = time.time()\n",
    "\n",
    "    # calculate the train loss and update the parameters\n",
    "    train_loss = train(model, train_iter, optimizer, criterion, CLIP)\n",
    "\n",
    "    # calculate the loss on the validation set\n",
    "    valid_loss = evaluate(model, valid_iter, criterion)\n",
    "\n",
    "    end_time = time.time()\n",
    "\n",
    "    # calculate how long the epoch took\n",
    "    epoch_mins, epoch_secs = epoch_time(start_time, end_time)\n",
    "\n",
    "    print(f\"Epoch: {epoch+1:02} | Time: {epoch_mins}m {epoch_secs}s\")\n",
    "    print(f\"\\tTrain Loss: {train_loss:.3f} | Train PPL: {math.exp(train_loss):7.3f}\")\n",
    "    print(f\"\\t Val. Loss: {valid_loss:.3f} |  Val. PPL: {math.exp(valid_loss):7.3f}\")\n",
    "\n",
    "    # save the model when it performs better than the previous run\n",
    "    if valid_loss < best_valid_loss:\n",
    "        best_valid_loss = valid_loss\n",
    "        torch.save(model.state_dict(), \"transformer-model.pt\")\n",
    "        patience_count = 0\n",
    "    else:\n",
    "        patience_count += 1\n",
    "        print(f\"Loss did not improve\")\n",
    "\n",
    "    if patience is not None and patience_count == patience:\n",
    "        print(\n",
    "            f\"Early Stopping since validation count did not decrease after {patience_count} patience!\"\n",
    "        )\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "1530e17e",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:29:26.493717Z",
     "iopub.status.busy": "2024-01-23T10:29:26.493361Z",
     "iopub.status.idle": "2024-01-23T10:29:26.747764Z",
     "shell.execute_reply": "2024-01-23T10:29:26.746792Z"
    },
    "papermill": {
     "duration": 0.284305,
     "end_time": "2024-01-23T10:29:26.749872",
     "exception": false,
     "start_time": "2024-01-23T10:29:26.465567",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test Loss: 1.723 | Test PPL:   5.601\n"
     ]
    }
   ],
   "source": [
    "model.load_state_dict(torch.load(\"transformer-model.pt\"))\n",
    "\n",
    "# calculate the loss on the test set\n",
    "test_loss = evaluate(model, test_iter, criterion)\n",
    "\n",
    "print(f\"Test Loss: {test_loss:.3f} | Test PPL: {math.exp(test_loss):7.3f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "56650185",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:29:26.806200Z",
     "iopub.status.busy": "2024-01-23T10:29:26.805860Z",
     "iopub.status.idle": "2024-01-23T10:29:26.816330Z",
     "shell.execute_reply": "2024-01-23T10:29:26.815345Z"
    },
    "papermill": {
     "duration": 0.041466,
     "end_time": "2024-01-23T10:29:26.819042",
     "exception": false,
     "start_time": "2024-01-23T10:29:26.777576",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def translate_sentence(sentence, model, device, max_length=50):\n",
    "    \"\"\"\n",
    "    Translate a German sentence to its English equivalent.\n",
    "\n",
    "    Args:\n",
    "        sentence:     German sentence to be translated to English; list or str\n",
    "        model:        Transformer model used for translation\n",
    "        device:       device to perform translation on\n",
    "        max_length:   maximum token length for translation\n",
    "\n",
    "    Returns:\n",
    "        src:                  return the tokenized input\n",
    "        trg_input:            return the input to the decoder before the final output\n",
    "        trg_output:           return the final translation, shifted right\n",
    "        attn_probs:           return the attention scores for the decoder heads\n",
    "        masked_attn_probs:    return the masked attention scores for the decoder heads\n",
    "    \"\"\"\n",
    "\n",
    "    model.eval()\n",
    "\n",
    "    # tokenize and index the provided string\n",
    "    if isinstance(sentence, str):\n",
    "        src = (\n",
    "            [\"<bos>\"] + [token.text.lower() for token in spacy_de(sentence)] + [\"<eos>\"]\n",
    "        )\n",
    "    else:\n",
    "        src = [\"<bos>\"] + sentence + [\"<eos>\"]\n",
    "\n",
    "    # convert to integers\n",
    "    src_indexes = [vocab_src[token] for token in src]\n",
    "\n",
    "    # convert list to tensor\n",
    "    src_tensor = torch.tensor(src_indexes).int().unsqueeze(0).to(device)\n",
    "\n",
    "    # set <bos> token for target generation\n",
    "    trg_indexes = [vocab_trg.get_stoi()[\"<bos>\"]]\n",
    "\n",
    "    # generate new tokens\n",
    "    for i in range(max_length):\n",
    "        # convert the list to a tensor\n",
    "        trg_tensor = torch.tensor(trg_indexes).int().unsqueeze(0).to(device)\n",
    "\n",
    "        # generate the next token\n",
    "        with torch.no_grad():\n",
    "            # generate the logits\n",
    "            logits = model.forward(src_tensor, trg_tensor)\n",
    "\n",
    "            # select the newly predicted token\n",
    "            pred_token = logits.argmax(2)[:, -1].item()\n",
    "\n",
    "            # if <eos> token or max length, stop generating\n",
    "            if pred_token == vocab_trg.get_stoi()[\"<eos>\"] or i == (max_length - 1):\n",
    "                # decoder input\n",
    "                trg_input = vocab_trg.lookup_tokens(trg_indexes)\n",
    "\n",
    "                # decoder output\n",
    "                trg_output = vocab_trg.lookup_tokens(\n",
    "                    logits.argmax(2).squeeze(0).tolist()\n",
    "                )\n",
    "\n",
    "                return src, trg_input, trg_output\n",
    "\n",
    "            # else, continue generating\n",
    "            else:\n",
    "                # add the token\n",
    "                trg_indexes.append(pred_token)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "4367f2c5",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:29:26.874610Z",
     "iopub.status.busy": "2024-01-23T10:29:26.874300Z",
     "iopub.status.idle": "2024-01-23T10:29:27.055499Z",
     "shell.execute_reply": "2024-01-23T10:29:27.054221Z"
    },
    "papermill": {
     "duration": 0.21117,
     "end_time": "2024-01-23T10:29:27.057521",
     "exception": false,
     "start_time": "2024-01-23T10:29:26.846351",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "source = ['<bos>', 'eine', 'frau', 'mit', 'einer', 'großen', 'geldbörse', 'geht', 'an', 'einem', 'tor', 'vorbei', '.', '<eos>']\n",
      "target input = ['<bos>', 'a', 'woman', 'with', 'a', 'large', 'purse', 'walks', 'past', 'a', 'gate', 'by', 'a', 'gate', '.']\n",
      "target output = ['a', 'woman', 'with', 'a', 'large', 'purse', 'walks', 'past', 'a', 'gate', 'by', 'a', 'gate', '.', '<eos>']\n"
     ]
    }
   ],
   "source": [
    "# 'a woman with a large purse is walking by a gate'\n",
    "src = [\n",
    "    \"eine\",\n",
    "    \"frau\",\n",
    "    \"mit\",\n",
    "    \"einer\",\n",
    "    \"großen\",\n",
    "    \"geldbörse\",\n",
    "    \"geht\",\n",
    "    \"an\",\n",
    "    \"einem\",\n",
    "    \"tor\",\n",
    "    \"vorbei\",\n",
    "    \".\",\n",
    "]\n",
    "trg = [\"a\", \"woman\", \"with\", \"a\", \"large\", \"purse\", \"is\", \"walking\", \"by\", \"a\", \"gate\"]\n",
    "\n",
    "src, trg_input, trg_output = translate_sentence(src, model, device)\n",
    "\n",
    "print(f\"source = {src}\")\n",
    "print(f\"target input = {trg_input}\")\n",
    "print(f\"target output = {trg_output}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a66ebbd3",
   "metadata": {
    "papermill": {
     "duration": 0.027364,
     "end_time": "2024-01-23T10:29:27.112932",
     "exception": false,
     "start_time": "2024-01-23T10:29:27.085568",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Calculate Bleu Score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "f2c82f89",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:29:27.169686Z",
     "iopub.status.busy": "2024-01-23T10:29:27.169189Z",
     "iopub.status.idle": "2024-01-23T10:29:27.177030Z",
     "shell.execute_reply": "2024-01-23T10:29:27.176057Z"
    },
    "papermill": {
     "duration": 0.038849,
     "end_time": "2024-01-23T10:29:27.179129",
     "exception": false,
     "start_time": "2024-01-23T10:29:27.140280",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def compute_metrics(model, iterator):\n",
    "    \"\"\"\n",
    "    Generate predictions for the provided iterator.\n",
    "\n",
    "    Args:\n",
    "        model:        Transformer model to be trained\n",
    "        iterator:     data to be evaluated\n",
    "\n",
    "    Returns:\n",
    "        predictions:  list of predictions, which are tokenized strings\n",
    "        labels:       list of expected output, which are tokenized strings\n",
    "    \"\"\"\n",
    "\n",
    "    # set the model to evaluation mode\n",
    "    model.eval()\n",
    "\n",
    "    predictions = []\n",
    "    labels = []\n",
    "\n",
    "    # evaluate without updating gradients\n",
    "    with torch.no_grad():\n",
    "        # loop through each batch in the iterator\n",
    "        for i, batch in tqdm(enumerate(iterator)):\n",
    "            # set the source and target batches\n",
    "            src, trg = batch\n",
    "\n",
    "            # predict the output\n",
    "            src_out, trg_input, trg_output = translate_sentence(\n",
    "                vocab_src.lookup_tokens(src.tolist()), model, device\n",
    "            )\n",
    "\n",
    "            # prediction | remove <eos> token\n",
    "            predictions.append(trg_output[:-1])\n",
    "\n",
    "            # expected output | add extra dim for calculation\n",
    "            labels.append([vocab_trg.lookup_tokens(trg.tolist())])\n",
    "\n",
    "    # return the average loss for the epoch\n",
    "    return predictions, labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "faf70141",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:29:27.235599Z",
     "iopub.status.busy": "2024-01-23T10:29:27.235233Z",
     "iopub.status.idle": "2024-01-23T10:31:44.355647Z",
     "shell.execute_reply": "2024-01-23T10:31:44.354619Z"
    },
    "papermill": {
     "duration": 137.151516,
     "end_time": "2024-01-23T10:31:44.358023",
     "exception": false,
     "start_time": "2024-01-23T10:29:27.206507",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "1000it [02:17,  7.29it/s]\n"
     ]
    }
   ],
   "source": [
    "predictions, labels = compute_metrics(model, test_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "80d983a3",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:31:44.553277Z",
     "iopub.status.busy": "2024-01-23T10:31:44.552948Z",
     "iopub.status.idle": "2024-01-23T10:31:44.987640Z",
     "shell.execute_reply": "2024-01-23T10:31:44.986306Z"
    },
    "papermill": {
     "duration": 0.534271,
     "end_time": "2024-01-23T10:31:44.989929",
     "exception": false,
     "start_time": "2024-01-23T10:31:44.455658",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Bleu score on test data is 30.82650601863861\n"
     ]
    }
   ],
   "source": [
    "from torchtext.data.metrics import bleu_score\n",
    "\n",
    "score = bleu_score(predictions, labels)\n",
    "score = score * 100\n",
    "print(f\"Bleu score on test data is {score}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "1b05f5f6",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:31:45.202392Z",
     "iopub.status.busy": "2024-01-23T10:31:45.201957Z",
     "iopub.status.idle": "2024-01-23T10:31:45.208663Z",
     "shell.execute_reply": "2024-01-23T10:31:45.207750Z"
    },
    "papermill": {
     "duration": 0.11707,
     "end_time": "2024-01-23T10:31:45.210551",
     "exception": false,
     "start_time": "2024-01-23T10:31:45.093481",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'30.8'"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "str_score = f\"{score:.1f}\"\n",
    "str_score"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "999bc4f5",
   "metadata": {
    "papermill": {
     "duration": 0.100601,
     "end_time": "2024-01-23T10:31:45.411099",
     "exception": false,
     "start_time": "2024-01-23T10:31:45.310498",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Save model and input param"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "ac03652b",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:31:45.607318Z",
     "iopub.status.busy": "2024-01-23T10:31:45.606386Z",
     "iopub.status.idle": "2024-01-23T10:31:45.708338Z",
     "shell.execute_reply": "2024-01-23T10:31:45.707299Z"
    },
    "papermill": {
     "duration": 0.203409,
     "end_time": "2024-01-23T10:31:45.710981",
     "exception": false,
     "start_time": "2024-01-23T10:31:45.507572",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "torch.save(model.state_dict(), f\"transformer-model-{str_score}.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "acd1a642",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:31:45.908235Z",
     "iopub.status.busy": "2024-01-23T10:31:45.907916Z",
     "iopub.status.idle": "2024-01-23T10:31:45.915228Z",
     "shell.execute_reply": "2024-01-23T10:31:45.914416Z"
    },
    "papermill": {
     "duration": 0.106267,
     "end_time": "2024-01-23T10:31:45.917080",
     "exception": false,
     "start_time": "2024-01-23T10:31:45.810813",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'dk': 32,\n",
       " 'dv': 32,\n",
       " 'h': 8,\n",
       " 'src_vocab_size': 8500,\n",
       " 'target_vocab_size': 6500,\n",
       " 'src_pad_idx': 2,\n",
       " 'target_pad_idx': 2,\n",
       " 'num_encoders': 3,\n",
       " 'num_decoders': 3,\n",
       " 'dim_multiplier': 4,\n",
       " 'pdropout': 0.1,\n",
       " 'lr': 0.0003,\n",
       " 'N_EPOCHS': 50,\n",
       " 'CLIP': 1,\n",
       " 'patience': 5}"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import json\n",
    "\n",
    "# Convert and write JSON object to file\n",
    "with open(\"input_params.json\", \"w\") as outfile:\n",
    "    json.dump(config, outfile)\n",
    "config"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4e496c4b",
   "metadata": {
    "papermill": {
     "duration": 0.09831,
     "end_time": "2024-01-23T10:31:46.112658",
     "exception": false,
     "start_time": "2024-01-23T10:31:46.014348",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Store model on HuggingFace Hub"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "b5384258",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:31:46.309635Z",
     "iopub.status.busy": "2024-01-23T10:31:46.308998Z",
     "iopub.status.idle": "2024-01-23T10:31:46.313095Z",
     "shell.execute_reply": "2024-01-23T10:31:46.312192Z"
    },
    "papermill": {
     "duration": 0.104491,
     "end_time": "2024-01-23T10:31:46.315296",
     "exception": false,
     "start_time": "2024-01-23T10:31:46.210805",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# ! pip install --upgrade huggingface_hub -q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "b19afe12",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:31:46.513435Z",
     "iopub.status.busy": "2024-01-23T10:31:46.513078Z",
     "iopub.status.idle": "2024-01-23T10:31:46.517119Z",
     "shell.execute_reply": "2024-01-23T10:31:46.516164Z"
    },
    "papermill": {
     "duration": 0.105331,
     "end_time": "2024-01-23T10:31:46.519052",
     "exception": false,
     "start_time": "2024-01-23T10:31:46.413721",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# from huggingface_hub import notebook_login\n",
    "# notebook_login()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "0f95d318",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:31:46.714223Z",
     "iopub.status.busy": "2024-01-23T10:31:46.713874Z",
     "iopub.status.idle": "2024-01-23T10:31:46.717801Z",
     "shell.execute_reply": "2024-01-23T10:31:46.716934Z"
    },
    "papermill": {
     "duration": 0.103024,
     "end_time": "2024-01-23T10:31:46.719589",
     "exception": false,
     "start_time": "2024-01-23T10:31:46.616565",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# from huggingface_hub import HfApi\n",
    "# api = HfApi()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "eefc6837",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-01-23T10:31:46.914852Z",
     "iopub.status.busy": "2024-01-23T10:31:46.914495Z",
     "iopub.status.idle": "2024-01-23T10:31:46.918521Z",
     "shell.execute_reply": "2024-01-23T10:31:46.917637Z"
    },
    "papermill": {
     "duration": 0.10462,
     "end_time": "2024-01-23T10:31:46.920414",
     "exception": false,
     "start_time": "2024-01-23T10:31:46.815794",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# api.upload_folder(\n",
    "#     folder_path=\"/kaggle/working/.\",\n",
    "#     repo_id=\"Rzoro/Transformer_de_en_multi30K\",\n",
    "#     repo_type=\"space\",\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "507b1591",
   "metadata": {
    "papermill": {
     "duration": 0.09659,
     "end_time": "2024-01-23T10:31:47.114073",
     "exception": false,
     "start_time": "2024-01-23T10:31:47.017483",
     "status": "completed"
    },
    "tags": []
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kaggle": {
   "accelerator": "gpu",
   "dataSources": [
    {
     "datasetId": 1800581,
     "sourceId": 2936819,
     "sourceType": "datasetVersion"
    }
   ],
   "dockerImageVersionId": 30636,
   "isGpuEnabled": true,
   "isInternetEnabled": true,
   "language": "python",
   "sourceType": "notebook"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  },
  "papermill": {
   "default_parameters": {},
   "duration": 418.352794,
   "end_time": "2024-01-23T10:31:48.734295",
   "environment_variables": {},
   "exception": null,
   "input_path": "__notebook__.ipynb",
   "output_path": "__notebook__.ipynb",
   "parameters": {},
   "start_time": "2024-01-23T10:24:50.381501",
   "version": "2.4.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}