NaturalQuestionQA_with_Seq2Seq_with_attention.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lxMY2HmHn0Ka"
   },
   "source": [
    "Natural Question Dataset:\n",
    "Official Link:\n",
    "This dataset contains really long context/passages and a single question based on the passage. The answer to the questions are subparts of the passage given in the form of <start_token> and <end_token>.\n",
    "What makes this dataset really challenging is the passage and answer length.\n",
    "\n",
    "I chose this dataset because it will be interesting to work with transformers with such long sequnces. When I tried this dataset with a seq2seq I kept getting OOM. The reason being exterme lenghts of contexts and their answers. The context are entire wikipedia pages in HTML format so there length is typically over 2000.\n",
    "\n",
    "I started experimenting with original dataset. I kept watering down the task. First reducing the number of examples, then answer length, batch_size but I kept getting OOM. Finally I restricted context lengths to 1000.\n",
    "\n",
    "The problem was in the sequence length of context. If I restrict them to 1000, the model has good enough memory to work with. No of samples has not much effect. And batch size does. I will remove this 1000 length restriction with transformers.\n",
    "\n",
    "Note: this notebook is just about experimenting with this dataset. Its not about getting SOTA or admirable results. For example, in this notebook we only consider contexts with lenghts less than 1000.\n",
    "\n",
    "Also, in this notebook we add attention components for answer_decoder. Attention uses only the context encoder hidden states and not question encoder hidden states."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "k7_hoAmMLduI"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import json\n",
    "from torchtext import data\n",
    "from itertools import chain\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import time\n",
    "from torch.nn import Embedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "dgMi39JRLvQY",
    "outputId": "226f474a-eb3b-4d2f-f9c9-4e4e90d4f984"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current\n",
      "                                 Dload  Upload   Total   Spent    Left  Speed\n",
      "\r",
      "  0     0    0     0    0     0      0      0 --:--:-- --:--:-- --:--:--     0\r",
      "  0 4497M    0 32.0M    0     0  34.8M      0  0:02:08 --:--:--  0:02:08 34.8M\r",
      "  2 4497M    2  114M    0     0  60.9M      0  0:01:13  0:00:01  0:01:12 60.9M\r",
      "  4 4497M    4  200M    0     0  63.3M      0  0:01:11  0:00:03  0:01:08 63.3M\r",
      "  6 4497M    6  280M    0     0  71.3M      0  0:01:03  0:00:03  0:01:00 71.2M\r",
      "  8 4497M    8  364M    0     0  74.6M      0  0:01:00  0:00:04  0:00:56 74.6M\r",
      " 10 4497M   10  464M    0     0  77.5M      0  0:00:57  0:00:05  0:00:52 85.2M\r",
      " 11 4497M   11  512M    0     0  73.8M      0  0:01:00  0:00:06  0:00:54 78.6M\r",
      " 14 4497M   14  631M    0     0  80.0M      0  0:00:56  0:00:07  0:00:49 91.2M\r",
      " 15 4497M   15  719M    0     0  80.9M      0  0:00:55  0:00:08  0:00:47 88.5M\r",
      " 17 4497M   17  800M    0     0  78.1M      0  0:00:57  0:00:10  0:00:47 81.3M\r",
      " 19 4497M   19  856M    0     0  78.4M      0  0:00:57  0:00:10  0:00:47 79.4M\r",
      " 20 4497M   20  936M    0     0  76.6M      0  0:00:58  0:00:12  0:00:46 80.2M\r",
      " 22 4497M   22 1000M    0     0  76.3M      0  0:00:58  0:00:13  0:00:45 70.7M\r",
      " 23 4497M   23 1064M    0     0  76.0M      0  0:00:59  0:00:13  0:00:46 67.6M\r",
      " 25 4497M   25 1143M    0     0  76.7M      0  0:00:58  0:00:14  0:00:44 73.8M\r",
      " 27 4497M   27 1228M    0     0  77.3M      0  0:00:58  0:00:15  0:00:43 74.9M\r",
      " 30 4497M   30 1352M    0     0  77.9M      0  0:00:57  0:00:17  0:00:40 81.2M\r",
      " 31 4497M   31 1417M    0     0  79.2M      0  0:00:56  0:00:17  0:00:39 87.2M\r",
      " 34 4497M   34 1536M    0     0  80.6M      0  0:00:55  0:00:19  0:00:36 93.5M\r",
      " 36 4497M   36 1632M    0     0  81.8M      0  0:00:54  0:00:19  0:00:35 96.8M\r",
      " 38 4497M   38 1742M    0     0  83.4M      0  0:00:53  0:00:20  0:00:33  102M\r",
      " 40 4497M   40 1824M    0     0  83.2M      0  0:00:54  0:00:21  0:00:33  103M\r",
      " 42 4497M   42 1928M    0     0  83.8M      0  0:00:53  0:00:23  0:00:30 99.8M\r",
      " 44 4497M   44 2016M    0     0  84.1M      0  0:00:53  0:00:23  0:00:30 97.3M\r",
      " 46 4497M   46 2112M    0     0  84.7M      0  0:00:53  0:00:24  0:00:29 95.9M\r",
      " 48 4497M   48 2176M    0     0  82.1M      0  0:00:54  0:00:26  0:00:28 77.6M\r",
      " 49 4497M   49 2208M    0     0  82.0M      0  0:00:54  0:00:26  0:00:28 76.8M\r",
      " 51 4497M   51 2307M    0     0  82.6M      0  0:00:54  0:00:27  0:00:27 77.2M\r",
      " 52 4497M   52 2342M    0     0  81.0M      0  0:00:55  0:00:28  0:00:27 66.1M\r",
      " 52 4497M   52 2378M    0     0  79.5M      0  0:00:56  0:00:29  0:00:27 53.5M\r",
      " 53 4497M   53 2411M    0     0  78.0M      0  0:00:57  0:00:30  0:00:27 53.0M\r",
      " 54 4497M   54 2439M    0     0  76.4M      0  0:00:58  0:00:31  0:00:27 46.3M\r",
      " 54 4497M   54 2472M    0     0  74.6M      0  0:01:00  0:00:33  0:00:27 31.5M\r",
      " 55 4497M   55 2511M    0     0  74.1M      0  0:01:00  0:00:33  0:00:27 33.7M\r",
      " 56 4497M   56 2539M    0     0  72.7M      0  0:01:01  0:00:34  0:00:27 32.2M\r",
      " 57 4497M   57 2568M    0     0  71.5M      0  0:01:02  0:00:35  0:00:27 31.5M\r",
      " 57 4497M   57 2603M    0     0  70.5M      0  0:01:03  0:00:36  0:00:27 32.7M\r",
      " 58 4497M   58 2652M    0     0  69.9M      0  0:01:04  0:00:37  0:00:27 37.6M\r",
      " 59 4497M   59 2674M    0     0  68.7M      0  0:01:05  0:00:38  0:00:27 32.6M\r",
      " 60 4497M   60 2720M    0     0  68.1M      0  0:01:06  0:00:39  0:00:27 36.0M\r",
      " 61 4497M   61 2776M    0     0  67.4M      0  0:01:06  0:00:41  0:00:25 39.6M\r",
      " 62 4497M   62 2817M    0     0  67.2M      0  0:01:06  0:00:41  0:00:25 42.7M\r",
      " 63 4497M   63 2848M    0     0  66.3M      0  0:01:07  0:00:42  0:00:25 39.0M\r",
      " 63 4497M   63 2863M    0     0  65.1M      0  0:01:09  0:00:43  0:00:26 37.2M\r",
      " 64 4497M   64 2892M    0     0  64.4M      0  0:01:09  0:00:44  0:00:25 34.6M\r",
      " 64 4497M   64 2915M    0     0  63.5M      0  0:01:10  0:00:45  0:00:25 29.3M\r",
      " 65 4497M   65 2944M    0     0  62.6M      0  0:01:11  0:00:47  0:00:24 24.8M\r",
      " 66 4497M   66 2968M    0     0  61.6M      0  0:01:12  0:00:48  0:00:24 23.1M\r",
      " 66 4497M   66 3000M    0     0  61.3M      0  0:01:13  0:00:48  0:00:25 27.7M\r",
      " 67 4497M   67 3056M    0     0  60.5M      0  0:01:14  0:00:50  0:00:24 29.5M\r",
      " 68 4497M   68 3072M    0     0  60.1M      0  0:01:14  0:00:51  0:00:23 30.2M\r",
      " 69 4497M   69 3112M    0     0  59.9M      0  0:01:14  0:00:51  0:00:23 34.3M\r",
      " 70 4497M   70 3154M    0     0  59.6M      0  0:01:15  0:00:52  0:00:23 38.7M\r",
      " 71 4497M   71 3208M    0     0  59.4M      0  0:01:15  0:00:53  0:00:22 41.0M\r",
      " 72 4497M   72 3277M    0     0  59.6M      0  0:01:15  0:00:54  0:00:21 49.6M\r",
      " 73 4497M   73 3308M    0     0  59.1M      0  0:01:16  0:00:55  0:00:21 48.7M\r",
      " 74 4497M   74 3328M    0     0  58.4M      0  0:01:16  0:00:56  0:00:20 43.0M\r",
      " 74 4497M   74 3351M    0     0  57.9M      0  0:01:17  0:00:57  0:00:20 39.7M\r",
      " 74 4497M   74 3368M    0     0  57.1M      0  0:01:18  0:00:58  0:00:20 32.1M\r",
      " 75 4497M   75 3400M    0     0  56.4M      0  0:01:19  0:01:00  0:00:19 22.7M\r",
      " 75 4497M   75 3416M    0     0  56.0M      0  0:01:20  0:01:00  0:00:20 21.3M\r",
      " 77 4497M   77 3464M    0     0  55.4M      0  0:01:21  0:01:02  0:00:19 24.3M\r",
      " 77 4497M   77 3480M    0     0  55.2M      0  0:01:21  0:01:02  0:00:19 25.3M\r",
      " 78 4497M   78 3536M    0     0  54.5M      0  0:01:22  0:01:04  0:00:18 28.7M\r",
      " 78 4497M   78 3536M    0     0  54.4M      0  0:01:22  0:01:04  0:00:18 29.3M\r",
      " 80 4497M   80 3624M    0     0  54.8M      0  0:01:22  0:01:06  0:00:16 40.7M\r",
      " 81 4497M   81 3672M    0     0  54.4M      0  0:01:22  0:01:07  0:00:15 41.5M\r",
      " 83 4497M   83 3736M    0     0  54.9M      0  0:01:21  0:01:08  0:00:13 50.3M\r",
      " 84 4497M   84 3817M    0     0  55.4M      0  0:01:21  0:01:08  0:00:13 68.4M\r",
      " 85 4497M   85 3850M    0     0  55.0M      0  0:01:21  0:01:09  0:00:12 63.0M\r",
      " 86 4497M   86 3880M    0     0  54.7M      0  0:01:22  0:01:10  0:00:12 53.4M\r",
      " 86 4497M   86 3904M    0     0  54.2M      0  0:01:22  0:01:11  0:00:11 52.5M\r",
      " 87 4497M   87 3922M    0     0  53.8M      0  0:01:23  0:01:12  0:00:11 38.4M\r",
      " 88 4497M   88 3960M    0     0  53.4M      0  0:01:24  0:01:14  0:00:10 27.5M\r",
      " 88 4497M   88 3991M    0     0  53.2M      0  0:01:24  0:01:14  0:00:10 28.3M\r",
      " 89 4497M   89 4008M    0     0  52.7M      0  0:01:25  0:01:15  0:00:10 25.1M\r",
      " 90 4497M   90 4056M    0     0  52.7M      0  0:01:25  0:01:16  0:00:09 30.5M\r",
      " 91 4497M   91 4113M    0     0  52.8M      0  0:01:25  0:01:17  0:00:08 38.0M\r",
      " 92 4497M   92 4160M    0     0  52.7M      0  0:01:25  0:01:18  0:00:07 41.1M\r",
      " 93 4497M   93 4208M    0     0  52.4M      0  0:01:25  0:01:20  0:00:05 40.8M\r",
      " 94 4497M   94 4256M    0     0  52.5M      0  0:01:25  0:01:20  0:00:05 50.0M\r",
      " 95 4497M   95 4295M    0     0  52.4M      0  0:01:25  0:01:21  0:00:04 47.7M\r",
      " 95 4497M   95 4314M    0     0  52.0M      0  0:01:26  0:01:22  0:00:04 39.8M\r",
      " 96 4497M   96 4344M    0     0  51.6M      0  0:01:27  0:01:24  0:00:03 35.2M\r",
      " 97 4497M   97 4368M    0     0  51.4M      0  0:01:27  0:01:24  0:00:03 33.4M\r",
      " 97 4497M   97 4405M    0     0  51.2M      0  0:01:27  0:01:25  0:00:02 30.1M\r",
      " 99 4497M   99 4458M    0     0  51.3M      0  0:01:27  0:01:26  0:00:01 32.7M\r",
      " 99 4497M   99 4494M    0     0  51.1M      0  0:01:27  0:01:27 --:--:-- 36.3M\r",
      "100 4497M  100 4497M    0     0  51.1M      0  0:01:28  0:01:28 --:--:-- 39.8M\n"
     ]
    }
   ],
   "source": [
    "%%bash\n",
    "file=/content/train.jsonl\n",
    "if [ ! -f \"$file\" ]; then\n",
    "  # Since the dataset itself is huge, we will make train and test set from the original train file itself\n",
    "  # A simple wget wont be able to get the natural questions dataset present at https://ai.google.com/research/NaturalQuestions/download\n",
    "  # Use the advice given in https://www.kaggle.com/c/deepfake-detection-challenge/discussion/121194 \n",
    "  # the download link is https://storage.cloud.google.com/natural_questions/v1.0-simplified/simplified-nq-train.jsonl.gz \n",
    "  \n",
    "  # put curl command here\n",
    "  # command -o /content/train.jsonl.gz\n",
    "  \n",
    "  zcat /content/train.jsonl.gz > /content/train.jsonl\n",
    "  \n",
    "fi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "fe9fVmeXsoS2"
   },
   "outputs": [],
   "source": [
    "# In the beginning restrciting total number of examples to 10k. \n",
    "total_limit = 1000\n",
    "context_length = 1000\n",
    "span_length = 200\n",
    "\n",
    "file = \"/content/train.jsonl\"\n",
    "f = open(file)\n",
    "\n",
    "total = 0\n",
    "examples = []\n",
    "for i,line in enumerate(f):\n",
    "    if total >total_limit:\n",
    "        break\n",
    "    ak = json.loads(line)\n",
    "    context = ak['document_text']\n",
    "    question = ak['question_text']\n",
    "    start = ak['annotations'][0]['long_answer']['start_token']\n",
    "    end = ak['annotations'][0]['long_answer']['end_token']\n",
    "    try:\n",
    "        assert start < end and end-start > span_length and len(context.split(\" \")) <context_length\n",
    "    except AssertionError:\n",
    "        continue\n",
    "    answer = \" \".join(context.split(\" \")[start:end])\n",
    "    examples.append([context,question,answer])\n",
    "    total += 1\n",
    "f.close()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "q-sy2SeXA4v3"
   },
   "outputs": [],
   "source": [
    "f = open(\"/content/pruned_examples\",\"w\")\n",
    "for i in examples:\n",
    "  f.write(str(i))\n",
    "  f.write(\"\\n\")\n",
    "\n",
    "f.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "id": "Bax6kpFwA08b"
   },
   "outputs": [],
   "source": [
    "# Will do the architecture with seperate encoders for context and question and a decoder for answer\n",
    "context = data.Field(sequential=True, tokenize='spacy', init_token='<sos>', eos_token='<eos>')\n",
    "question = data.Field(sequential=True, tokenize='spacy', init_token='<sos>', eos_token='<eos>')\n",
    "answer = data.Field(sequential=True, tokenize='spacy', init_token='<sos>', eos_token='<eos>')\n",
    "\n",
    "fields = [('context', context), ('question', question), ('answer', answer)]\n",
    "\n",
    "Examples = [data.Example.fromlist([i[0], i[1], i[2]], fields) for i in examples]\n",
    "Dataset = data.Dataset(Examples, fields)\n",
    "\n",
    "train_dataset,valid_dataset = Dataset.split(split_ratio=[0.85,0.15])\n",
    "\n",
    "context.build_vocab(train_dataset,min_freq=2,max_size = 20000)\n",
    "question.build_vocab(train_dataset,min_freq=2,max_size = 5000)\n",
    "answer.build_vocab(train_dataset,min_freq=2,max_size = 10000)\n",
    "\n",
    "\n",
    "BATCH_SIZE = 32\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "train_iterator, valid_iterator = data.BucketIterator.splits((train_dataset, valid_dataset), batch_size=32,\n",
    "                                                            sort_key=lambda x: len(x.context),\n",
    "                                                            sort_within_batch=True,device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "F6NRhae8tmy_",
    "outputId": "048a81d1-c34e-47ac-c549-c9a1be3477fd"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<torchtext.data.example.Example object at 0x7fd174d055f8>\n",
      "247\n"
     ]
    }
   ],
   "source": [
    "print(Examples[0])\n",
    "print(len(Examples))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Ckb4iFY43Umf"
   },
   "source": [
    "Only 247 examples with context length less than 1000!!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "id": "Os7EnyUltw2_"
   },
   "outputs": [],
   "source": [
    "class Encoder(nn.Module):\n",
    "  def __init__(self,input_dim,emb_dim,hid_dim,n_layers,dropout,bidirectional):\n",
    "    super().__init__()\n",
    "    self.hid_dim = hid_dim\n",
    "    \n",
    "    self.embedding = nn.Embedding(input_dim,emb_dim)\n",
    "\n",
    "    self.rnn = nn.LSTM(input_size=emb_dim,hidden_size = hid_dim,num_layers= n_layers,dropout= dropout,bidirectional = bidirectional)\n",
    "\n",
    "    self.dropout = nn.Dropout(dropout)\n",
    "  \n",
    "  def forward(self,input,hidden=None,cell_state=None):\n",
    "    \n",
    "    embedded = self.dropout(self.embedding(input))\n",
    "\n",
    "    if not hidden == None:\n",
    "      outputs, (hidden,cell_state) = self.rnn(embedded,(hidden,cell_state))\n",
    "    else:\n",
    "      outputs, (hidden,cell_state) = self.rnn(embedded)\n",
    "\n",
    "    return outputs ,hidden,cell_state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "id": "Lhriarwom_Jg"
   },
   "outputs": [],
   "source": [
    "class Attention(nn.Module):\n",
    "  def __init__(self, attn_vector_size,hid_dim,bidirectional,n_layers, type_='dot'):\n",
    "    super().__init__()\n",
    "    self.attn_vector_size = attn_vector_size\n",
    "    self.bidirectional = bidirectional\n",
    "    self.n_layers = n_layers\n",
    "    self.type_ = type_\n",
    "    self.hid_dim = hid_dim\n",
    "    self.directions = 2 if bidirectional else 1\n",
    "    # The final vector is concatenation of (n_layers*directions) different vectors and the final size is 'attn_vector_size'\n",
    "    # so it should be divisible with (n_layers*directions) \n",
    "    assert attn_vector_size%(n_layers*self.directions) == 0 \n",
    "\n",
    "    self.transform_ = nn.Linear(self.directions*self.hid_dim,int(attn_vector_size/(n_layers*self.directions)))\n",
    "  \n",
    "  def forward(self,encoder_outputs,decoder_hidden_state):\n",
    "    # Interesting: what happens when there are multiple layers and bidirectionality.. Do I find attention vector for each layer and append it to \n",
    "    # input of each layer.. but we give LSTMs input only once..so I think I need a attention vector including information from all layers\n",
    "    # given as a single input once.\n",
    "\n",
    "    # Since I am already working with very long sequences..to avoid too many params, I am using dot product as a score function\n",
    "    if not self.type_ == 'dot': raise NotImplementedError\n",
    "    \n",
    "    # The idea is to find attention for each layer and direction of decoder hidden state and later combine all these and give a final vector\n",
    "    # so the attention input has attention information from all decoder hidden states\n",
    "\n",
    "    # encoder_outputs : [src_len,batch_size,num_of_directions*hid_dim]\n",
    "    # decoder_hidden_state : [n_layers*no_of_directions,batch_size,hid_dim]\n",
    "\n",
    "    with torch.no_grad():\n",
    "      decoder_hidden_state_len = len(decoder_hidden_state)\n",
    "      src_len = encoder_outputs.shape[0]\n",
    "\n",
    "      encoder_view = encoder_outputs.view(-1,self.directions,self.hid_dim) # encoder_view : [src_len*batch_size,num_of_directions,hid_dim]\n",
    "\n",
    "      encoder_outputs = encoder_outputs.permute(1,0,2) # encoder_outputs : [batch_size,src_len,num_of_directions*hid_dim]\n",
    "      all_attention_vectors = []\n",
    "\n",
    "    for i in range(decoder_hidden_state_len):\n",
    "      \n",
    "      with torch.no_grad():\n",
    "        # Calculating alpha\n",
    "        hidden_state = decoder_hidden_state[i] # hidden state : [batch_size,hid_dim]\n",
    "\n",
    "        hidden_state = hidden_state.unsqueeze(0) # hidden_state : [1,batch_size,hid_dim]\n",
    "        hidden_state = hidden_state.permute(1,2,0) # hidden_state :[batch_size,hid_dim, 1]\n",
    "        hidden_state = hidden_state.repeat(src_len,1,1) # hidden_state :[src_len*batch_size,hid_dim, 1]\n",
    "\n",
    "        temp = torch.bmm(encoder_view,hidden_state) # temp : [src_len*batch_size,num_of_directions,1]\n",
    "        temp = temp.sum(1) # temp : [src_len*batch_size,1]\n",
    "        temp = temp.reshape(-1,1,src_len) # temp :[batch_size,1,src_len]\n",
    "\n",
    "        alpha = F.softmax(temp,dim=-1)  # alpha :[batch_size,1,src_len]\n",
    "\n",
    "        # Calculating attention vector\n",
    "        c = torch.bmm(alpha,encoder_outputs) # s: [batch_size,1,num_of_directions*hid_dim]\n",
    "\n",
    "      transform_c = self.transform_(c) # transform_s : [batch_size,1,attn_vector_size/(n_layers*self.directions)]\n",
    "      transform_c = transform_c.permute(1,0,2) # transform_s : [1,batch_size,attn_vector_size/(n_layers*self.directions)]\n",
    "      all_attention_vectors.append(transform_c)\n",
    "\n",
    "\n",
    "    # return vector of shape [1,batch_size,self.attn_vector_size(emb_dim)]\n",
    "    return torch.cat(all_attention_vectors,dim = -1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "id": "G2oFrolIuuMF"
   },
   "outputs": [],
   "source": [
    "class Decoder(nn.Module):\n",
    "  def __init__(self,output_dim,emd_dim,hid_dim,n_layers,bidirectional,dropout):\n",
    "    super().__init__()\n",
    "\n",
    "    self.embedded = nn.Embedding(output_dim,emb_dim)\n",
    "\n",
    "    # Input dims will be 2*emb_size because it will also recieve attention information\n",
    "    self.rnn = nn.LSTM(input_size=2*emb_dim,hidden_size=hid_dim,num_layers=n_layers,bidirectional=bidirectional,dropout=dropout)\n",
    "\n",
    "    self.attention = Attention(emd_dim,hid_dim,bidirectional,n_layers)\n",
    "\n",
    "    self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    no_of_directions = 2 if bidirectional else 1\n",
    "\n",
    "    # I am not passing attention vector to help in output prediction\n",
    "    self.fc_out = nn.Linear(no_of_directions*hid_dim,output_dim)\n",
    "\n",
    "  def forward(self,input,encoder_outputs,hidden,cell_state):\n",
    "    input = input.unsqueeze(0)\n",
    "\n",
    "    input = self.dropout(self.embedded(input))\n",
    "\n",
    "    # Attention\n",
    "    attention_vector = self.attention(encoder_outputs,hidden)\n",
    "\n",
    "    # concat input and attention vector\n",
    "    input = torch.cat([input,attention_vector],dim=-1)\n",
    "\n",
    "    output, (hidden,cell_state) = self.rnn(input,(hidden,cell_state))\n",
    "\n",
    "    output = output.squeeze(0)\n",
    "    \n",
    "    prediction = self.fc_out(output)\n",
    "\n",
    "    return prediction, hidden , cell_state\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "id": "pnp6p4uTxJFE"
   },
   "outputs": [],
   "source": [
    "import random\n",
    "class Seq2seq(nn.Module):\n",
    "  def __init__(self,context_dim,question_dim,answer_dim,emd_dim,hid_dim,n_layers,bidirectional,dropout):\n",
    "    super().__init__()\n",
    "\n",
    "    self.context_encoder = Encoder(context_dim,emd_dim,hid_dim,n_layers,dropout,bidirectional)\n",
    "    self.question_encoder = Encoder(question_dim,emd_dim,hid_dim,n_layers,dropout,bidirectional)\n",
    "    self.answer_decoder = Decoder(answer_dim,emd_dim,hid_dim,n_layers,bidirectional,dropout)\n",
    "\n",
    "    self.answer_dim = answer_dim\n",
    "\n",
    "  def forward(self,context,question,answer,teacher_forcing =0.5):\n",
    "\n",
    "    encoder_outputs, hidden,cell_state = self.context_encoder(context)\n",
    "\n",
    "    _,hidden,cell_state = self.question_encoder(question,hidden,cell_state)\n",
    "\n",
    "    answer_len = len(answer)\n",
    "    batch_size = answer.shape[1]\n",
    "\n",
    "    outputs = torch.zeros(answer_len,batch_size,self.answer_dim).to(device)\n",
    "\n",
    "    for i,j in enumerate(range(answer_len)):\n",
    "      k = answer[j]\n",
    "      if i != 0:\n",
    "        k = prediction.argmax(1) if random.random() < teacher_forcing else k\n",
    "      prediction, hidden, cell_state = self.answer_decoder(k,encoder_outputs,hidden,cell_state)\n",
    "      outputs[j] = prediction\n",
    "\n",
    "    return outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "7a64ckaL1r80",
    "outputId": "8b436c72-b3f6-42c2-8e53-4bc2aa2f5318"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/rnn.py:61: UserWarning: dropout option adds dropout after all but last recurrent layer, so non-zero dropout expects num_layers greater than 1, but got dropout=0.5 and num_layers=1\n",
      "  \"num_layers={}\".format(dropout, num_layers))\n"
     ]
    }
   ],
   "source": [
    "context_dim = len(context.vocab)\n",
    "question_dim = len(question.vocab)\n",
    "answer_dim = len(answer.vocab)\n",
    "emb_dim = 100\n",
    "hid_dim = 100\n",
    "n_layers = 1\n",
    "bidirectional = False\n",
    "dropout = 0.5\n",
    "\n",
    "model = Seq2seq(context_dim,question_dim,answer_dim,emb_dim,hid_dim,n_layers,bidirectional,dropout).to(device)\n",
    "\n",
    "def init_weights(m):\n",
    "    for name, param in m.named_parameters():\n",
    "      # if not isinstance(m, Embedding):\n",
    "      nn.init.normal_(param.data, mean=0, std=0.01)\n",
    "        \n",
    "model.apply(init_weights)\n",
    "\n",
    "optimizer = optim.Adam(model.parameters())\n",
    "\n",
    "\n",
    "TRG_PAD_IDX = answer.vocab.stoi[answer.pad_token]\n",
    "\n",
    "criterion = nn.CrossEntropyLoss(ignore_index = TRG_PAD_IDX)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "uM3lPXZDeKFp",
    "outputId": "e8abccad-4cb6-401e-fd66-9345a01a4f34"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The model has 1,870,156 trainable parameters\n"
     ]
    }
   ],
   "source": [
    "def count_parameters(model):\n",
    "    return sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "\n",
    "print(f'The model has {count_parameters(model):,} trainable parameters')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "id": "SqDxwXf_2D0i"
   },
   "outputs": [],
   "source": [
    "def train(model, iterator, optimizer, criterion, clip):\n",
    "    \n",
    "    model.train()\n",
    "    \n",
    "    epoch_loss = 0\n",
    "    \n",
    "    for i, batch in enumerate(iterator):\n",
    "        \n",
    "        context_ = batch.context\n",
    "        question_ = batch.question\n",
    "        answer_ = batch.answer\n",
    "        \n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        output = model(context_, question_,answer_)\n",
    "        \n",
    "        trg = answer_\n",
    "        #trg = [trg len, batch size]\n",
    "        #output = [trg len, batch size, output dim]\n",
    "        \n",
    "        output_dim = output.shape[-1]\n",
    "        \n",
    "        output = output[1:].view(-1, output_dim)\n",
    "        trg = trg[1:].view(-1)\n",
    "        \n",
    "        #trg = [(trg len - 1) * batch size]\n",
    "        #output = [(trg len - 1) * batch size, output dim]\n",
    "        \n",
    "        loss = criterion(output, trg)\n",
    "        \n",
    "        loss.backward()\n",
    "        \n",
    "        torch.nn.utils.clip_grad_norm_(model.parameters(), clip)\n",
    "        \n",
    "        optimizer.step()\n",
    "        \n",
    "        epoch_loss += loss.item()\n",
    "        \n",
    "    return epoch_loss / len(iterator)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "id": "WmRFlHIW2QQ2"
   },
   "outputs": [],
   "source": [
    "def evaluate(model, iterator, criterion):\n",
    "    \n",
    "    model.eval()\n",
    "    \n",
    "    epoch_loss = 0\n",
    "    \n",
    "    with torch.no_grad():\n",
    "    \n",
    "        for i, batch in enumerate(iterator):\n",
    "\n",
    "            context_ = batch.context\n",
    "            question_ = batch.question\n",
    "            answer_ = batch.answer\n",
    "        \n",
    "            output = model(context_, question_,answer_,0) #turn off teacher forcing\n",
    "\n",
    "            trg = answer_\n",
    "            #trg = [trg len, batch size]\n",
    "            #output = [trg len, batch size, output dim]\n",
    "\n",
    "            output_dim = output.shape[-1]\n",
    "            \n",
    "            output = output[1:].view(-1, output_dim)\n",
    "            trg = trg[1:].view(-1)\n",
    "\n",
    "            #trg = [(trg len - 1) * batch size]\n",
    "            #output = [(trg len - 1) * batch size, output dim]\n",
    "\n",
    "            loss = criterion(output, trg)\n",
    "\n",
    "            epoch_loss += loss.item()\n",
    "        \n",
    "    return epoch_loss / len(iterator)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "id": "iHsk7wYu2Tdj"
   },
   "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": "code",
   "execution_count": 26,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "kPvtYPOw2VqN",
    "outputId": "85804c98-8c1c-4f8c-c351-1dfbcd484176"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 01 | Time: 0m 26s\n",
      "\tTrain Loss: 8.348 | Train PPL: 4222.266\n",
      "\t Val. Loss: 8.325 |  Val. PPL: 4124.963\n",
      "Epoch: 02 | Time: 0m 26s\n",
      "\tTrain Loss: 8.281 | Train PPL: 3949.416\n",
      "\t Val. Loss: 8.129 |  Val. PPL: 3391.024\n",
      "Epoch: 03 | Time: 0m 26s\n",
      "\tTrain Loss: 7.323 | Train PPL: 1514.181\n",
      "\t Val. Loss: 6.159 |  Val. PPL: 473.176\n",
      "Epoch: 04 | Time: 0m 26s\n",
      "\tTrain Loss: 5.937 | Train PPL: 378.985\n",
      "\t Val. Loss: 4.935 |  Val. PPL: 139.111\n",
      "Epoch: 05 | Time: 0m 26s\n",
      "\tTrain Loss: 5.099 | Train PPL: 163.874\n",
      "\t Val. Loss: 4.047 |  Val. PPL:  57.237\n",
      "Epoch: 06 | Time: 0m 26s\n",
      "\tTrain Loss: 4.550 | Train PPL:  94.597\n",
      "\t Val. Loss: 3.593 |  Val. PPL:  36.334\n",
      "Epoch: 07 | Time: 0m 26s\n",
      "\tTrain Loss: 4.407 | Train PPL:  82.062\n",
      "\t Val. Loss: 3.447 |  Val. PPL:  31.419\n",
      "Epoch: 08 | Time: 0m 26s\n",
      "\tTrain Loss: 4.422 | Train PPL:  83.293\n",
      "\t Val. Loss: 3.418 |  Val. PPL:  30.512\n",
      "Epoch: 09 | Time: 0m 27s\n",
      "\tTrain Loss: 4.410 | Train PPL:  82.303\n",
      "\t Val. Loss: 3.406 |  Val. PPL:  30.131\n",
      "Epoch: 10 | Time: 0m 26s\n",
      "\tTrain Loss: 4.404 | Train PPL:  81.774\n",
      "\t Val. Loss: 3.418 |  Val. PPL:  30.511\n",
      "Epoch: 11 | Time: 0m 26s\n",
      "\tTrain Loss: 4.388 | Train PPL:  80.492\n",
      "\t Val. Loss: 3.441 |  Val. PPL:  31.208\n",
      "Epoch: 12 | Time: 0m 26s\n",
      "\tTrain Loss: 4.383 | Train PPL:  80.071\n",
      "\t Val. Loss: 3.451 |  Val. PPL:  31.547\n",
      "Epoch: 13 | Time: 0m 26s\n",
      "\tTrain Loss: 4.391 | Train PPL:  80.691\n",
      "\t Val. Loss: 3.482 |  Val. PPL:  32.523\n",
      "Epoch: 14 | Time: 0m 26s\n",
      "\tTrain Loss: 4.382 | Train PPL:  79.961\n",
      "\t Val. Loss: 3.477 |  Val. PPL:  32.374\n",
      "Epoch: 15 | Time: 0m 26s\n",
      "\tTrain Loss: 4.385 | Train PPL:  80.270\n",
      "\t Val. Loss: 3.449 |  Val. PPL:  31.476\n",
      "Epoch: 16 | Time: 0m 26s\n",
      "\tTrain Loss: 4.382 | Train PPL:  79.958\n",
      "\t Val. Loss: 3.443 |  Val. PPL:  31.288\n",
      "Epoch: 17 | Time: 0m 26s\n",
      "\tTrain Loss: 4.389 | Train PPL:  80.524\n",
      "\t Val. Loss: 3.465 |  Val. PPL:  31.971\n",
      "Epoch: 18 | Time: 0m 26s\n",
      "\tTrain Loss: 4.381 | Train PPL:  79.897\n",
      "\t Val. Loss: 3.465 |  Val. PPL:  31.963\n",
      "Epoch: 19 | Time: 0m 26s\n",
      "\tTrain Loss: 4.382 | Train PPL:  79.959\n",
      "\t Val. Loss: 3.443 |  Val. PPL:  31.268\n",
      "Epoch: 20 | Time: 0m 26s\n",
      "\tTrain Loss: 4.381 | Train PPL:  79.915\n",
      "\t Val. Loss: 3.450 |  Val. PPL:  31.508\n"
     ]
    }
   ],
   "source": [
    "import math\n",
    "import torch.nn.functional as F\n",
    "\n",
    "N_EPOCHS = 20\n",
    "CLIP = 1\n",
    "\n",
    "best_valid_loss = float('inf')\n",
    "\n",
    "for epoch in range(N_EPOCHS):\n",
    "    \n",
    "    start_time = time.time()\n",
    "    \n",
    "    train_loss = train(model, train_iterator, optimizer, criterion, CLIP)\n",
    "    valid_loss = evaluate(model, valid_iterator, criterion)\n",
    "    \n",
    "    end_time = time.time()\n",
    "    \n",
    "    epoch_mins, epoch_secs = epoch_time(start_time, end_time)\n",
    "    \n",
    "    if valid_loss < best_valid_loss:\n",
    "        best_valid_loss = valid_loss\n",
    "        torch.save(model.state_dict(), 'tut2-model.pt')\n",
    "    \n",
    "    print(f'Epoch: {epoch+1:02} | Time: {epoch_mins}m {epoch_secs}s')\n",
    "    print(f'\\tTrain Loss: {train_loss:.3f} | Train PPL: {math.exp(train_loss):7.3f}')\n",
    "    print(f'\\t Val. Loss: {valid_loss:.3f} |  Val. PPL: {math.exp(valid_loss):7.3f}')"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "name": "NaturalQuestionQA with Seq2Seq with attention.ipynb",
   "provenance": []
  },
  "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.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}