nmt-transformers-glove.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "93ca8d46-0e2e-4d9b-b9a7-de35eb0763cb",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\father\\.conda\\envs\\maq\\lib\\site-packages\\tqdm\\auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import random\n",
    "import math\n",
    "import time\n",
    "import gc\n",
    "import copy\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.nn import Transformer\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "import torchtext\n",
    "from torchtext.vocab import GloVe\n",
    "from torchtext.data.metrics import bleu_score\n",
    "\n",
    "import spacy\n",
    "import numpy as np\n",
    "\n",
    "from tqdm import tqdm\n",
    "\n",
    "from dataset import nmtDataset\n",
    "import helpers as utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6e93a05a-37ec-49cf-9fd6-8920dd074953",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "# Seeding for consistency in reproducibility\n",
    "SEED = 1234\n",
    "\n",
    "random.seed(SEED)\n",
    "torch.manual_seed(SEED)\n",
    "torch.cuda.manual_seed(SEED)\n",
    "torch.backends.cudnn.deterministic = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b016e667-484b-41f7-89b4-4e3f8b8da896",
   "metadata": {},
   "outputs": [],
   "source": [
    "log = utils.Logger('logs/transformers-glove.out')  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "98efca44-4448-448d-b1b5-93f60bb40852",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting en-core-web-sm==3.3.0\n",
      "  Downloading https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.3.0/en_core_web_sm-3.3.0-py3-none-any.whl (12.8 MB)\n",
      "Requirement already satisfied: spacy<3.4.0,>=3.3.0.dev0 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from en-core-web-sm==3.3.0) (3.3.0)\n",
      "Requirement already satisfied: requests<3.0.0,>=2.13.0 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from spacy<3.4.0,>=3.3.0.dev0->en-core-web-sm==3.3.0) (2.27.1)\n",
      "Requirement already satisfied: jinja2 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from spacy<3.4.0,>=3.3.0.dev0->en-core-web-sm==3.3.0) (3.1.2)\n",
      "Requirement already satisfied: spacy-legacy<3.1.0,>=3.0.9 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from spacy<3.4.0,>=3.3.0.dev0->en-core-web-sm==3.3.0) (3.0.9)\n",
      "Requirement already satisfied: setuptools in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from spacy<3.4.0,>=3.3.0.dev0->en-core-web-sm==3.3.0) (61.2.0)\n",
      "Requirement already satisfied: srsly<3.0.0,>=2.4.3 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from spacy<3.4.0,>=3.3.0.dev0->en-core-web-sm==3.3.0) (2.4.3)\n",
      "Requirement already satisfied: cymem<2.1.0,>=2.0.2 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from spacy<3.4.0,>=3.3.0.dev0->en-core-web-sm==3.3.0) (2.0.6)\n",
      "Requirement already satisfied: pathy>=0.3.5 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from spacy<3.4.0,>=3.3.0.dev0->en-core-web-sm==3.3.0) (0.6.1)\n",
      "Requirement already satisfied: preshed<3.1.0,>=3.0.2 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from spacy<3.4.0,>=3.3.0.dev0->en-core-web-sm==3.3.0) (3.0.6)\n",
      "Requirement already satisfied: spacy-loggers<2.0.0,>=1.0.0 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from spacy<3.4.0,>=3.3.0.dev0->en-core-web-sm==3.3.0) (1.0.2)\n",
      "Requirement already satisfied: catalogue<2.1.0,>=2.0.6 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from spacy<3.4.0,>=3.3.0.dev0->en-core-web-sm==3.3.0) (2.0.6)\n",
      "Requirement already satisfied: packaging>=20.0 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from spacy<3.4.0,>=3.3.0.dev0->en-core-web-sm==3.3.0) (21.3)\n",
      "Requirement already satisfied: blis<0.8.0,>=0.4.0 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from spacy<3.4.0,>=3.3.0.dev0->en-core-web-sm==3.3.0) (0.7.7)\n",
      "Requirement already satisfied: langcodes<4.0.0,>=3.2.0 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from spacy<3.4.0,>=3.3.0.dev0->en-core-web-sm==3.3.0) (3.3.0)\n",
      "Requirement already satisfied: murmurhash<1.1.0,>=0.28.0 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from spacy<3.4.0,>=3.3.0.dev0->en-core-web-sm==3.3.0) (1.0.7)\n",
      "Requirement already satisfied: numpy>=1.15.0 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from spacy<3.4.0,>=3.3.0.dev0->en-core-web-sm==3.3.0) (1.22.3)\n",
      "Requirement already satisfied: wasabi<1.1.0,>=0.9.1 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from spacy<3.4.0,>=3.3.0.dev0->en-core-web-sm==3.3.0) (0.9.1)\n",
      "Requirement already satisfied: thinc<8.1.0,>=8.0.14 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from spacy<3.4.0,>=3.3.0.dev0->en-core-web-sm==3.3.0) (8.0.17)\n",
      "Requirement already satisfied: typer<0.5.0,>=0.3.0 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from spacy<3.4.0,>=3.3.0.dev0->en-core-web-sm==3.3.0) (0.4.1)\n",
      "Requirement already satisfied: pydantic!=1.8,!=1.8.1,<1.9.0,>=1.7.4 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from spacy<3.4.0,>=3.3.0.dev0->en-core-web-sm==3.3.0) (1.8.2)\n",
      "Requirement already satisfied: tqdm<5.0.0,>=4.38.0 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from spacy<3.4.0,>=3.3.0.dev0->en-core-web-sm==3.3.0) (4.64.0)\n",
      "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from packaging>=20.0->spacy<3.4.0,>=3.3.0.dev0->en-core-web-sm==3.3.0) (3.0.9)\n",
      "Requirement already satisfied: smart-open<6.0.0,>=5.0.0 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from pathy>=0.3.5->spacy<3.4.0,>=3.3.0.dev0->en-core-web-sm==3.3.0) (5.2.1)\n",
      "Requirement already satisfied: typing-extensions>=3.7.4.3 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from pydantic!=1.8,!=1.8.1,<1.9.0,>=1.7.4->spacy<3.4.0,>=3.3.0.dev0->en-core-web-sm==3.3.0) (4.1.1)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from requests<3.0.0,>=2.13.0->spacy<3.4.0,>=3.3.0.dev0->en-core-web-sm==3.3.0) (1.26.9)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from requests<3.0.0,>=2.13.0->spacy<3.4.0,>=3.3.0.dev0->en-core-web-sm==3.3.0) (2022.5.18.1)\n",
      "Requirement already satisfied: charset-normalizer~=2.0.0 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from requests<3.0.0,>=2.13.0->spacy<3.4.0,>=3.3.0.dev0->en-core-web-sm==3.3.0) (2.0.4)\n",
      "Requirement already satisfied: idna<4,>=2.5 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from requests<3.0.0,>=2.13.0->spacy<3.4.0,>=3.3.0.dev0->en-core-web-sm==3.3.0) (3.3)\n",
      "Requirement already satisfied: colorama in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from tqdm<5.0.0,>=4.38.0->spacy<3.4.0,>=3.3.0.dev0->en-core-web-sm==3.3.0) (0.4.4)\n",
      "Requirement already satisfied: click<9.0.0,>=7.1.1 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from typer<0.5.0,>=0.3.0->spacy<3.4.0,>=3.3.0.dev0->en-core-web-sm==3.3.0) (7.1.2)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from jinja2->spacy<3.4.0,>=3.3.0.dev0->en-core-web-sm==3.3.0) (2.1.1)\n",
      "[+] Download and installation successful\n",
      "You can now load the package via spacy.load('en_core_web_sm')\n",
      "Collecting de-core-news-sm==3.3.0\n",
      "  Downloading https://github.com/explosion/spacy-models/releases/download/de_core_news_sm-3.3.0/de_core_news_sm-3.3.0-py3-none-any.whl (14.6 MB)\n",
      "Requirement already satisfied: spacy<3.4.0,>=3.3.0.dev0 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from de-core-news-sm==3.3.0) (3.3.0)\n",
      "Requirement already satisfied: setuptools in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from spacy<3.4.0,>=3.3.0.dev0->de-core-news-sm==3.3.0) (61.2.0)\n",
      "Requirement already satisfied: spacy-legacy<3.1.0,>=3.0.9 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from spacy<3.4.0,>=3.3.0.dev0->de-core-news-sm==3.3.0) (3.0.9)\n",
      "Requirement already satisfied: murmurhash<1.1.0,>=0.28.0 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from spacy<3.4.0,>=3.3.0.dev0->de-core-news-sm==3.3.0) (1.0.7)\n",
      "Requirement already satisfied: srsly<3.0.0,>=2.4.3 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from spacy<3.4.0,>=3.3.0.dev0->de-core-news-sm==3.3.0) (2.4.3)\n",
      "Requirement already satisfied: blis<0.8.0,>=0.4.0 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from spacy<3.4.0,>=3.3.0.dev0->de-core-news-sm==3.3.0) (0.7.7)\n",
      "Requirement already satisfied: numpy>=1.15.0 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from spacy<3.4.0,>=3.3.0.dev0->de-core-news-sm==3.3.0) (1.22.3)\n",
      "Requirement already satisfied: cymem<2.1.0,>=2.0.2 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from spacy<3.4.0,>=3.3.0.dev0->de-core-news-sm==3.3.0) (2.0.6)\n",
      "Requirement already satisfied: spacy-loggers<2.0.0,>=1.0.0 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from spacy<3.4.0,>=3.3.0.dev0->de-core-news-sm==3.3.0) (1.0.2)\n",
      "Requirement already satisfied: thinc<8.1.0,>=8.0.14 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from spacy<3.4.0,>=3.3.0.dev0->de-core-news-sm==3.3.0) (8.0.17)\n",
      "Requirement already satisfied: typer<0.5.0,>=0.3.0 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from spacy<3.4.0,>=3.3.0.dev0->de-core-news-sm==3.3.0) (0.4.1)\n",
      "Requirement already satisfied: wasabi<1.1.0,>=0.9.1 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from spacy<3.4.0,>=3.3.0.dev0->de-core-news-sm==3.3.0) (0.9.1)\n",
      "Requirement already satisfied: pathy>=0.3.5 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from spacy<3.4.0,>=3.3.0.dev0->de-core-news-sm==3.3.0) (0.6.1)\n",
      "Requirement already satisfied: preshed<3.1.0,>=3.0.2 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from spacy<3.4.0,>=3.3.0.dev0->de-core-news-sm==3.3.0) (3.0.6)\n",
      "Requirement already satisfied: tqdm<5.0.0,>=4.38.0 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from spacy<3.4.0,>=3.3.0.dev0->de-core-news-sm==3.3.0) (4.64.0)\n",
      "Requirement already satisfied: catalogue<2.1.0,>=2.0.6 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from spacy<3.4.0,>=3.3.0.dev0->de-core-news-sm==3.3.0) (2.0.6)\n",
      "Requirement already satisfied: langcodes<4.0.0,>=3.2.0 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from spacy<3.4.0,>=3.3.0.dev0->de-core-news-sm==3.3.0) (3.3.0)\n",
      "Requirement already satisfied: requests<3.0.0,>=2.13.0 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from spacy<3.4.0,>=3.3.0.dev0->de-core-news-sm==3.3.0) (2.27.1)\n",
      "Requirement already satisfied: pydantic!=1.8,!=1.8.1,<1.9.0,>=1.7.4 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from spacy<3.4.0,>=3.3.0.dev0->de-core-news-sm==3.3.0) (1.8.2)\n",
      "Requirement already satisfied: jinja2 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from spacy<3.4.0,>=3.3.0.dev0->de-core-news-sm==3.3.0) (3.1.2)\n",
      "Requirement already satisfied: packaging>=20.0 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from spacy<3.4.0,>=3.3.0.dev0->de-core-news-sm==3.3.0) (21.3)\n",
      "Requirement already satisfied: pyparsing!=3.0.5,>=2.0.2 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from packaging>=20.0->spacy<3.4.0,>=3.3.0.dev0->de-core-news-sm==3.3.0) (3.0.9)\n",
      "Requirement already satisfied: smart-open<6.0.0,>=5.0.0 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from pathy>=0.3.5->spacy<3.4.0,>=3.3.0.dev0->de-core-news-sm==3.3.0) (5.2.1)\n",
      "Requirement already satisfied: typing-extensions>=3.7.4.3 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from pydantic!=1.8,!=1.8.1,<1.9.0,>=1.7.4->spacy<3.4.0,>=3.3.0.dev0->de-core-news-sm==3.3.0) (4.1.1)\n",
      "Requirement already satisfied: idna<4,>=2.5 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from requests<3.0.0,>=2.13.0->spacy<3.4.0,>=3.3.0.dev0->de-core-news-sm==3.3.0) (3.3)\n",
      "Requirement already satisfied: urllib3<1.27,>=1.21.1 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from requests<3.0.0,>=2.13.0->spacy<3.4.0,>=3.3.0.dev0->de-core-news-sm==3.3.0) (1.26.9)\n",
      "Requirement already satisfied: charset-normalizer~=2.0.0 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from requests<3.0.0,>=2.13.0->spacy<3.4.0,>=3.3.0.dev0->de-core-news-sm==3.3.0) (2.0.4)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from requests<3.0.0,>=2.13.0->spacy<3.4.0,>=3.3.0.dev0->de-core-news-sm==3.3.0) (2022.5.18.1)\n",
      "Requirement already satisfied: colorama in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from tqdm<5.0.0,>=4.38.0->spacy<3.4.0,>=3.3.0.dev0->de-core-news-sm==3.3.0) (0.4.4)\n",
      "Requirement already satisfied: click<9.0.0,>=7.1.1 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from typer<0.5.0,>=0.3.0->spacy<3.4.0,>=3.3.0.dev0->de-core-news-sm==3.3.0) (7.1.2)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in c:\\users\\father\\.conda\\envs\\maq\\lib\\site-packages (from jinja2->spacy<3.4.0,>=3.3.0.dev0->de-core-news-sm==3.3.0) (2.1.1)\n",
      "[+] Download and installation successful\n",
      "You can now load the package via spacy.load('de_core_news_sm')\n"
     ]
    }
   ],
   "source": [
    "!python -m spacy download en_core_web_sm\n",
    "!python -m spacy download de_core_news_sm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "13d5143e-4bb3-4fe1-8b7d-cc0fdf976d69",
   "metadata": {},
   "outputs": [],
   "source": [
    "hyp_params = {\n",
    "    \"batch_size\": 128,\n",
    "    \"lr\": 0.0005,\n",
    "    \"num_epochs\": 10,\n",
    "    \n",
    "    # Same as presented in paper\n",
    "    \"d_model\": 300,\n",
    "    \n",
    "    # No. of multi-head attention block (aka paralle self-attention layers)\n",
    "    # Same as presented in paper\n",
    "    \"n_head\": 6,\n",
    "    \n",
    "    # N in the paper and they used 6 of each\n",
    "    \"num_encoder_layers\": 3,\n",
    "    \"num_decoder_layers\": 3,\n",
    "    \n",
    "    \"feedforward_dim\": 128,\n",
    "    \n",
    "    # Following paper\n",
    "    \"dropout\": 0.1\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3407d825-9a5f-4b71-a315-825b17256f6d",
   "metadata": {},
   "outputs": [],
   "source": [
    "nmtds_train = nmtDataset('datasets/Multi30k/', 'train')\n",
    "nmtds_valid = nmtDataset('datasets/Multi30k/', 'val', nmtds_train)\n",
    "nmtds_test = nmtDataset('datasets/Multi30k/', 'test', nmtds_train)\n",
    "\n",
    "SRC_PAD_IDX = nmtds_train.src_vocab[\"<pad>\"]\n",
    "TRG_PAD_IDX = nmtds_train.trg_vocab[\"<pad>\"]\n",
    "\n",
    "train_dataloader = DataLoader(nmtds_train, batch_size=hyp_params[\"batch_size\"], shuffle=True,\n",
    "                              collate_fn=lambda batch_size: utils.collate_fn(batch_size, SRC_PAD_IDX, device))\n",
    "\n",
    "valid_dataloader = DataLoader(nmtds_valid, batch_size=hyp_params[\"batch_size\"], shuffle=True,\n",
    "                              collate_fn=lambda batch_size: utils.collate_fn(batch_size, SRC_PAD_IDX, device))\n",
    "\n",
    "hyp_params[\"src_vocab_size\"] = len(nmtds_train.src_vocab)\n",
    "hyp_params[\"trg_vocab_size\"] = len(nmtds_train.trg_vocab)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "7e79d461-7517-46a9-95db-581de366e977",
   "metadata": {},
   "outputs": [],
   "source": [
    "class PositionalEncoding(nn.Module):\n",
    "    def __init__(self, d_model, dropout, maxlen = 5000):\n",
    "        super(PositionalEncoding, self).__init__()\n",
    "        \n",
    "        # A tensor consists of all the possible positions (index) e.g 0, 1, 2, ... max length of input\n",
    "        # Shape (pos) --> [max len, 1]\n",
    "        pos = torch.arange(0, maxlen).unsqueeze(1)\n",
    "        \n",
    "        pos_encoding = torch.zeros((maxlen, d_model))\n",
    "        \n",
    "        # In the paper, they had 2i in the positional encoding formula\n",
    "        # where i is the dimension \n",
    "        sin_den = 10000 ** (torch.arange(0, d_model, 2)/d_model) # sin for even item of position's dimension\n",
    "        cos_den = 10000 ** (torch.arange(1, d_model, 2)/d_model) # cos for odd \n",
    "        \n",
    "        pos_encoding[:, 0::2] = torch.sin(pos / sin_den) \n",
    "        pos_encoding[:, 1::2] = torch.cos(pos / cos_den)\n",
    "        \n",
    "        # Shape (pos_embedding) --> [max len, d_model]\n",
    "        pos_encoding = pos_encoding.unsqueeze(-2)\n",
    "        # Shape (pos_embedding) --> [max len, 1, d_model]\n",
    "\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "        # We want pos_encoding be saved and restored in the `state_dict`, but not trained by the optimizer\n",
    "        # hence registering it!\n",
    "        # Source & credits: https://discuss.pytorch.org/t/what-is-the-difference-between-register-buffer-and-register-parameter-of-nn-module/32723/2\n",
    "        self.register_buffer('pos_encoding', pos_encoding)\n",
    "\n",
    "    def forward(self, token_embedding):\n",
    "        \n",
    "        # shape (token_embedding) --> [sentence len, batch size, d_model]\n",
    "        \n",
    "        # Concatenating embeddings with positional encodings\n",
    "        # Note: As we made positional encoding with the size max length of sentence in our dataset \n",
    "        #       hence here we are picking till the sentence length in a batch\n",
    "        #       Another thing to notice is in the paper they used FIXED positional encoding, there are\n",
    "        #       methods where we can also learn them but we are doing as presented in the paper\n",
    "        return self.dropout(token_embedding + self.pos_encoding[:token_embedding.size(0), :])\n",
    "    \n",
    "class InputEmbedding(nn.Module):\n",
    "    def __init__(self, vocab_size, d_model, src_vocab=None):\n",
    "        super(InputEmbedding, self).__init__()\n",
    "        \n",
    "        if src_vocab:\n",
    "            glove_vectors = GloVe(name=\"6B\")\n",
    "            vecs = glove_vectors.get_vecs_by_tokens(src_vocab.get_itos())\n",
    "            self.embedding = nn.Embedding.from_pretrained(vecs, freeze=False)\n",
    "        else:\n",
    "            self.embedding = nn.Embedding(vocab_size, d_model)\n",
    "        self.d_model = d_model\n",
    "\n",
    "    def forward(self, tokens):\n",
    "        # shape (tokens) --> [sentence len, batch size]\n",
    "        # shape (inp_emb) --> [sentence len, batch size, d_model]\n",
    "        # Multiplying with square root of d_model as they mentioned in the paper\n",
    "        inp_emb = self.embedding(tokens.long()) * math.sqrt(self.d_model)\n",
    "        return inp_emb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "1cba2350-fb0b-4f6b-a41a-591d03cdca9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Seq2SeqTransformer(nn.Module):\n",
    "    def __init__(self, \n",
    "                 src_vocab, \n",
    "                 trg_vocab_size, \n",
    "                 d_model, \n",
    "                 dropout,\n",
    "                 nhead,\n",
    "                 num_encoder_layers,\n",
    "                 num_decoder_layers,\n",
    "                 dim_feedforward,\n",
    "                 src_pad_idx,\n",
    "                 trg_pad_idx\n",
    "                ):\n",
    "        super(Seq2SeqTransformer, self).__init__()\n",
    "        \n",
    "        self.src_pad_idx = src_pad_idx\n",
    "        self.trg_pad_idx = trg_pad_idx\n",
    "        \n",
    "        self.src_inp_emb = InputEmbedding(len(src_vocab), d_model, src_vocab)\n",
    "        self.trg_inp_emb = InputEmbedding(trg_vocab_size, d_model)\n",
    "        \n",
    "        self.positional_encoding = PositionalEncoding(d_model, dropout=dropout)\n",
    "        \n",
    "        self.transformer = Transformer(d_model=d_model,\n",
    "                                       nhead=nhead,\n",
    "                                       num_encoder_layers=num_encoder_layers,\n",
    "                                       num_decoder_layers=num_decoder_layers,\n",
    "                                       dim_feedforward=dim_feedforward,\n",
    "                                       dropout=dropout)\n",
    "        \n",
    "        self.linear = nn.Linear(d_model, trg_vocab_size)\n",
    "    \n",
    "    def forward(self, src, trg):\n",
    "        src_emb = self.positional_encoding(self.src_inp_emb(src))\n",
    "        trg_emb = self.positional_encoding(self.trg_inp_emb(trg))\n",
    "        \n",
    "        src_mask, tgt_mask, src_padding_mask, tgt_padding_mask = self.create_mask(src, trg)\n",
    "        \n",
    "        outs = self.transformer(src = src_emb, \n",
    "                                tgt = trg_emb, \n",
    "                                src_mask = src_mask,\n",
    "                                tgt_mask = tgt_mask, \n",
    "                                src_key_padding_mask = src_padding_mask, \n",
    "                                tgt_key_padding_mask = tgt_padding_mask,\n",
    "                                memory_key_padding_mask = src_padding_mask\n",
    "                               )\n",
    "        return self.linear(outs)\n",
    "        \n",
    "\n",
    "    def create_mask(self, src, trg):\n",
    "        src_seq_len = src.shape[0]\n",
    "        trg_seq_len = trg.shape[0]\n",
    "\n",
    "        # Subsequent mask aka \"look ahead mask\" is important as it wont let Decoder\n",
    "        # to peek into future tokens.\n",
    "        trg_mask = self.transformer.generate_square_subsequent_mask(trg_seq_len).to(device)\n",
    "        \n",
    "        #\n",
    "        src_mask = torch.zeros((src_seq_len, src_seq_len),device=device).type(torch.bool) # All False hence unchanged\n",
    "\n",
    "        # Padding masking will allow attention to ignore padding <pad> tokens \n",
    "        src_padding_mask = (src == self.src_pad_idx).transpose(0, 1)\n",
    "        trg_padding_mask = (trg == self.trg_pad_idx).transpose(0, 1)\n",
    "        \n",
    "        return src_mask, trg_mask, src_padding_mask, trg_padding_mask\n",
    "    \n",
    "    # These two functions will only used while inferring\n",
    "    \n",
    "    def encode(self, src):\n",
    "        src_mask = torch.zeros((src.shape[0], src.shape[0]),device=device).type(torch.bool)\n",
    "        src_padding_mask = (src == self.src_pad_idx).transpose(0, 1)\n",
    "        \n",
    "        return self.transformer.encoder(self.positional_encoding(self.src_inp_emb(src)), src_mask, src_padding_mask)\n",
    "\n",
    "    def decode(self, trg, memory):\n",
    "        # memory is the output from the encoder \n",
    "        trg_seq_len = trg.shape[0]\n",
    "        trg_mask = self.transformer.generate_square_subsequent_mask(trg_seq_len).type(torch.bool).to(device)\n",
    "        trg_padding_mask = (trg == self.trg_pad_idx).transpose(0, 1)\n",
    "        \n",
    "        return self.transformer.decoder(tgt = self.positional_encoding(self.trg_inp_emb(trg)), \n",
    "                                        memory = memory,\n",
    "                                        tgt_mask = trg_mask,\n",
    "                                        tgt_key_padding_mask = trg_padding_mask\n",
    "                                       )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "2993bac6-a5c2-47fe-b190-4f39953fe817",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model(model, train_dataloader, criterion, optimizer):\n",
    "    model.train()\n",
    "    epoch_loss = 0\n",
    "    for batch_idx, batch in enumerate(tqdm(train_dataloader)):\n",
    "        # shape (src, trg) --> [seq len, batch size]\n",
    "        src = batch[\"src\"]\n",
    "        trg = batch[\"trg\"]\n",
    "\n",
    "        # Clear the accumulating gradients\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # shape (trg_inp, trg_out) --> [seq len - 1, batch size]\n",
    "        trg_inp = trg[:-1, :]\n",
    "        trg_out = trg[1:, :]\n",
    "\n",
    "        # shape --> (seq len - 1) * batch size \n",
    "        # Making all target seqeunces in 1d tensor\n",
    "        trg_out = trg_out.reshape(-1)\n",
    "\n",
    "        # shape (logits) --> [seq len - 1, batch size, trg vocab size]\n",
    "        logits = model(src, trg_inp)\n",
    "\n",
    "        # shape (logits) --> [(seq len - 1) * batch size, trg vocab size]\n",
    "        logits = logits.reshape(-1, logits.shape[-1])\n",
    "\n",
    "        loss = criterion(logits, trg_out)\n",
    "\n",
    "        loss.backward()\n",
    "\n",
    "        optimizer.step()\n",
    "        epoch_loss += loss.detach().cpu()\n",
    "\n",
    "    return epoch_loss/len(train_dataloader)\n",
    "\n",
    "def evaluate_model(model, valid_dataloader, criterion):\n",
    "    model.eval()\n",
    "    epoch_loss = 0\n",
    "    with torch.no_grad():\n",
    "        for batch_idx, batch in enumerate(valid_dataloader):\n",
    "            # shape (src, trg) --> [seq len, batch size]\n",
    "            src = batch[\"src\"]\n",
    "            trg = batch[\"trg\"]\n",
    "\n",
    "            # shape (trg_inp, trg_out) --> [seq len - 1, batch size]\n",
    "            trg_inp = trg[:-1, :]\n",
    "            trg_out = trg[1:, :]\n",
    "\n",
    "            # shape --> (seq len - 1) * batch size \n",
    "            # Making all target seqeunces in 1d tensor\n",
    "            trg_out = trg_out.reshape(-1)\n",
    "\n",
    "            # shape (logits) --> [seq len - 1, batch size, trg vocab size]\n",
    "            logits = model(src, trg_inp)\n",
    "\n",
    "            # shape (logits) --> [(seq len - 1) * batch size, trg vocab size]\n",
    "            logits = logits.reshape(-1, logits.shape[-1])\n",
    "\n",
    "            loss = criterion(logits, trg_out)\n",
    "\n",
    "            epoch_loss += loss.detach().cpu()\n",
    "    \n",
    "    return epoch_loss/len(valid_dataloader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "83a78a5d-596b-4672-9b84-75596fd6e725",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      ".vector_cache\\glove.6B.zip: 862MB [02:44, 5.25MB/s]                                                                    \n",
      "100%|██████████████████████████████████████████████████████████████████████▉| 399999/400000 [00:27<00:00, 14736.55it/s]\n"
     ]
    }
   ],
   "source": [
    "model = Seq2SeqTransformer(nmtds_train.src_vocab,\n",
    "                                hyp_params[\"trg_vocab_size\"],\n",
    "                                hyp_params[\"d_model\"],\n",
    "                                hyp_params[\"dropout\"],\n",
    "                                hyp_params[\"n_head\"],\n",
    "                                hyp_params[\"num_encoder_layers\"],\n",
    "                                hyp_params[\"num_decoder_layers\"],\n",
    "                                hyp_params[\"feedforward_dim\"],\n",
    "                                SRC_PAD_IDX,\n",
    "                                TRG_PAD_IDX\n",
    "                                ).to(device)\n",
    "\n",
    "# They did not mention it in paper however tutorials from \n",
    "# bentrevett and others uses it so -_(-.-)_-\n",
    "for p in model.parameters():\n",
    "    if p.dim() > 1:\n",
    "        nn.init.xavier_uniform_(p)\n",
    "\n",
    "criterion = nn.CrossEntropyLoss(ignore_index=TRG_PAD_IDX).to(device)\n",
    "\n",
    "# They did not used fixed learning rate in the paper infact \n",
    "# their optimizer would look like \n",
    "#   optimizer = torch.optim.Adam(transformer.parameters(), lr=lr, betas=(0.9, 0.98), eps=1e-9)\n",
    "# with variable learning rate as they presented in the paper\n",
    "# but for the sake of simplicity and also fixed lr works fine\n",
    "optimizer = optim.Adam(model.parameters(), lr=hyp_params[\"lr\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "6dd1a932-6579-4272-9a23-dceae964232c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 227/227 [00:13<00:00, 16.96it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 1, Train loss: 4.450022220611572, Eval loss: 3.107335090637207, patience: 1. Time 13.805654764175415\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 227/227 [00:12<00:00, 18.42it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 2, Train loss: 2.7562944889068604, Eval loss: 2.196779251098633, patience: 1. Time 12.599369525909424\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 227/227 [00:12<00:00, 17.96it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 3, Train loss: 2.0567777156829834, Eval loss: 1.8117517232894897, patience: 1. Time 12.920658111572266\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 227/227 [00:12<00:00, 18.28it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 4, Train loss: 1.6460583209991455, Eval loss: 1.6064262390136719, patience: 1. Time 12.70738935470581\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 227/227 [00:12<00:00, 18.21it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 5, Train loss: 1.3749070167541504, Eval loss: 1.4789458513259888, patience: 1. Time 12.750669002532959\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 227/227 [00:12<00:00, 18.30it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 6, Train loss: 1.1709470748901367, Eval loss: 1.4230008125305176, patience: 1. Time 12.690899848937988\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 227/227 [00:12<00:00, 18.29it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 7, Train loss: 1.01650071144104, Eval loss: 1.4143142700195312, patience: 1. Time 12.69911813735962\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 227/227 [00:12<00:00, 18.26it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 8, Train loss: 0.8906280398368835, Eval loss: 1.3949575424194336, patience: 1. Time 12.724719047546387\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 227/227 [00:12<00:00, 18.34it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 9, Train loss: 0.7918173670768738, Eval loss: 1.4209078550338745, patience: 1. Time 12.663723468780518\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████████| 227/227 [00:12<00:00, 18.38it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 10, Train loss: 0.7089118957519531, Eval loss: 1.419948697090149, patience: 2. Time 12.627846717834473\n",
      "Best epoch was 8 with 1.3949575424194336 eval loss\n"
     ]
    }
   ],
   "source": [
    "min_el = math.inf\n",
    "patience = 1\n",
    "best_model = {}\n",
    "best_epoch = 0\n",
    "\n",
    "epoch_loss = 0\n",
    "for epoch in range(hyp_params[\"num_epochs\"]):\n",
    "    start = time.time()\n",
    "    \n",
    "    gc.collect()\n",
    "    torch.cuda.empty_cache()\n",
    "    \n",
    "    epoch_loss = train_model(model, train_dataloader, criterion, optimizer)\n",
    "    eval_loss = evaluate_model(model, valid_dataloader, criterion)\n",
    "    \n",
    "    log.log(f\"Epoch: {epoch+1}, Train loss: {epoch_loss}, Eval loss: {eval_loss}, patience: {patience}. Time {time.time() - start}\")\n",
    "\n",
    "    \n",
    "    if eval_loss < min_el:\n",
    "        best_epoch = epoch+1\n",
    "        min_el = eval_loss\n",
    "        best_model = copy.deepcopy(model)\n",
    "        torch.save({\n",
    "            'model_state_dict': model.state_dict(),\n",
    "            'optimizer_state_dict': optimizer.state_dict(),\n",
    "            'eval_loss': min_el\n",
    "        }, 'model-transformer-glove.pt')\n",
    "        patience = 1\n",
    "    else:\n",
    "        patience += 1\n",
    "    \n",
    "    if patience == 10:\n",
    "        log.log(\"[STOPPING] Early stopping in action..\")\n",
    "        log.log(f\"Best epoch was {best_epoch} with {min_el} eval loss\")\n",
    "        break\n",
    "        \n",
    "log.log(f\"Best epoch was {best_epoch} with {min_el} eval loss\")\n",
    "log.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "0c3bac57-24ba-46bc-9075-f8017e195243",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Seq2SeqTransformer(\n",
       "  (src_inp_emb): InputEmbedding(\n",
       "    (embedding): Embedding(5893, 300)\n",
       "  )\n",
       "  (trg_inp_emb): InputEmbedding(\n",
       "    (embedding): Embedding(7853, 300)\n",
       "  )\n",
       "  (positional_encoding): PositionalEncoding(\n",
       "    (dropout): Dropout(p=0.1, inplace=False)\n",
       "  )\n",
       "  (transformer): Transformer(\n",
       "    (encoder): TransformerEncoder(\n",
       "      (layers): ModuleList(\n",
       "        (0): TransformerEncoderLayer(\n",
       "          (self_attn): MultiheadAttention(\n",
       "            (out_proj): NonDynamicallyQuantizableLinear(in_features=300, out_features=300, bias=True)\n",
       "          )\n",
       "          (linear1): Linear(in_features=300, out_features=128, bias=True)\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "          (linear2): Linear(in_features=128, out_features=300, bias=True)\n",
       "          (norm1): LayerNorm((300,), eps=1e-05, elementwise_affine=True)\n",
       "          (norm2): LayerNorm((300,), eps=1e-05, elementwise_affine=True)\n",
       "          (dropout1): Dropout(p=0.1, inplace=False)\n",
       "          (dropout2): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "        (1): TransformerEncoderLayer(\n",
       "          (self_attn): MultiheadAttention(\n",
       "            (out_proj): NonDynamicallyQuantizableLinear(in_features=300, out_features=300, bias=True)\n",
       "          )\n",
       "          (linear1): Linear(in_features=300, out_features=128, bias=True)\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "          (linear2): Linear(in_features=128, out_features=300, bias=True)\n",
       "          (norm1): LayerNorm((300,), eps=1e-05, elementwise_affine=True)\n",
       "          (norm2): LayerNorm((300,), eps=1e-05, elementwise_affine=True)\n",
       "          (dropout1): Dropout(p=0.1, inplace=False)\n",
       "          (dropout2): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "        (2): TransformerEncoderLayer(\n",
       "          (self_attn): MultiheadAttention(\n",
       "            (out_proj): NonDynamicallyQuantizableLinear(in_features=300, out_features=300, bias=True)\n",
       "          )\n",
       "          (linear1): Linear(in_features=300, out_features=128, bias=True)\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "          (linear2): Linear(in_features=128, out_features=300, bias=True)\n",
       "          (norm1): LayerNorm((300,), eps=1e-05, elementwise_affine=True)\n",
       "          (norm2): LayerNorm((300,), eps=1e-05, elementwise_affine=True)\n",
       "          (dropout1): Dropout(p=0.1, inplace=False)\n",
       "          (dropout2): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (norm): LayerNorm((300,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (decoder): TransformerDecoder(\n",
       "      (layers): ModuleList(\n",
       "        (0): TransformerDecoderLayer(\n",
       "          (self_attn): MultiheadAttention(\n",
       "            (out_proj): NonDynamicallyQuantizableLinear(in_features=300, out_features=300, bias=True)\n",
       "          )\n",
       "          (multihead_attn): MultiheadAttention(\n",
       "            (out_proj): NonDynamicallyQuantizableLinear(in_features=300, out_features=300, bias=True)\n",
       "          )\n",
       "          (linear1): Linear(in_features=300, out_features=128, bias=True)\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "          (linear2): Linear(in_features=128, out_features=300, bias=True)\n",
       "          (norm1): LayerNorm((300,), eps=1e-05, elementwise_affine=True)\n",
       "          (norm2): LayerNorm((300,), eps=1e-05, elementwise_affine=True)\n",
       "          (norm3): LayerNorm((300,), eps=1e-05, elementwise_affine=True)\n",
       "          (dropout1): Dropout(p=0.1, inplace=False)\n",
       "          (dropout2): Dropout(p=0.1, inplace=False)\n",
       "          (dropout3): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "        (1): TransformerDecoderLayer(\n",
       "          (self_attn): MultiheadAttention(\n",
       "            (out_proj): NonDynamicallyQuantizableLinear(in_features=300, out_features=300, bias=True)\n",
       "          )\n",
       "          (multihead_attn): MultiheadAttention(\n",
       "            (out_proj): NonDynamicallyQuantizableLinear(in_features=300, out_features=300, bias=True)\n",
       "          )\n",
       "          (linear1): Linear(in_features=300, out_features=128, bias=True)\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "          (linear2): Linear(in_features=128, out_features=300, bias=True)\n",
       "          (norm1): LayerNorm((300,), eps=1e-05, elementwise_affine=True)\n",
       "          (norm2): LayerNorm((300,), eps=1e-05, elementwise_affine=True)\n",
       "          (norm3): LayerNorm((300,), eps=1e-05, elementwise_affine=True)\n",
       "          (dropout1): Dropout(p=0.1, inplace=False)\n",
       "          (dropout2): Dropout(p=0.1, inplace=False)\n",
       "          (dropout3): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "        (2): TransformerDecoderLayer(\n",
       "          (self_attn): MultiheadAttention(\n",
       "            (out_proj): NonDynamicallyQuantizableLinear(in_features=300, out_features=300, bias=True)\n",
       "          )\n",
       "          (multihead_attn): MultiheadAttention(\n",
       "            (out_proj): NonDynamicallyQuantizableLinear(in_features=300, out_features=300, bias=True)\n",
       "          )\n",
       "          (linear1): Linear(in_features=300, out_features=128, bias=True)\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "          (linear2): Linear(in_features=128, out_features=300, bias=True)\n",
       "          (norm1): LayerNorm((300,), eps=1e-05, elementwise_affine=True)\n",
       "          (norm2): LayerNorm((300,), eps=1e-05, elementwise_affine=True)\n",
       "          (norm3): LayerNorm((300,), eps=1e-05, elementwise_affine=True)\n",
       "          (dropout1): Dropout(p=0.1, inplace=False)\n",
       "          (dropout2): Dropout(p=0.1, inplace=False)\n",
       "          (dropout3): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (norm): LayerNorm((300,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "  )\n",
       "  (linear): Linear(in_features=300, out_features=7853, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_l = Seq2SeqTransformer(nmtds_train.src_vocab,\n",
    "                                hyp_params[\"trg_vocab_size\"],\n",
    "                                hyp_params[\"d_model\"],\n",
    "                                hyp_params[\"dropout\"],\n",
    "                                hyp_params[\"n_head\"],\n",
    "                                hyp_params[\"num_encoder_layers\"],\n",
    "                                hyp_params[\"num_decoder_layers\"],\n",
    "                                hyp_params[\"feedforward_dim\"],\n",
    "                                SRC_PAD_IDX,\n",
    "                                TRG_PAD_IDX\n",
    "                                ).to(device)\n",
    "\n",
    "model_l.load_state_dict(torch.load('model-transformer-glove.pt', map_location=device)[\"model_state_dict\"])\n",
    "model_l.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "6cf7be70-bda6-45bc-a179-1cb9f40f66d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def translate(snt, dataset, model, device):\n",
    "    snt = torch.tensor(snt).view(-1,1).to(device)\n",
    "    \n",
    "    num_tokens = snt.shape[0]\n",
    "    max_len = 50\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        memory = model.encode(snt).to(device)\n",
    "    \n",
    "    ys = torch.LongTensor([dataset.trg_vocab['<sos>']]).unsqueeze(0).to(device)\n",
    "\n",
    "    for i in range(max_len):\n",
    "        with torch.no_grad():\n",
    "            out = model.decode(ys, memory)\n",
    "            \n",
    "        out = out.transpose(0, 1)\n",
    "        prob = model.linear(out[:, -1])\n",
    "        next_word = prob.argmax().detach().item()\n",
    "\n",
    "        ys = torch.cat([ys, torch.tensor([next_word]).unsqueeze(1).to(device)])\n",
    "        \n",
    "        \n",
    "        if next_word == dataset.trg_vocab['<eos>']:\n",
    "            break\n",
    "    \n",
    "    return dataset.trg_vocab.lookup_tokens(ys.squeeze().cpu().numpy())\n",
    "\n",
    "\n",
    "def bleu(model, dataset, device):\n",
    "    targets = []\n",
    "    outputs = []\n",
    "\n",
    "    for example in tqdm(dataset):\n",
    "        src = example[\"src\"]\n",
    "        trg = example[\"trg\"]\n",
    "        \n",
    "        trg = dataset.trg_vocab.lookup_tokens(trg)    \n",
    "        prediction = translate(src, dataset, model, device)\n",
    "        \n",
    "        prediction = prediction[1:-1]  # removing <sos> and <eos> tokens\n",
    "        trg = trg[1:-1]\n",
    "        \n",
    "        targets.append([trg])\n",
    "        outputs.append(prediction)\n",
    "\n",
    "    return bleu_score(outputs, targets)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "ff28e153-7e78-4864-88fd-759fe22a3fb3",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████████████████████████████████████████████████████████████████████████| 1000/1000 [00:37<00:00, 26.93it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "0.3799642933115588"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bleu(model_l, nmtds_test, device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4f68e85-5d3d-4706-8039-e0989622360e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01a6898a-982d-49da-bd7c-dd947e93fa54",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}