1. DrQA.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## DrQA \n",
    "\n",
    "This notebook implements model proposed in the paper: [Reading Wikipedia to Answer Open-Domain Questions](https://arxiv.org/abs/1704.00051) which is called DrQA by the authors. Specifically, DrQA is an end-to-end system for open domain question answering which involves an information retrieval system as well. This notebook however only explains the deep learning model proposed by them. This model is very similar to the one explained in [this](https://arxiv.org/abs/1606.02858) paper. The first authors in both the papers are also the same. The latter model is also known as \"Stanford Attentive Reader\" and is one of the models that is explained in Chris Manning's lecture on QA.   \n",
    "The flow of all the notebooks will be as follows:\n",
    "* Data Preprocessing: This section prepares the data for training and involves trademark NLP preprocessing steps. The functions being called here are imported from the script `preprocess.py`.\n",
    "* Model: I've tried to explain the intuition behind each layer/component of the model.\n",
    " > *Texts coming from the paper are in block quotes like this.*   \n",
    " \n",
    " Along with the intuition, I explain how equations written in paper can be transformed into code to the best of my ability.\n",
    "* Training: While I have tried my best to use similar training procedures as mentioned in the paper, there can be some changes. I do not have unlimited access to GPUs. \n",
    "* References: I provide an exhaustive list of resources/references that I used at the end of each notebook.\n",
    "\n",
    "\n",
    "### Tensor Based Approach\n",
    "All the notebooks are based on this approach. Ultimately, building neural nets is all about working with tensors. Knowing the shape and contents of each tensor is something that I have found very useful. Hence, after each line of code, I have commented the tensor shape and changes that happen due to the transformations in code. This makes the process of understanding whats going on in neural nets more intuitive.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data Preprocessing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import torchtext\n",
    "import torch\n",
    "from torch import nn\n",
    "import json, re, unicodedata, string, typing, time\n",
    "import torch.nn.functional as F\n",
    "import spacy\n",
    "from collections import Counter\n",
    "import pickle\n",
    "from nltk import word_tokenize\n",
    "nlp = spacy.load('en')\n",
    "from preprocess import *\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load dataset json files\n",
    "\n",
    "train_data = load_json('./data/squad_train.json')\n",
    "valid_data = load_json('./data/squad_dev.json')\n",
    "\n",
    "# parse the json structure to return the data as a list of dictionaries\n",
    "\n",
    "train_list = parse_data(train_data)\n",
    "valid_list = parse_data(valid_data)\n",
    "\n",
    "print('Train list len: ',len(train_list))\n",
    "print('Valid list len: ',len(valid_list))\n",
    "\n",
    "# converting the lists into dataframes\n",
    "\n",
    "train_df = pd.DataFrame(train_list)\n",
    "valid_df = pd.DataFrame(valid_list)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def normalize_spaces(text):\n",
    "    '''\n",
    "    Removes extra white spaces from the context.\n",
    "    '''\n",
    "    text = re.sub(r'\\s', ' ', text)\n",
    "    return text\n",
    "\n",
    "train_df.context = train_df.context.apply(normalize_spaces)\n",
    "valid_df.context = valid_df.context.apply(normalize_spaces)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "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>id</th>\n",
       "      <th>context</th>\n",
       "      <th>question</th>\n",
       "      <th>label</th>\n",
       "      <th>answer</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>5733be284776f41900661182</td>\n",
       "      <td>Architecturally, the school has a Catholic cha...</td>\n",
       "      <td>To whom did the Virgin Mary allegedly appear i...</td>\n",
       "      <td>[515, 541]</td>\n",
       "      <td>Saint Bernadette Soubirous</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>5733be284776f4190066117f</td>\n",
       "      <td>Architecturally, the school has a Catholic cha...</td>\n",
       "      <td>What is in front of the Notre Dame Main Building?</td>\n",
       "      <td>[188, 213]</td>\n",
       "      <td>a copper statue of Christ</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>5733be284776f41900661180</td>\n",
       "      <td>Architecturally, the school has a Catholic cha...</td>\n",
       "      <td>The Basilica of the Sacred heart at Notre Dame...</td>\n",
       "      <td>[279, 296]</td>\n",
       "      <td>the Main Building</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>5733be284776f41900661181</td>\n",
       "      <td>Architecturally, the school has a Catholic cha...</td>\n",
       "      <td>What is the Grotto at Notre Dame?</td>\n",
       "      <td>[381, 420]</td>\n",
       "      <td>a Marian place of prayer and reflection</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>5733be284776f4190066117e</td>\n",
       "      <td>Architecturally, the school has a Catholic cha...</td>\n",
       "      <td>What sits on top of the Main Building at Notre...</td>\n",
       "      <td>[92, 126]</td>\n",
       "      <td>a golden statue of the Virgin Mary</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                         id  \\\n",
       "0  5733be284776f41900661182   \n",
       "1  5733be284776f4190066117f   \n",
       "2  5733be284776f41900661180   \n",
       "3  5733be284776f41900661181   \n",
       "4  5733be284776f4190066117e   \n",
       "\n",
       "                                             context  \\\n",
       "0  Architecturally, the school has a Catholic cha...   \n",
       "1  Architecturally, the school has a Catholic cha...   \n",
       "2  Architecturally, the school has a Catholic cha...   \n",
       "3  Architecturally, the school has a Catholic cha...   \n",
       "4  Architecturally, the school has a Catholic cha...   \n",
       "\n",
       "                                            question       label  \\\n",
       "0  To whom did the Virgin Mary allegedly appear i...  [515, 541]   \n",
       "1  What is in front of the Notre Dame Main Building?  [188, 213]   \n",
       "2  The Basilica of the Sacred heart at Notre Dame...  [279, 296]   \n",
       "3                  What is the Grotto at Notre Dame?  [381, 420]   \n",
       "4  What sits on top of the Main Building at Notre...   [92, 126]   \n",
       "\n",
       "                                    answer  \n",
       "0               Saint Bernadette Soubirous  \n",
       "1                a copper statue of Christ  \n",
       "2                        the Main Building  \n",
       "3  a Marian place of prayer and reflection  \n",
       "4       a golden statue of the Virgin Mary  "
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wall time: 616 ms\n"
     ]
    }
   ],
   "source": [
    "# gather text to build vocabularies\n",
    "\n",
    "%time vocab_text = gather_text_for_vocab([train_df, valid_df])\n",
    "print(\"Number of sentences in dataset: \", len(vocab_text))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "raw-vocab: 111086\n",
      "vocab-length: 111088\n",
      "word2idx-length: 111088\n",
      "Wall time: 42.5 s\n"
     ]
    }
   ],
   "source": [
    "# build word vocabulary\n",
    "\n",
    "%time word2idx, idx2word, word_vocab = build_word_vocab(vocab_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wall time: 1min 49s\n",
      "Wall time: 45.8 s\n",
      "Wall time: 10.6 s\n",
      "Wall time: 4.07 s\n"
     ]
    }
   ],
   "source": [
    "# numericalize context and questions for training and validation set\n",
    "\n",
    "\n",
    "%time train_df['context_ids'] = train_df.context.apply(context_to_ids, word2idx=word2idx)\n",
    "%time valid_df['context_ids'] = valid_df.context.apply(context_to_ids, word2idx=word2idx)\n",
    "\n",
    "%time train_df['question_ids'] = train_df.question.apply(question_to_ids,  word2idx=word2idx)\n",
    "%time valid_df['question_ids'] = valid_df.question.apply(question_to_ids,  word2idx=word2idx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of error indices: 1002\n",
      "Number of error indices: 431\n"
     ]
    }
   ],
   "source": [
    "# get indices with tokenization errors and drop those indices \n",
    "\n",
    "train_err = get_error_indices(train_df, idx2word)\n",
    "valid_err = get_error_indices(valid_df, idx2word)\n",
    "\n",
    "train_df.drop(train_err, inplace=True)\n",
    "valid_df.drop(valid_err, inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# get start and end positions of answers from the context\n",
    "# this is basically the label for training QA models\n",
    "\n",
    "train_label_idx = train_df.apply(index_answer, axis=1, idx2word=idx2word)\n",
    "valid_label_idx = valid_df.apply(index_answer, axis=1, idx2word=idx2word)\n",
    "\n",
    "train_df['label_idx'] = train_label_idx\n",
    "valid_df['label_idx'] = valid_label_idx"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Dump data to pickle files \n",
    "This ensures that we can directly access the preprocessed dataframe next time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "with open('drqastoi.pickle','wb') as handle:\n",
    "    pickle.dump(word2idx, handle)\n",
    "    \n",
    "train_df.to_pickle('drqatrain.pkl')\n",
    "valid_df.to_pickle('drqavalid.pkl')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Read data from pickle files\n",
    "\n",
    "You only need to run the preprocessing once. Some preprocessing functions can take upto 3 mins. Therefore, pickling preprocessed data can save a lot of time.\n",
    "Once the preprocessed files are saved, you can directly start from here."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df = pd.read_pickle('drqatrain.pkl')\n",
    "valid_df = pd.read_pickle('drqavalid.pkl')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dataset/ Dataloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SquadDataset:\n",
    "    '''\n",
    "    -Divides the dataframe in batches.\n",
    "    -Pads the contexts and questions dynamically for each batch by padding \n",
    "     the examples to the maximum-length sequence in that batch.\n",
    "    -Calculates masks for context and question.\n",
    "    -Calculates spans for contexts.\n",
    "    '''\n",
    "    \n",
    "    def __init__(self, data, batch_size):\n",
    "        \n",
    "        self.batch_size = batch_size\n",
    "        data = [data[i:i+self.batch_size] for i in range(0, len(data), self.batch_size)]\n",
    "        self.data = data\n",
    "    \n",
    "    def get_span(self, text):\n",
    "        \n",
    "        text = nlp(text, disable=['parser','tagger','ner'])\n",
    "        span = [(w.idx, w.idx+len(w.text)) for w in text]\n",
    "\n",
    "        return span\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "    \n",
    "    def __iter__(self):\n",
    "        '''\n",
    "        Creates batches of data and yields them.\n",
    "        \n",
    "        Each yield comprises of:\n",
    "        :padded_context: padded tensor of contexts for each batch \n",
    "        :padded_question: padded tensor of questions for each batch \n",
    "        :context_mask & question_mask: zero-mask for question and context\n",
    "        :label: start and end index wrt context_ids\n",
    "        :context_text,answer_text: used while validation to calculate metrics\n",
    "        :context_spans: spans of context text\n",
    "        :ids: question_ids used in evaluation\n",
    "        '''\n",
    "        \n",
    "        for batch in self.data:\n",
    "                            \n",
    "            spans = []\n",
    "            context_text = []\n",
    "            answer_text = []\n",
    "            \n",
    "            max_context_len = max([len(ctx) for ctx in batch.context_ids])\n",
    "            padded_context = torch.LongTensor(len(batch), max_context_len).fill_(1)\n",
    "            \n",
    "            for ctx in batch.context:\n",
    "                context_text.append(ctx)\n",
    "                spans.append(self.get_span(ctx))\n",
    "            \n",
    "            for ans in batch.answer:\n",
    "                answer_text.append(ans)\n",
    "                \n",
    "            for i, ctx in enumerate(batch.context_ids):\n",
    "                padded_context[i, :len(ctx)] = torch.LongTensor(ctx)\n",
    "            \n",
    "            max_question_len = max([len(ques) for ques in batch.question_ids])\n",
    "            padded_question = torch.LongTensor(len(batch), max_question_len).fill_(1)\n",
    "            \n",
    "            for i, ques in enumerate(batch.question_ids):\n",
    "                padded_question[i,: len(ques)] = torch.LongTensor(ques)\n",
    "                \n",
    "            \n",
    "            label = torch.LongTensor(list(batch.label_idx))\n",
    "            context_mask = torch.eq(padded_context, 1)\n",
    "            question_mask = torch.eq(padded_question, 1)\n",
    "            \n",
    "            ids = list(batch.id)  \n",
    "            \n",
    "            yield (padded_context, padded_question, context_mask, \n",
    "                   question_mask, label, context_text, answer_text, ids)\n",
    "            \n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = SquadDataset(train_df, 32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "valid_dataset = SquadDataset(valid_df, 32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [],
   "source": [
    "a = next(iter(train_dataset))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(torch.Size([32, 253]),\n",
       " torch.Size([32, 19]),\n",
       " torch.Size([32, 253]),\n",
       " torch.Size([32, 19]),\n",
       " torch.Size([32, 2]))"
      ]
     },
     "execution_count": 59,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a[0].shape, a[1].shape, a[2].shape, a[3].shape, a[4].shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before we dive deep into the intricacies of the model, let's set up the notations. An input example during training is comprised of \n",
    "* a paragraph / context $p$ consisting of $l$ tokens { $p_{1}$, $p_{2}$,..., $p_{l}$ }\n",
    "* a question $q$ consisting of $m$ tokens { $q_{1}$, $q_{2}$,..., $q_{m}$ }\n",
    "* a start and and end position that comes from the context itself. More specifically, the start and end indices of the answer from the context  \n",
    "\n",
    "The following flowchart shows the flow of the model. It might not make sense now, but as we progress down the chart and build all the components, things will become clearer.\n",
    "\n",
    "\n",
    "<img src=\"images/drqaflow.PNG\" width=\"700\" height=\"800\"/>\n",
    "\n",
    "\n",
    "### Word Embedding\n",
    "\n",
    "The first transformation for both the question and the context tokens is that they are passed through an embedding layer initialized with pre-trained GloVe word vectors. 300-dimensional vectors from 840B web crawl version are used here. This version of GLoVe has a vocabulary of 2.2M words. Out of vocabulary or OOV words are initialized by a zero vector. OOV words are the words that are present in your dataset but not in the pretrained vocabulary of GLoVe.  \n",
    "These word vectors are used to project/convert a word into a floating point vector which encodes various features associated with the word into its dimensions. Such a conversion is necassary since computers cannot process words as strings but can seamlessly work with a large number of floating point matrices.   \n",
    "A dot product between vectors of word that are semantically similar is close to 1 and vice-versa.\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_glove_matrix():\n",
    "    '''\n",
    "    Parses the glove word vectors text file and returns a dictionary with the words as\n",
    "    keys and their respective pretrained word vectors as values.\n",
    "\n",
    "    '''\n",
    "    glove_dict = {}\n",
    "    with open(\"./glove.840B.300d/glove.840B.300d.txt\", \"r\", encoding=\"utf-8\") as f:\n",
    "        for line in f:\n",
    "            values = line.split(' ')\n",
    "            word = values[0]\n",
    "            vector = np.asarray(values[1:], dtype=\"float32\")\n",
    "            glove_dict[word] = vector\n",
    "\n",
    "    f.close()\n",
    "    \n",
    "    return glove_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "glove_dict = create_glove_matrix()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_word_embedding(glove_dict):\n",
    "    '''\n",
    "    Creates a weight matrix of the words that are common in the GloVe vocab and\n",
    "    the dataset's vocab. Initializes OOV words with a zero vector.\n",
    "    '''\n",
    "    weights_matrix = np.zeros((len(word_vocab), 300))\n",
    "    words_found = 0\n",
    "    for i, word in enumerate(word_vocab):\n",
    "        try:\n",
    "            weights_matrix[i] = glove_dict[word]\n",
    "            words_found += 1\n",
    "        except:\n",
    "            pass\n",
    "    return weights_matrix, words_found"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "weights_matrix, words_found = create_word_embedding(glove_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Total words found in glove vocab: \", words_found)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.save('drqaglove_vt.npy',weights_matrix)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Align Question Embedding\n",
    "\n",
    "The paper has different encoding procedures for the context and the question. The context/paragraph encoding is more exhaustive and comprises of following additional features:\n",
    "\n",
    "* exact match : encodes a binary feature if $p$ can be exactly matched to one word in question in its original, lemma or lowercase form\n",
    "* token features : Includes POS, NER and TF of context tokens and\n",
    "* aligned question embedding ($f_{align}$) .  \n",
    "\n",
    "In this re-implementation I have only implemented the aligned question embedding. The other features can be added easily but they do not affect the metrics by a large margin(~2).  \n",
    "$f_{align}$ has been formulated as shown below:\n",
    "\n",
    "$$ f_{align} = \\sum_{j}a_{i,j}E(q_{j}) $$ \n",
    "\n",
    "where $E()$ represents the glove embeddings and\n",
    "\n",
    "<img src=\"images/drqa1.PNG\" width=\"400\" height=\"200\"/>\n",
    "\n",
    "where $\\alpha()$ is a single dense layer with relu non-linearity. This transformation can be thought of as a projection to a new vector sub-space. The weights of the projection matrix will be learnt via backpropogation.\n",
    "These equations can be converted into code quite easily. Lets break this down into smaller chunks and understand what's going on actually. \n",
    "<img src=\"images/drqa2.PNG\" width=\"200\" height=\"150\"/>\n",
    "This is simply the product of projections of glove embeddings of the context and the question. Careful inspection of the equation for $a_{i,j}$ reveals that it is actually a softmax of the above product. The equations above depict everything at token level where $i$ represents a context token and $j$ represents a question token. Practically we usually vectorize our computations and deal with tensors directly.\n",
    "$f_{align}$ is a weighted representation of the question embeddings. $a_{i,j}$ represents the weights and hence a softmax function is necessary.  \n",
    "#### Intuition\n",
    "This feature enables the model to understand what portion of the context is more important or relevant with respect to the question. The products of projections taken at token level ensure a higher value when similar words from the question and context are multiplied. Quoting the paper,\n",
    "> *these features add soft alignments between similar but non-identical words (e.g., car and vehicle).* \n",
    "\n",
    "This is achieved via backpropation and training the weights of the dense layer. While this might seem a bit weird initially, we have to trust the process of backpropogation.   \n",
    "\n",
    "While implementing, we first calculate the projections of context and question vectors. We then use `torch.bmm` to calculate the product in the numerator of $a_{i,j}$, mask the product and then pass it through the softmax function to get $a_{i,j}$. Finally, we multiply this with the question embeddings. The output of this layer is an additional context embedding which is then concatenated with the glove embeddings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class AlignQuestionEmbedding(nn.Module):\n",
    "    \n",
    "    def __init__(self, input_dim):        \n",
    "        \n",
    "        super().__init__()\n",
    "        \n",
    "        self.linear = nn.Linear(input_dim, input_dim)\n",
    "        \n",
    "        self.relu = nn.ReLU()\n",
    "        \n",
    "    def forward(self, context, question, question_mask):\n",
    "        \n",
    "        # context = [bs, ctx_len, emb_dim]\n",
    "        # question = [bs, qtn_len, emb_dim]\n",
    "        # question_mask = [bs, qtn_len]\n",
    "    \n",
    "        ctx_ = self.linear(context)\n",
    "        ctx_ = self.relu(ctx_)\n",
    "        # ctx_ = [bs, ctx_len, emb_dim]\n",
    "        \n",
    "        qtn_ = self.linear(question)\n",
    "        qtn_ = self.relu(qtn_)\n",
    "        # qtn_ = [bs, qtn_len, emb_dim]\n",
    "        \n",
    "        qtn_transpose = qtn_.permute(0,2,1)\n",
    "        # qtn_transpose = [bs, emb_dim, qtn_len]\n",
    "        \n",
    "        align_scores = torch.bmm(ctx_, qtn_transpose)\n",
    "        # align_scores = [bs, ctx_len, qtn_len]\n",
    "        \n",
    "        qtn_mask = question_mask.unsqueeze(1).expand(align_scores.size())\n",
    "        # qtn_mask = [bs, 1, qtn_len] => [bs, ctx_len, qtn_len]\n",
    "        \n",
    "        # Fills elements of self tensor(align_scores) with value(-float(inf)) where mask is True. \n",
    "        # The shape of mask must be broadcastable with the shape of the underlying tensor.\n",
    "        align_scores = align_scores.masked_fill(qtn_mask == 1, -float('inf'))\n",
    "        # align_scores = [bs, ctx_len, qtn_len]\n",
    "        \n",
    "        align_scores_flat = align_scores.view(-1, question.size(1))\n",
    "        # align_scores = [bs*ctx_len, qtn_len]\n",
    "        \n",
    "        alpha = F.softmax(align_scores_flat, dim=1)\n",
    "        alpha = alpha.view(-1, context.shape[1], question.shape[1])\n",
    "        # alpha = [bs, ctx_len, qtn_len]\n",
    "        \n",
    "        align_embedding = torch.bmm(alpha, question)\n",
    "        # align = [bs, ctx_len, emb_dim]\n",
    "        \n",
    "        return align_embedding"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stacked BiLSTM\n",
    "\n",
    "The paragraph/context encoding which now has two features (glove and $f_{align}$) is then passed to a multilayer (3 layers) bidirectional LSTM. According to the paper,\n",
    "\n",
    "> *Speciﬁcally, we choose to use a multi-layer bidirectional long short-term memory network (LSTM), and take the concatenation of each layer’s hidden units in the end. *\n",
    "\n",
    "To achieve this functionality we cannot directly use the pytorch recurrent layers. Every recurrent layer in pytorch returns a tuple `[output, hidden]` where `output` holds the hidden states of all the timesteps from the __last layer only__. We need to access the hidden states of intermediate layers and then concatenate them at the end.\n",
    "The following figure illustrates this point in more detail.\n",
    "\n",
    "<img src=\"images/bilstm.png\" width=\"700\" height=\"600\"/>\n",
    "\n",
    "This figure shows a 3-layer bidirectional LSTM with an input sequence of size $n$. The green blocks denote the forward LSTMs and the blue blocks backward. Each block is labelled with the value that it calculates. The subscript denotes the time-step and the superscript denotes the depth or the layer-number. For example $hf_{1}^{(0)}$ calculates the first hidden state in forward LSTM in the first layer. \n",
    "\n",
    "As highlighted in the diagram, we need the intermediate hidden states passed between the layers along with the final output. To create this in code, we create a `nn.ModuleList` and add 3 LSTM layers to it. The input size of the first layer remains the same but for subsequent LSTMs the input size must be twice the hidden size. This is because the `output` of the first LSTM will have the dimension of `[batch_size, seq_len, hidden_size*num_directions]` and `num_directions` is 2 in our case. In the forward method, we loop through the LSTMs, store the hidden states of each layer and finally return the concatenated output. \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class StackedBiLSTM(nn.Module):\n",
    "    \n",
    "    def __init__(self, input_dim, hidden_dim, num_layers, dropout):\n",
    "        \n",
    "        super().__init__()\n",
    "        \n",
    "        self.dropout = dropout\n",
    "        \n",
    "        self.num_layers = num_layers\n",
    "        \n",
    "        self.lstms = nn.ModuleList()\n",
    "        \n",
    "        for i in range(self.num_layers):\n",
    "            \n",
    "            input_dim = input_dim if i == 0 else hidden_dim * 2\n",
    "            \n",
    "            self.lstms.append(nn.LSTM(input_dim, hidden_dim,\n",
    "                                      batch_first=True, bidirectional=True))\n",
    "           \n",
    "    \n",
    "    def forward(self, x):\n",
    "        # x = [bs, seq_len, feature_dim]\n",
    "\n",
    "        outputs = [x]\n",
    "        for i in range(self.num_layers):\n",
    "\n",
    "            lstm_input = outputs[-1]\n",
    "            lstm_out = F.dropout(lstm_input, p=self.dropout)\n",
    "            lstm_out, (hidden, cell) = self.lstms[i](lstm_input)\n",
    "           \n",
    "            outputs.append(lstm_out)\n",
    "\n",
    "    \n",
    "        output = torch.cat(outputs[1:], dim=2)\n",
    "        # [bs, seq_len, num_layers*num_dir*hidden_dim]\n",
    "        \n",
    "        output = F.dropout(output, p=self.dropout)\n",
    "      \n",
    "        return output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Linear Attention Layer\n",
    "\n",
    "The paper does not mention specifically the name of this layer and I have named this based on my understanding of its functionality. The previous layers were majorly about encoding and representing the context. This layer is used to encode the question and is much simpler than the previous layers. The question tokens are first passed through the glove embedding layer, then passed through the bilstm layer and finally reach this layer. \n",
    "This layer is used to calculate the importance of each word in the question. This can be achieved by simply taking a softmax over the input. However to add more learning capacity to the model, the inputs are multiplied by a trainable weight vector $w$ and then passed through a softmax function.  \n",
    "This layer calculates the weights as \n",
    "<img src=\"images/drqab.PNG\" width=\"300\" height=\"300\"/>\n",
    "\n",
    "Essentially the layer is performing \"attention\" on inputs. The $w$ in code is characterized by a linear layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LinearAttentionLayer(nn.Module):\n",
    "    \n",
    "    def __init__(self, input_dim):\n",
    "        \n",
    "        super().__init__()\n",
    "        \n",
    "        self.linear = nn.Linear(input_dim, 1)\n",
    "        \n",
    "    def forward(self, question, question_mask):\n",
    "        \n",
    "        # question = [bs, qtn_len, input_dim] = [bs, qtn_len, bi_lstm_hid_dim]\n",
    "        # question_mask = [bs,  qtn_len]\n",
    "        \n",
    "        qtn = question.view(-1, question.shape[-1])\n",
    "        # qtn = [bs*qtn_len, hid_dim]\n",
    "        \n",
    "        attn_scores = self.linear(qtn)\n",
    "        # attn_scores = [bs*qtn_len, 1]\n",
    "        \n",
    "        attn_scores = attn_scores.view(question.shape[0], question.shape[1])\n",
    "        # attn_scores = [bs, qtn_len]\n",
    "        \n",
    "        attn_scores = attn_scores.masked_fill(question_mask == 1, -float('inf'))\n",
    "        \n",
    "        alpha = F.softmax(attn_scores, dim=1)\n",
    "        # alpha = [bs, qtn_len]\n",
    "        \n",
    "        return alpha\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The following function just multiplies the weights calculated in the previous layer by the outputs of the question bilstm layer. This allows the model to assign higher values to important words in each question.\n",
    "\n",
    "$$ q = \\sum_{j} b_{j} q_{j} $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def weighted_average(x, weights):\n",
    "    # x = [bs, len, dim]\n",
    "    # weights = [bs, len]\n",
    "    \n",
    "    weights = weights.unsqueeze(1)\n",
    "    # weights = [bs, 1, len]\n",
    "    \n",
    "    w = weights.bmm(x).squeeze(1)\n",
    "    # w = [bs, 1, dim] => [bs, dim]\n",
    "    \n",
    "    return w"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Introduction to Attention in NLP\n",
    "\n",
    "Attention as a concept in NLP was introduced in 2015 by Bahdanau et al. to improve neural machine translation systems. Before this paper, NMT systems were largely based on seq2seq architectures which had an encoder to encode a representation of the source language and a decoder which to decode this representation into the target language. Such models were trained on large quantities of parallel text data of two languages. One major drawback of this architecture was that it didn't work well for longer documents/sequences. This is because the entire information in the source sentence was being crammed into a single vector. If this vector fails to capture the important information from the source language, the system is going to perform poorly.   \n",
    "<img src=\"images/seq2seq.PNG\" width=\"600\" height=\"700\"/>\n",
    "\n",
    "When we claim that these neural nets mimic the human brain, this is certainly not how the human brain works. While *learning* about some topic, we do not simply read 2-3 pages of content and expect our brain to remember all the details in the first go. We usually revisit various concepts, recollect and refer to the material back and forth before mastering it. The attention mechanism in NMT was designed to do this. While decoding at any particular time step, encoder hidden states from all the time-steps are made available to the decoder. The decoder then can look back at the encoder hidden states or the source language and make a more informed prediction at a particular time-step. This alieviates the problem of all the information from source language being crammed into a single vector.  \n",
    "To illustrate this with equations, consider that the hidden states of the encoder RNN are represented by $H$ = {$h_{1}, h_{2}, h_{3},...,h_{t}$}. While decoding the token at position $t$, the input to the decoder unit is hidden state from previous unit $s_{t-1}$ and an attention vector which is a selective summary of the encoder hidden states and helps the decoder to pay more attention to a particular encoder state. \n",
    "The similarity between the encoder hidden states $H$ and the decoder hidden state so far $s_{t-1}$ is computed by,  \n",
    "$$ \\alpha = tanh (W [H ; s_{t-1}]) $$   \n",
    "\n",
    "$\\alpha$ is then passed through a softmax layer to obtain attention distribution such that $\\sum_{t} \\alpha_{t}$ = 1.\n",
    "The final step is calculating the attention vector by taking a weighted sum of the encoder hidden states,\n",
    "$$ \\sum_{t} \\alpha_{t} h_{t} $$\n",
    "\n",
    "The following diagram illustrates this process.  \n",
    " \n",
    "<img src=\"images/attnkj.PNG\" width=\"600\" height=\"100\"/>\n",
    "\n",
    "\n",
    "Since then, many different forms of attention have been proposed and used in the literature. Attention is not limited to NMT systems and has evolved into a more general concept in NLP. At the heart of it attention is about summarizing a particular entity/representation by *attending* to the important parts of this representation. \n",
    "A more general definition of attention is as follows:\n",
    "\n",
    "> *__Given a set of vectors `values`, and a single vector `query`, attention is a method to calculate a weighted sum of the values, where the query determines which values to focus on.__*\n",
    "\n",
    "It is a way to obtain a fixed size representation of an arbitrary set of representations (values), dependent on some other representation (query).   \n",
    "\n",
    "In our earlier NMT example, the encoder hidden states {$h_{1}, h_{2}, h_{3},...,h_{t}$} are the __*values*__ and the decoder hidden state $s_{t-1}$ is the __*query*__.  \n",
    "\n",
    "### A More General Take On Attention\n",
    "\n",
    "In general there are 3 steps when calculating the attention. Consider that values are represented by {$h_{1}, h_{2}, h_{3},..h_{n}$} and query is $s$. Then attention always involves,\n",
    "\n",
    "1. Calculating the energy $e$ or attention scores between these 2 vectors,\n",
    "$e$   $ \\epsilon$  $ R^{N} $\n",
    "2. Taking softmax to get an attention distribution $\\alpha$, $\\alpha$ $\\epsilon$ $R^{N}$\n",
    "\n",
    "$$ \\alpha = softmax(e)$$ \n",
    "$$ \\sum_{t}^{N} \\alpha_{t} = 1 $$\n",
    "\n",
    "3. Taking the weighted sum of the `values` by using $\\alpha$\n",
    "$$ a = \\sum_{t}^{N}\\alpha_{t}h_{t} $$\n",
    "\n",
    "\n",
    "Now there are different ways to calculate the energy between `query` and `values`. \n",
    "* **Basic Dot Product Attention**    \n",
    "$$ e_{t} = s^{T}h_{t}$$      \n",
    "* **Additive Attention**\n",
    "$$ e_{t} = v^{T} tanh (W [h_{t};s])$$  \n",
    "This is nothing but the Bahdanau attention first proposed for NMT systems.\n",
    "* **Scaled Dot Product Attention**\n",
    "$$ e_{t} = s^{T}h_{t}/\\sqrt n$$\n",
    "where $n$ is the model size. A modified version of this proposed in the Transformers paper by Vaswani et al. is now employed in almost every NLP system.\n",
    "\n",
    "* **Bilinear Attention**\n",
    "$$ e_{t} = s^{T} W h_{t}$$\n",
    "where $W$ is a trainable weight vector.\n",
    "This is the method used in this paper to predict the start and end position of the answer from the context.    \n",
    "\n",
    "\n",
    "To implement this layer, we characterise $W$ by a linear layer.\n",
    "First the linear layer is applied to the question, which is equivalent to the product $W.q$. This product is then multiplied by the context using `torch.bmm`.   \n",
    "Note that softmax is not taken over here to get the weights. This is taken care of when we calculate the loss using cross entropy. The following layer does not actually calculate the attention as a weighted sum. It just uses the bilinear term's representation to predict the span. However the intuition behind the bilinear term still remains the same."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BilinearAttentionLayer(nn.Module):\n",
    "    \n",
    "    def __init__(self, context_dim, question_dim):\n",
    "        \n",
    "        super().__init__()\n",
    "        \n",
    "        self.linear = nn.Linear(question_dim, context_dim)\n",
    "        \n",
    "    def forward(self, context, question, context_mask):\n",
    "        \n",
    "        # context = [bs, ctx_len, ctx_hid_dim] = [bs, ctx_len, hid_dim*6] = [bs, ctx_len, 768]\n",
    "        # question = [bs, qtn_hid_dim] = [bs, qtn_len, 768]\n",
    "        # context_mask = [bs, ctx_len]\n",
    "        \n",
    "        qtn_proj = self.linear(question)\n",
    "        # qtn_proj = [bs, ctx_hid_dim]\n",
    "        \n",
    "        qtn_proj = qtn_proj.unsqueeze(2)\n",
    "        # qtn_proj = [bs, ctx_hid_dim, 1]\n",
    "        \n",
    "        scores = context.bmm(qtn_proj)\n",
    "        # scores = [bs, ctx_len, 1]\n",
    "        \n",
    "        scores = scores.squeeze(2)\n",
    "        # scores = [bs, ctx_len]\n",
    "        \n",
    "        scores = scores.masked_fill(context_mask == 1, -float('inf'))\n",
    "        \n",
    "        #alpha = F.log_softmax(scores, dim=1)\n",
    "        # alpha = [bs, ctx_len]\n",
    "        \n",
    "        return scores"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Putting it together\n",
    "\n",
    "The following module brings all the components discussed so far together. It takes in the context and question tokens as inputs and returns the start and end positions of the answer from the context.  \n",
    "\n",
    "<img src=\"images/drqaflow.PNG\" width=\"600\" height=\"600\"/>\n",
    "\n",
    "  \n",
    "Going down the flowchart, following steps are performed in sequence:  \n",
    "* The context and question tokens are passed through the Glove embedding layer. The glove embeddings are partially finetuned during training. According to the paper,  \n",
    "> *We keep most of the pre-trained word embeddings ﬁxed and only ﬁne-tune the 1000 most frequent question words because the representations of some key words such as what, how, which, many could be crucial for QA systems.*   \n",
    "\n",
    "In code, this is done by using hooks in pytorch. Hooks work as a callback functions and are executed after `forward` or `backward` function is called for a particular tensor. You should read more about this in their documentation.\n",
    "\n",
    "* Aligned question embedding is calculated for the context vector and concatenated (using `torch.cat`) to the context representation. If $d$ is the embedding dimension then context $\\epsilon$ $R^{2d}$ and question $\\epsilon$ $R^{d}$.\n",
    "* Context and question representations are then passed to bilstm layers to obtain tensors of dimension `[batch_size, seq_len, hidden_dim*6]` since the LSTM is bidirectional and there are 3 layers of it.\n",
    "* The embedded question is also passed through the linear attention layer and a weighted sum of its output is taken with the biLSTM output.\n",
    "* Both these representations are finally passed through the bilinear attention layer to predict the start and end position of the answer.   \n",
    "\n",
    "An intriguing point here is that the same set of weights are passed to the bilinear attention layers. Yet how do they predict different things. This is left over to the neural network to learn. Our loss function ensures that our objective is to predict different positions from the context. It is now the neural net's responsibility to learn different weights for each layer. It is sort of a \"black-box\" and we have to trust the process of backpropogation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DocumentReader(nn.Module):\n",
    "    \n",
    "    def __init__(self, hidden_dim, embedding_dim, num_layers, num_directions, dropout, device):\n",
    "        \n",
    "        super().__init__()\n",
    "        \n",
    "        self.device = device\n",
    "        \n",
    "        #self.embedding = self.get_glove_embedding()\n",
    "        \n",
    "        self.context_bilstm = StackedBiLSTM(embedding_dim * 2, hidden_dim, num_layers, dropout)\n",
    "        \n",
    "        self.question_bilstm = StackedBiLSTM(embedding_dim, hidden_dim, num_layers, dropout)\n",
    "        \n",
    "        self.glove_embedding = self.get_glove_embedding()\n",
    "        \n",
    "        def tune_embedding(grad, words=1000):\n",
    "            grad[words:] = 0\n",
    "            return grad\n",
    "        \n",
    "        self.glove_embedding.weight.register_hook(tune_embedding)\n",
    "        \n",
    "        self.align_embedding = AlignQuestionEmbedding(embedding_dim)\n",
    "        \n",
    "        self.linear_attn_question = LinearAttentionLayer(hidden_dim*num_layers*num_directions) \n",
    "        \n",
    "        self.bilinear_attn_start = BilinearAttentionLayer(hidden_dim*num_layers*num_directions, \n",
    "                                                          hidden_dim*num_layers*num_directions)\n",
    "        \n",
    "        self.bilinear_attn_end = BilinearAttentionLayer(hidden_dim*num_layers*num_directions,\n",
    "                                                        hidden_dim*num_layers*num_directions)\n",
    "        \n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "   \n",
    "        \n",
    "    def get_glove_embedding(self):\n",
    "        \n",
    "        weights_matrix = np.load('drqaglove_vt.npy')\n",
    "        num_embeddings, embedding_dim = weights_matrix.shape\n",
    "        embedding = nn.Embedding.from_pretrained(torch.FloatTensor(weights_matrix).to(self.device),freeze=False)\n",
    "\n",
    "        return embedding\n",
    "    \n",
    "    \n",
    "    def forward(self, context, question, context_mask, question_mask):\n",
    "       \n",
    "        # context = [bs, len_c]\n",
    "        # question = [bs, len_q]\n",
    "        # context_mask = [bs, len_c]\n",
    "        # question_mask = [bs, len_q]\n",
    "        \n",
    "        \n",
    "        ctx_embed = self.glove_embedding(context)\n",
    "        # ctx_embed = [bs, len_c, emb_dim]\n",
    "        \n",
    "        ques_embed = self.glove_embedding(question)\n",
    "        # ques_embed = [bs, len_q, emb_dim]\n",
    "        \n",
    "\n",
    "        ctx_embed = self.dropout(ctx_embed)\n",
    "     \n",
    "        ques_embed = self.dropout(ques_embed)\n",
    "             \n",
    "        align_embed = self.align_embedding(ctx_embed, ques_embed, question_mask)\n",
    "        # align_embed = [bs, len_c, emb_dim]  \n",
    "        \n",
    "        ctx_bilstm_input = torch.cat([ctx_embed, align_embed], dim=2)\n",
    "        # ctx_bilstm_input = [bs, len_c, emb_dim*2]\n",
    "                \n",
    "        ctx_outputs = self.context_bilstm(ctx_bilstm_input)\n",
    "        # ctx_outputs = [bs, len_c, hid_dim*layers*dir] = [bs, len_c, hid_dim*6]\n",
    "       \n",
    "        qtn_outputs = self.question_bilstm(ques_embed)\n",
    "        # qtn_outputs = [bs, len_q, hid_dim*6]\n",
    "    \n",
    "        qtn_weights = self.linear_attn_question(qtn_outputs, question_mask)\n",
    "        # qtn_weights = [bs, len_q]\n",
    "            \n",
    "        qtn_weighted = weighted_average(qtn_outputs, qtn_weights)\n",
    "        # qtn_weighted = [bs, hid_dim*6]\n",
    "        \n",
    "        start_scores = self.bilinear_attn_start(ctx_outputs, qtn_weighted, context_mask)\n",
    "        # start_scores = [bs, len_c]\n",
    "         \n",
    "        end_scores = self.bilinear_attn_end(ctx_outputs, qtn_weighted, context_mask)\n",
    "        # end_scores = [bs, len_c]\n",
    "        \n",
    "      \n",
    "        return start_scores, end_scores"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "###  Hyperparameters\n",
    "\n",
    "> *We use 3-layer bidirectional LSTMs with h = 128 hidden units for both paragraph and question encoding. Dropout with p = 0.3 is applied to word embeddings and all the hidden units of LSTMs. *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device('cuda')\n",
    "HIDDEN_DIM = 128\n",
    "EMB_DIM = 300\n",
    "NUM_LAYERS = 3\n",
    "NUM_DIRECTIONS = 2\n",
    "DROPOUT = 0.3\n",
    "device = torch.device('cuda')\n",
    "\n",
    "model = DocumentReader(HIDDEN_DIM,\n",
    "                       EMB_DIM, \n",
    "                       NUM_LAYERS, \n",
    "                       NUM_DIRECTIONS, \n",
    "                       DROPOUT, \n",
    "                       device).to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = torch.optim.Adamax(model.parameters())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The model has 37,367,549 trainable parameters\n"
     ]
    }
   ],
   "source": [
    "def count_parameters(model):\n",
    "    '''Returns the number of trainable parameters in the 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": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(model, train_dataset):\n",
    "    '''\n",
    "    Trains the model.\n",
    "    '''\n",
    "    \n",
    "    print(\"Starting training ........\")\n",
    "    \n",
    "    train_loss = 0.\n",
    "    batch_count = 0\n",
    "    \n",
    "    # put the model in training mode\n",
    "    model.train()\n",
    "    \n",
    "    # iterate through training data\n",
    "    for batch in train_dataset:\n",
    "\n",
    "        if batch_count % 500 == 0:\n",
    "            print(f\"Starting batch: {batch_count}\")\n",
    "        batch_count += 1\n",
    "\n",
    "        context, question, context_mask, question_mask, label, ctx, ans, ids = batch\n",
    "        \n",
    "        # place the tensors on GPU\n",
    "        context, context_mask, question, question_mask, label = context.to(device), context_mask.to(device),\\\n",
    "                                    question.to(device), question_mask.to(device), label.to(device)\n",
    "        \n",
    "        # forward pass, get the predictions\n",
    "        preds = model(context, question, context_mask, question_mask)\n",
    "\n",
    "        start_pred, end_pred = preds\n",
    "        \n",
    "        # separate labels for start and end position\n",
    "        start_label, end_label = label[:,0], label[:,1]\n",
    "        \n",
    "        # calculate loss\n",
    "        loss = F.cross_entropy(start_pred, start_label) + F.cross_entropy(end_pred, end_label)\n",
    "        \n",
    "        # backward pass, calculates the gradients\n",
    "        loss.backward()\n",
    "        \n",
    "        # gradient clipping\n",
    "        torch.nn.utils.clip_grad_norm_(model.parameters(), 10)\n",
    "        \n",
    "        # update the gradients\n",
    "        optimizer.step()\n",
    "        \n",
    "        # zero the gradients to prevent them from accumulating\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        train_loss += loss.item()\n",
    "\n",
    "    return train_loss/len(train_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "def valid(model, valid_dataset):\n",
    "    '''\n",
    "    Performs validation.\n",
    "    '''\n",
    "    \n",
    "    print(\"Starting validation .........\")\n",
    "   \n",
    "    valid_loss = 0.\n",
    "\n",
    "    batch_count = 0\n",
    "    \n",
    "    f1, em = 0., 0.\n",
    "    \n",
    "    # puts the model in eval mode. Turns off dropout\n",
    "    model.eval()\n",
    "    \n",
    "    predictions = {}\n",
    "    \n",
    "    for batch in valid_dataset:\n",
    "\n",
    "        if batch_count % 500 == 0:\n",
    "            print(f\"Starting batch {batch_count}\")\n",
    "        batch_count += 1\n",
    "\n",
    "        context, question, context_mask, question_mask, label, context_text, answers, ids = batch\n",
    "\n",
    "        context, context_mask, question, question_mask, label = context.to(device), context_mask.to(device),\\\n",
    "                                    question.to(device), question_mask.to(device), label.to(device)\n",
    "\n",
    "        with torch.no_grad():\n",
    "\n",
    "            preds = model(context, question, context_mask, question_mask)\n",
    "\n",
    "            p1, p2 = preds\n",
    "\n",
    "            y1, y2 = label[:,0], label[:,1]\n",
    "\n",
    "            loss = F.cross_entropy(p1, y1) + F.cross_entropy(p2, y2)\n",
    "\n",
    "            valid_loss += loss.item()\n",
    "\n",
    "            \n",
    "            # get the start and end index positions from the model preds\n",
    "            \n",
    "            batch_size, c_len = p1.size()\n",
    "            ls = nn.LogSoftmax(dim=1)\n",
    "            mask = (torch.ones(c_len, c_len) * float('-inf')).to(device).tril(-1).unsqueeze(0).expand(batch_size, -1, -1)\n",
    "            \n",
    "            score = (ls(p1).unsqueeze(2) + ls(p2).unsqueeze(1)) + mask\n",
    "            score, s_idx = score.max(dim=1)\n",
    "            score, e_idx = score.max(dim=1)\n",
    "            s_idx = torch.gather(s_idx, 1, e_idx.view(-1, 1)).squeeze()\n",
    "            \n",
    "            # stack predictions\n",
    "            for i in range(batch_size):\n",
    "                id = ids[i]\n",
    "                pred = context[i][s_idx[i]:e_idx[i]+1]\n",
    "                pred = ' '.join([idx2word[idx.item()] for idx in pred])\n",
    "                predictions[id] = pred\n",
    "            \n",
    "            \n",
    "            \n",
    "    em, f1 = evaluate(predictions)            \n",
    "    return valid_loss/len(valid_dataset), em, f1\n",
    "                "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(predictions):\n",
    "    '''\n",
    "    Gets a dictionary of predictions with question_id as key\n",
    "    and prediction as value. The validation dataset has multiple \n",
    "    answers for a single question. Hence we compare our prediction\n",
    "    with all the answers and choose the one that gives us\n",
    "    the maximum metric (em or f1). \n",
    "    This method first parses the JSON file, gets all the answers\n",
    "    for a given id and then passes the list of answers and the \n",
    "    predictions to calculate em, f1.\n",
    "    \n",
    "    \n",
    "    :param dict predictions\n",
    "    Returns\n",
    "    : exact_match: 1 if the prediction and ground truth \n",
    "      match exactly, 0 otherwise.\n",
    "    : f1_score: \n",
    "    '''\n",
    "    with open('./data/squad_dev.json','r',encoding='utf-8') as f:\n",
    "        dataset = json.load(f)\n",
    "        \n",
    "    dataset = dataset['data']\n",
    "    f1 = exact_match = total = 0\n",
    "    for article in dataset:\n",
    "        for paragraph in article['paragraphs']:\n",
    "            for qa in paragraph['qas']:\n",
    "                total += 1\n",
    "                if qa['id'] not in predictions:\n",
    "                    continue\n",
    "                \n",
    "                ground_truths = list(map(lambda x: x['text'], qa['answers']))\n",
    "                \n",
    "                prediction = predictions[qa['id']]\n",
    "                \n",
    "                exact_match += metric_max_over_ground_truths(\n",
    "                    exact_match_score, prediction, ground_truths)\n",
    "                \n",
    "                f1 += metric_max_over_ground_truths(\n",
    "                    f1_score, prediction, ground_truths)\n",
    "                \n",
    "    \n",
    "    exact_match = 100.0 * exact_match / total\n",
    "    f1 = 100.0 * f1 / total\n",
    "    \n",
    "    return exact_match, f1\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "def normalize_answer(s):\n",
    "    '''\n",
    "    Performs a series of cleaning steps on the ground truth and \n",
    "    predicted answer.\n",
    "    '''\n",
    "    def remove_articles(text):\n",
    "        return re.sub(r'\\b(a|an|the)\\b', ' ', text)\n",
    "\n",
    "    def white_space_fix(text):\n",
    "        return ' '.join(text.split())\n",
    "\n",
    "    def remove_punc(text):\n",
    "        exclude = set(string.punctuation)\n",
    "        return ''.join(ch for ch in text if ch not in exclude)\n",
    "\n",
    "    def lower(text):\n",
    "        return text.lower()\n",
    "\n",
    "    return white_space_fix(remove_articles(remove_punc(lower(s))))\n",
    "\n",
    "\n",
    "def metric_max_over_ground_truths(metric_fn, prediction, ground_truths):\n",
    "    '''\n",
    "    Returns maximum value of metrics for predicition by model against\n",
    "    multiple ground truths.\n",
    "    \n",
    "    :param func metric_fn: can be 'exact_match_score' or 'f1_score'\n",
    "    :param str prediction: predicted answer span by the model\n",
    "    :param list ground_truths: list of ground truths against which\n",
    "                               metrics are calculated. Maximum values of \n",
    "                               metrics are chosen.\n",
    "                            \n",
    "    \n",
    "    '''\n",
    "    scores_for_ground_truths = []\n",
    "    for ground_truth in ground_truths:\n",
    "        score = metric_fn(prediction, ground_truth)\n",
    "        scores_for_ground_truths.append(score)\n",
    "        \n",
    "    return max(scores_for_ground_truths)\n",
    "\n",
    "\n",
    "def f1_score(prediction, ground_truth):\n",
    "    '''\n",
    "    Returns f1 score of two strings.\n",
    "    '''\n",
    "    prediction_tokens = normalize_answer(prediction).split()\n",
    "    ground_truth_tokens = normalize_answer(ground_truth).split()\n",
    "    common = Counter(prediction_tokens) & Counter(ground_truth_tokens)\n",
    "    num_same = sum(common.values())\n",
    "    if num_same == 0:\n",
    "        return 0\n",
    "    precision = 1.0 * num_same / len(prediction_tokens)\n",
    "    recall = 1.0 * num_same / len(ground_truth_tokens)\n",
    "    f1 = (2 * precision * recall) / (precision + recall)\n",
    "    return f1\n",
    "\n",
    "\n",
    "def exact_match_score(prediction, ground_truth):\n",
    "    '''\n",
    "    Returns exact_match_score of two strings.\n",
    "    '''\n",
    "    return (normalize_answer(prediction) == normalize_answer(ground_truth))\n",
    "\n",
    "\n",
    "def epoch_time(start_time, end_time):\n",
    "    '''\n",
    "    Helper function to record epoch time.\n",
    "    '''\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": 23,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "\n",
    "train_losses = []\n",
    "valid_losses = []\n",
    "ems = []\n",
    "f1s = []\n",
    "epochs = 5\n",
    "\n",
    "for epoch in range(epochs):\n",
    "    print(f\"Epoch {epoch+1}\")\n",
    "    \n",
    "    start_time = time.time()\n",
    "    \n",
    "    train_loss = train(model, train_dataset)\n",
    "    valid_loss, em, f1 = valid(model, valid_dataset)\n",
    "    \n",
    "    end_time = time.time()\n",
    "    \n",
    "    epoch_mins, epoch_secs = epoch_time(start_time, end_time)\n",
    "    \n",
    "    train_losses.append(train_loss)\n",
    "    valid_losses.append(valid_loss)\n",
    "    ems.append(em)\n",
    "    f1s.append(f1)\n",
    "    \n",
    "    print(f\"Epoch train loss : {train_loss}| Time: {epoch_mins}m {epoch_secs}s\")\n",
    "    print(f\"Epoch valid loss: {valid_loss}\")\n",
    "    print(f\"Epoch EM: {em}\")\n",
    "    print(f\"Epoch F1: {f1}\")\n",
    "    print(\"====================================================================================\")\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## References\n",
    "\n",
    "* Papers read/referenced\n",
    "    1. https://arxiv.org/abs/1704.00051\n",
    "    2. https://arxiv.org/abs/1606.02858\n",
    "    3. https://arxiv.org/abs/1409.0473\n",
    "* Other helpful links\n",
    "    1. https://lilianweng.github.io/lil-log/2018/06/24/attention-attention.html\n",
    "    2. https://github.com/facebookresearch/DrQA\n",
    "    3. https://github.com/hitvoice/DrQA. Special thanks to [Runqi Yang](https://github.com/hitvoice) who helped me clarify some doubts with respect to preprocessing the SQUAD dataset.\n",
    "    4. https://towardsdatascience.com/the-definitive-guide-to-bidaf-part-3-attention-92352bbdcb07. Good introduction to attention.\n",
    "    5. https://web.stanford.edu/class/archive/cs/cs224n/cs224n.1184/lectures/lecture10.pdf. The attention section of this notebook is largely inspired and derived from these slides.\n",
    "* Following links are related to debugging neural nets. Something on which I was stuck for quite some time during training these models.\n",
    "    1. https://datascience.stackexchange.com/questions/410/choosing-a-learning-rate\n",
    "    2. https://www.jeremyjordan.me/nn-learning-rate/\n",
    "    3. https://towardsdatascience.com/estimating-optimal-learning-rate-for-a-deep-neural-network-ce32f2556ce0\n",
    "    4. https://towardsdatascience.com/learning-rate-schedules-and-adaptive-learning-rate-methods-for-deep-learning-2c8f433990d1\n",
    "    5. https://towardsdatascience.com/checklist-for-debugging-neural-networks-d8b2a9434f21\n",
    "    6. https://arxiv.org/abs/1708.07120\n",
    "    7. https://sgugger.github.io/how-do-you-find-a-good-learning-rate.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}