3. QANet.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## QANet\n",
    "\n",
    "### Introduction \n",
    "\n",
    "The papers that we've seen so far have been heavily based on recurrent neural nets and attention. However, RNNs are slow to train given their sequential nature and are also slow for inference. QANet was proposed in early 2018. This paper does away with recurrence and is only based on self-attention and convolutions. This paper derives its major ideas from the \"Attention is all you need\".  \n",
    "Progression of ideas for any NLP task is usually derived from progress in other fields. For instance, most of the QA models employ methods that have been proven successful in Machine Translation (RNN with Attenion, self-attention), Language Modeling (BERT, ALBERT etc) etc. \n",
    "\n",
    ">  *We instead exclusively use convolutions and self-attentions as the building blocks of encoders that separately encodes the query and context. Then we learn the interactions between context and question by standard attentions. The resulting representation is encoded again with our recurrency-free encoder before ﬁnally decoding to the probability of each position being the start or end of the answer span. We call this architecture QANet*\n",
    "\n",
    "> *The key motivation behind the design of our model is the following: convolution captures the local structure of the text, while the self-attention learns the global interaction between each pair of words. *\n",
    "\n",
    "Let's get into the model.  \n",
    "Note: An exhaustive list of resources/references that I followed during this has been given in the end."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch import nn\n",
    "import torch\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import pickle, time\n",
    "import re, os, string, typing, gc, json\n",
    "import torch.nn.functional as F\n",
    "import spacy\n",
    "from sklearn.model_selection import train_test_split\n",
    "from collections import Counter\n",
    "nlp = spacy.load('en')\n",
    "from preprocess import *\n",
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data preprocessing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Length of data:  442\n",
      "Data Keys:  dict_keys(['title', 'paragraphs'])\n",
      "Title:  University_of_Notre_Dame\n",
      "Length of data:  48\n",
      "Data Keys:  dict_keys(['title', 'paragraphs'])\n",
      "Title:  Super_Bowl_50\n"
     ]
    }
   ],
   "source": [
    "# load dataset json files\n",
    "\n",
    "train_data = load_json('data/squad_train.json')\n",
    "valid_data = load_json('data/squad_dev.json')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train list len:  87599\n",
      "Valid list len:  34726\n"
     ]
    }
   ],
   "source": [
    "print('Train list len: ',len(train_list))\n",
    "print('Valid list len: ',len(valid_list))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 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": 42,
   "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": 42,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wall time: 2min 43s\n",
      "Wall time: 1min 3s\n"
     ]
    }
   ],
   "source": [
    "# get indices of outliers and drop them from the dataframe\n",
    "\n",
    "%time drop_ids_train = filter_large_examples(train_df)\n",
    "train_df.drop(list(drop_ids_train), inplace=True)\n",
    "\n",
    "%time drop_ids_valid = filter_large_examples(valid_df)\n",
    "valid_df.drop(list(drop_ids_valid), inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of sentences in the dataset:  118441\n"
     ]
    }
   ],
   "source": [
    "# gather text to build vocabularies\n",
    "\n",
    "vocab_text = gather_text_for_vocab([train_df, valid_df])\n",
    "print(\"Number of sentences in the dataset: \", len(vocab_text))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "raw-vocab: 110479\n",
      "glove-vocab: 110479\n",
      "vocab-length: 110481\n",
      "word2idx-length: 110481\n",
      "Wall time: 38.1 s\n",
      "----------------------------------\n",
      "raw-char-vocab: 1401\n",
      "char-vocab-intersect: 232\n",
      "char2idx-length: 234\n",
      "Wall time: 3.3 s\n"
     ]
    }
   ],
   "source": [
    "# build word and character-level vocabularies\n",
    "\n",
    "%time word2idx, idx2word, word_vocab = build_word_vocab(vocab_text)\n",
    "print(\"----------------------------------\")\n",
    "%time char2idx, char_vocab = build_char_vocab(vocab_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wall time: 44.6 s\n",
      "Wall time: 9.77 s\n",
      "Wall time: 4.06 s\n"
     ]
    }
   ],
   "source": [
    "# numericalize context and questions for training and validation set\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",
    "%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": 50,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of error indices: 1000\n",
      "Number of error indices: 428\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": 51,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(86335, 34011)"
      ]
     },
     "execution_count": 51,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(train_df), len(valid_df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "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": 44,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df.to_pickle('qanettrain.pkl')\n",
    "valid_df.to_pickle('qanetvalid.pkl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "with open('qanetw2id.pickle','wb') as handle:\n",
    "    pickle.dump(word2idx, handle)\n",
    "\n",
    "with open('qanetc2id.pickle','wb') as handle:\n",
    "    pickle.dump(char2idx, handle)"
   ]
  },
  {
   "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pickle\n",
    "\n",
    "with open('qanetw2id.pickle','rb') as handle:\n",
    "    word2idx = pickle.load(handle)\n",
    "with open('qanetc2id.pickle','rb') as handle:\n",
    "    char2idx = pickle.load(handle)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df = pd.read_pickle('qanettrain.pkl')\n",
    "valid_df = pd.read_pickle('qanetvalid.pkl')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "idx2word = {v:k for k,v in word2idx.items()}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating the dataloader\n",
    "\n",
    "This class takes care of batching, creating character vectors and returns all the things needed during training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "class SquadDataset:\n",
    "    '''\n",
    "    - Creates batches dynamically by padding to the length of largest example\n",
    "      in a given batch.\n",
    "    - Calulates character vectors for contexts and question.\n",
    "    - Returns tensors for training.\n",
    "    '''\n",
    "    def __init__(self, data, batch_size):\n",
    "        '''\n",
    "        data: dataframe\n",
    "        batch_size: int\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",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.data)\n",
    "    \n",
    "    def make_char_vector(self, max_sent_len, sentence, max_word_len=16):\n",
    "        \n",
    "        char_vec = torch.zeros(max_sent_len, max_word_len).type(torch.LongTensor)\n",
    "        \n",
    "        for i, word in enumerate(nlp(sentence, disable=['parser','tagger','ner'])):\n",
    "            for j, ch in enumerate(word.text):\n",
    "                if j == max_word_len:\n",
    "                    break\n",
    "                char_vec[i][j] = char2idx.get(ch, 0)\n",
    "        \n",
    "        return char_vec     \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",
    "    \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",
    "        :char_ctx & ques_ctx: character-level ids for context and question\n",
    "        :label: start and end index wrt context_ids\n",
    "        :context_text,answer_text: used while validation to calculate metrics\n",
    "        :ids: question_ids for evaluation\n",
    "        '''\n",
    "        \n",
    "        for batch in self.data:\n",
    "            \n",
    "            spans = []\n",
    "            ctx_text = []\n",
    "            answer_text = []\n",
    "            \n",
    "             \n",
    "            for ctx in batch.context:\n",
    "                ctx_text.append(ctx)\n",
    "                spans.append(self.get_span(ctx))\n",
    "            \n",
    "            for ans in batch.answer:\n",
    "                answer_text.append(ans)\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 i, ctx in enumerate(batch.context_ids):\n",
    "                padded_context[i, :len(ctx)] = torch.LongTensor(ctx)\n",
    "                \n",
    "            max_word_ctx = 16\n",
    "          \n",
    "            char_ctx = torch.zeros(len(batch), max_context_len, max_word_ctx).type(torch.LongTensor)\n",
    "            for i, context in enumerate(batch.context):\n",
    "                char_ctx[i] = self.make_char_vector(max_context_len, context)\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",
    "            max_word_ques = 16\n",
    "            \n",
    "            char_ques = torch.zeros(len(batch), max_question_len, max_word_ques).type(torch.LongTensor)\n",
    "            for i, question in enumerate(batch.question):\n",
    "                char_ques[i] = self.make_char_vector(max_question_len, question)\n",
    "            \n",
    "              \n",
    "            label = torch.LongTensor(list(batch.label_idx))\n",
    "            ids = list(batch.id)\n",
    "            \n",
    "            yield (padded_context, padded_question, char_ctx, char_ques, label, ctx_text, answer_text, ids)\n",
    "            \n",
    "            "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create dataloaders\n",
    "\n",
    "train_dataset = SquadDataset(train_df,16)\n",
    "valid_dataset = SquadDataset(valid_df,16)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([16, 253])\n",
      "torch.Size([16, 16])\n",
      "torch.Size([16, 253, 16])\n",
      "torch.Size([16, 16, 16])\n",
      "torch.Size([16, 2])\n",
      "16\n",
      "16\n",
      "16\n"
     ]
    }
   ],
   "source": [
    "# looking at the shapes of various tensors returned by the loader\n",
    "\n",
    "a = next(iter(train_dataset))\n",
    "for i in range(len(a)):\n",
    "    try:\n",
    "        print(a[i].shape)\n",
    "    except AttributeError:\n",
    "        print(len(a[i]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_glove_dict():\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": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "glove_dict = get_glove_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_weights_matrix(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",
    "\n",
    "    return weights_matrix, words_found\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "91195\n"
     ]
    }
   ],
   "source": [
    "weights_matrix, words_found = create_weights_matrix(glove_dict)\n",
    "print(\"Words found in the GloVe vocab: \" ,words_found)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save the weight matrix for future loading.\n",
    "# This matrix is the nn.Embedding's weight matrix.\n",
    "\n",
    "np.save('qanetglove_vt.npy', weights_matrix)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Depthwise Separable Convolutions\n",
    "\n",
    "Depthwise separable convolutions serve the same purpose as normal convolutions with the only difference being that they are faster because they reduce the number of multiplication operations. This is done by breaking the convolution operation into two parts: depthwise convolution and pointwise convolution.\n",
    "> *We use depthwise separable convolutions rather than traditional ones, as we observe that it is memory efﬁcient and has better generalization. *\n",
    "\n",
    "Let's understand why depthwise convolutions are faster than traditional convolution.\n",
    "Traditional convolution can be visualized as,\n",
    "\n",
    "<img src=\"images/conv2d.PNG\" width=\"700\" height=\"700\"/>\n",
    "\n",
    "Let's count the number of multiplications in a traditional convolution operation.  \n",
    "The number of multiplications for a single convolution operation is the number of elements inside the kernel. This is $D_{K}$ X $D_{K}$ X $M$ = $D_{K}^{2}$ X $M$.\n",
    "To get the output feature map, we slide or convolve this kernel over the input. Given the output dimensions, we perform $D_{O}$ covolutions along the width and the height of the input image. Therefore, the number of multiplications per kernel are $D_{O}^{2}$ X $D_{K}^{2}$ X $M$.   \n",
    "These calculations are for a single kernel. In convolutional neural networks, we usually use multiple kernels. Each kernel is expected to extract a unique feature from the input. If we use $N$ such filters, then number of multiplications become \n",
    "$N$ X $D_{O}^{2}$ X $D_{K}^{2}$ X $M$.  \n",
    "\n",
    "### Depthwise convolution\n",
    "\n",
    "<img src=\"images/depthconv.PNG\" width=\"800\" height=\"900\"/>\n",
    "\n",
    "In depthwise convolution we perform convolution using kernels of dimension $D_{K}$ X $D_{K}$ X 1. Therefore the number of multiplications in a single convolution operation would be $D_{K}^{2}$ X $1$. If the output dimension is $D_{O}$, then the number of multiplications per kernel are $D_{K}^{2}$ X $D_{O}^{2}$. If there are $M$ input channels, we need to use $M$ such kernels, one kernel for each input channel to get the all the features. For $M$ kernels, we then get $D_{K}^{2}$ X $D_{O}^{2}$ X $M$ multiplications. \n",
    "\n",
    "### Pointwise convolution\n",
    "\n",
    "<img src=\"images/pointconv.PNG\" width=\"700\" height=\"700\"/>\n",
    "\n",
    "This part takes the output from depthwise convolution and performs convolution operation with a kernel of size 1 X 1 X $N$, where $N$ is the desired number of output features/channels. Here similarly,   \n",
    "Multiplications per 1 convolution operation = 1 X 1 X $M$  \n",
    "Multiplications per kernel = $D_{O}^{2}$ X $M$  \n",
    "For N output features = $N$ X $D_{O}^{2}$ X $M$\n",
    "  \n",
    "   \n",
    "Adding up the number of multiplications from both the phases, we get, \n",
    "\n",
    "$$ =\\  N\\ .\\ D_{O}^{2} \\ .\\ M \\ +\\ D_{K}^{2}\\ .\\ D_{O}^{2}\\ .\\ M $$\n",
    "$$ =\\  D_{O}^{2}\\ .\\ M (N + D_{K}^{2}) $$\n",
    "\n",
    "Comparing this with traditional convolutions, \n",
    "\n",
    "$$ =\\ \\frac {D_{O}^{2}\\ .\\ M\\ (N + D_{K}^{2})} {D_{O}^{2}\\ .\\  M\\ .\\ D_{K}^{2}\\ .\\ N}$$  \n",
    "\n",
    "$$ =\\  \\frac{1}{D_{K}^{2}}\\ +\\ \\frac{1}{N} $$\n",
    "\n",
    "This clearly shows that the number of computations in depthwise separable convolutions are lesser than traditional ones.\n",
    "In code, the depthwise phase of the convolution is done by assigning `groups` as `in_channels`. According to the documentation, \n",
    "\n",
    "> *At groups= `in_channels`, each `nput channel is convolved with its own set of filters, of size: $\\left\\lfloor\\frac{out\\_channels}{in\\_channels}\\right\\rfloor$*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DepthwiseSeparableConvolution(nn.Module):\n",
    "    \n",
    "    def __init__(self, in_channels, out_channels, kernel_size, dim=1):\n",
    "        \n",
    "        super().__init__()\n",
    "        self.dim = dim\n",
    "        if dim == 2:\n",
    "            \n",
    "            self.depthwise_conv = nn.Conv2d(in_channels=in_channels, out_channels=in_channels,\n",
    "                                        kernel_size=kernel_size, groups=in_channels, padding=kernel_size//2)\n",
    "        \n",
    "            self.pointwise_conv = nn.Conv2d(in_channels, out_channels, kernel_size=1, padding=0)\n",
    "        \n",
    "    \n",
    "        else:\n",
    "        \n",
    "            self.depthwise_conv = nn.Conv1d(in_channels=in_channels, out_channels=in_channels,\n",
    "                                            kernel_size=kernel_size, groups=in_channels, padding=kernel_size//2,\n",
    "                                            bias=False)\n",
    "\n",
    "            self.pointwise_conv = nn.Conv1d(in_channels, out_channels, kernel_size=1, padding=0, bias=True)\n",
    "\n",
    "    \n",
    "    def forward(self, x):\n",
    "        # x = [bs, seq_len, emb_dim]\n",
    "        if self.dim == 1:\n",
    "            x = x.transpose(1,2)\n",
    "            x = self.pointwise_conv(self.depthwise_conv(x))\n",
    "            x = x.transpose(1,2)\n",
    "        else:\n",
    "            x = self.pointwise_conv(self.depthwise_conv(x))\n",
    "        #print(\"DepthWiseConv output: \", x.shape)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Highway Networks\n",
    "\n",
    "Highway networks were originally introduced to ease the training of deep neural networks. While researchers had cracked the code for optimizing shallow neural networks, training *deep* networks was still a challenging task owing to problems such as vanishing gradients etc. Quoting the paper,\n",
    "\n",
    ">  *We present a novel architecture that enables the optimization of networks with virtually arbitrary depth. This is accomplished through the use of a learned gating mechanism for regulating information ﬂow which is inspired by Long Short Term Memory recurrent neural networks. Due to this gating mechanism, a neural network can have paths along which information can ﬂow across several layers without attenuation. We call such paths information highways, and such networks highway networks.* \n",
    "\n",
    "This paper takes the key idea of learned gating mechanism from LSTMs which process information internally through a sequence of learned gates. The purpose of this layer is to *learn* to pass relevant information from the input. A highway network is a series of feed-forward or linear layers with a gating mechanism. The gating is implemented by using a sigmoid function which decides what amount of information should be transformed and what should be passed as it is.   \n",
    "\n",
    "A plain feed-forward layer is associated with a linear transform $H$ parameterized by ($W_{H}, b_{H}$), such that for input $x$, the output $y$ is  \n",
    "\n",
    "$$ y = g(W_{H}.x + b_{H})$$\n",
    "where $g$ is a non-linear activation.  \n",
    "For highway networks, two additional linear transforms are defined viz. $T$ ($W_{T},b_{T}$) and $C$ ($W_{C}$,$b_{C}$).\n",
    "Then,    \n",
    "  \n",
    "$$ y = T(x) . H(x) + x . C(x) $$ \n",
    "> *We refer to T as the transform gate and C as the carry gate, since they express how much of the output is produced by\n",
    "transforming the input and carrying it, respectively. For simplicity, in this paper we set C = 1 − T. *\n",
    "\n",
    "$$ y = T(x) . H(x) + x . (1 - T(x)) $$  \n",
    "  \n",
    "$$ y = T(x) . g(W_{H}.x + b_{H}) + x . (1 - T(x)) $$  \n",
    "where $T(x)$ = $\\sigma$ ($W_{T}$ . $x$ + $b_{T}$) and $g$ is relu activation.  \n",
    "\n",
    "The input to this layer is the concatenation of word and character embeddings of each word. To implement this we use `nn.ModuleList` to add multiple linear layers. This is done for the gate layer as well as for a normal linear transform. In code the `flow_layer` is the same as linear transform $H$ discussed above and `gate_layer` is $T$. In the forward method we loop through each layer and compute the output according to the highway equation described above.   \n",
    "  \n",
    "The output of this layer for context is $X$ $\\epsilon$ $R^{\\ d \\ X \\ T}$ and for query is $Q$ $\\epsilon$ $R^{\\ d \\ X \\ J}$, where $d$ is hidden size of the LSTM, $T$ is the context length, $J$ is the query length.  \n",
    "\n",
    "The structure discussed so far is a recurring pattern in many NLP systems. Although this might be out of favor now with the advent of transformers and large pretrained language models, you will find this pattern in many NLP systems before transformers came into being. The idea behind this is that adding highway layers enables the network to make more efficient use of character embeddings. If a particular word is not found in the pretrained word vector vocabulary (OOV word), it will most likely be initialized with a zero vector. It then makes much more sense to look at the character embedding of that word rather than the word embedding. The soft gating mechanism in highway layers helps the model to achieve this. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "class HighwayLayer(nn.Module):\n",
    "    \n",
    "    def __init__(self, layer_dim, num_layers=2):\n",
    "    \n",
    "        super().__init__()\n",
    "        self.num_layers = num_layers\n",
    "        \n",
    "        self.flow_layers = nn.ModuleList([nn.Linear(layer_dim, layer_dim) for _ in range(num_layers)])\n",
    "        self.gate_layers = nn.ModuleList([nn.Linear(layer_dim, layer_dim) for _ in range(num_layers)])\n",
    "    \n",
    "    def forward(self, x):\n",
    "        #print(\"Highway input: \", x.shape)\n",
    "        for i in range(self.num_layers):\n",
    "            \n",
    "            flow = self.flow_layers[i](x)\n",
    "            gate = torch.sigmoid(self.gate_layers[i](x))\n",
    "            \n",
    "            x = gate * flow + (1 - gate) * x\n",
    "            \n",
    "        #print(\"Highway output: \", x.shape)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Embedding Layer\n",
    "\n",
    "This layer:\n",
    "* converts word-level tokens into a 300-dim pre-trained glove embedding vector \n",
    "* creates trainable character embeddings using 2-D convolutions\n",
    "* concatenates character and word embeddings and passes them through a highway network  \n",
    "\n",
    "The details of calculating character embeddings has been discussed in detail in the previous notebook. The only difference here is that instead of max-pooling, `torch.max` is used to get a fixed-size representation of each word.\n",
    "\n",
    "> *Each character is represented as a trainable vector of dimension p2 = 200,meaning each word can be viewed as the concatenation of the embedding vectors for each of its characters. The length of each word is either truncated or padded to 16. We take maximum value of each row of this matrix to get a ﬁxed-size vector representation of each word.* "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EmbeddingLayer(nn.Module):\n",
    "    \n",
    "    def __init__(self, char_vocab_dim, char_emb_dim, kernel_size, device):\n",
    "        \n",
    "        super().__init__()\n",
    "        \n",
    "        self.device = device\n",
    "        \n",
    "        self.char_embedding = nn.Embedding(char_vocab_dim, char_emb_dim)\n",
    "        \n",
    "        self.word_embedding = self.get_glove_word_embedding()\n",
    "        \n",
    "        self.conv2d = DepthwiseSeparableConvolution(char_emb_dim, char_emb_dim, kernel_size,dim=2)\n",
    "        \n",
    "        self.highway = HighwayLayer(self.word_emb_dim + char_emb_dim)\n",
    "    \n",
    "        \n",
    "    def get_glove_word_embedding(self):\n",
    "        \n",
    "        weights_matrix = np.load('qanetglove.npy')\n",
    "        num_embeddings, embedding_dim = weights_matrix.shape\n",
    "        self.word_emb_dim = embedding_dim\n",
    "        embedding = nn.Embedding.from_pretrained(torch.FloatTensor(weights_matrix).to(self.device),freeze=True)\n",
    "\n",
    "        return embedding\n",
    "    \n",
    "    def forward(self, x, x_char):\n",
    "        # x = [bs, seq_len]\n",
    "        # x_char = [bs, seq_len, word_len(=16)]\n",
    "        \n",
    "        word_emb = self.word_embedding(x)\n",
    "        # word_emb = [bs, seq_len, word_emb_dim]\n",
    "        \n",
    "        word_emb = F.dropout(word_emb,p=0.1)\n",
    "        \n",
    "        char_emb = self.char_embedding(x_char)\n",
    "        # char_embed = [bs, seq_len, word_len, char_emb_dim]\n",
    "        \n",
    "        char_emb = F.dropout(char_emb.permute(0,3,1,2), p=0.05)\n",
    "        # [bs, char_emb_dim, seq_len, word_len] == [N, Cin, Hin, Win]\n",
    "        \n",
    "        conv_out = F.relu(self.conv2d(char_emb))\n",
    "        # [bs, char_emb_dim, seq_len, word_len] \n",
    "        # the depthwise separable conv does not change the shape of the input\n",
    "        \n",
    "        char_emb, _ = torch.max(conv_out, dim=3)\n",
    "        # [bs, char_emb_dim, seq_len]\n",
    "        \n",
    "        char_emb = char_emb.permute(0,2,1)\n",
    "        # [bs, seq_len, char_emb_dim]\n",
    "        \n",
    "        concat_emb = torch.cat([char_emb, word_emb], dim=2)\n",
    "        # [bs, seq_len, char_emb_dim + word_emb_dim]\n",
    "        \n",
    "        emb = self.highway(concat_emb)\n",
    "        # [bs, seq_len, char_emb_dim + word_emb_dim]\n",
    "        \n",
    "        #print(\"Embedding output: \", emb.shape)\n",
    "        return emb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Multiheaded Self Attention\n",
    "\n",
    "### Idea of Linear Projections\n",
    "\n",
    "Consider a system of an online book store like kindle, which lets you rent, buy and read books on its platform. Such platforms usually have a recommendation system (recsys) in place that enables them to understand their users' taste and preferences over time. This helps them in making personalized recommendations to users and in turn improve their revenue.\n",
    "For simplicity, let's assume that there are 10,000 books available on the platform and the system maintain a simple binary vector of size 10,000 for each user. If a user has read a particular book, the position in the vector corresponding to the book's id is 1 and 0 otherwise. A *books-read* vector for a user looks like,\n",
    "$$ [1,0,0,1,1,0,0,0,0,1,1,...,1] $$ \n",
    "\n",
    "  \n",
    "Now assume a projection matrix of dimension 10,000 X 100. When we multiply any user's books vector, we get a new low dimensional vector of size 100. This vector is totally different from the previous one and now represents the user's taste or preferences in books. It basically represents a user-profile for the recommendation system. Calculating this user-taster vector for different users enables the application to find users with similar taste and recommend books that they *might* like simply based on what the other \"similar\" user has read.  \n",
    "The weights or values of this projection matrix can be thought as representing certain features or properties that a book might possess. It might capture various genres like science, philsophy, fantasy novels, etc.  \n",
    "The question that still remains however, is how do we get such a projection matrix in the first place that can transform a represenation from one vector space to another that is somehow related to the original vector but has an entirely different interpretation.  \n",
    "This is exactly what deep learning is about. Neural networks work as this *universal function approximators* that helps in learning such transformations. The weights of such projection matrices are learned via backpropagation. We also need a lot of training data to achieve this.\n",
    "\n",
    "### Self Attention \n",
    "\n",
    "Much of what will follow is heavily derived from Jay Alammar's famous blog post: The Illustrated Transformer. The intuition and visualizations can be directly converted into code and that's my main motive here. To understand the details, we'll first look at self attention using vectors at a granular level. We'll then show how actually these computations are made using matrices which directly correspond to the code. For convenience, we'll explain how self attention works in the transformer model. The input to the self attention layer is an embedding vector.   \n",
    "The central idea of attention is the same as discussed in the first notebook. Even here we'll calculate the measure of similarity between two representations, convert them into an attention distribution and take a weighted sum with the values. However, there are certain details involved that need to be addressed.  \n",
    "Following steps involved in calculating self-attention.\n",
    "1. The first step is to project the input into 3 different vector spaces: key space, query space and value space. These projections give us a key vector, a query vector and a value vector. The weights of these projection matrices are learnt via backpropagation during training. The projection matrices for key, query and value are $W^{K}$, $W^{Q}$, $W^{V}$ respectively. These projections are exactly what we discussed above. Their values depend a lot on the training procedure and the training data.  \n",
    "<img src=\"images/selfattn1.PNG\" width=\"700\" height=\"700\"/>\n",
    "\n",
    "2. The next step is to calculate attention scores. This is basically the part where we determine how similar are two input vectors and hence how much attention/focus needs to be paid on one vector while summarizing the other. \n",
    "\n",
    " >*The score determines how much focus to place on other parts of the input sentence as we encode a word at a certain position.*     \n",
    " \n",
    " There are different ways to determine this. In this paper, a dot product between the query and the key is used. Consider the phrase \"Thinking Machines\". For the word \"Thinking\", we need to calculate a score with each word in the sentence including \"Thinking\" itself. Therefore, score for first position would be,\n",
    "$$ q_{1}\\ .\\ k_{1} $$\n",
    "The result of this product represents the amount of attention we need to pay to \"Thinking\" itself while encoding \"Thinking\".\n",
    "The score for next position would be,\n",
    "$$ q_{1}\\ .\\ k_{2} $$\n",
    "which captures the importance of \"Machines\" while encoding \"Thinking\".  \n",
    "3. We then divide the scores calculated in the previous step by $\\sqrt d_{k}$, where $d_{k}$ is the dimension of key vectors. This scaling was done to ensure that the gradients are stable during training. Next, these scores are passed through a softmax function to get an attention distribution. This means that for a sentence of length $n$, if $\\alpha_{t}$ represents the score at $t$-th position, then\n",
    "$$ \\sum_{t=1}^{n} \\alpha_{t} = 1$$\n",
    "4. The last step is to multiply the softmax output with the value vector at respective position and sum these products up. In effect this computes a weighted sum. For a sentence of length n,\n",
    "$$ \\sum_{t=1}^{n} \\alpha_{t}\\ v_{t}$$\n",
    "\n",
    "All the steps explained above can be summarized as,   \n",
    "<img src=\"images/selfattn2.PNG\" width=\"600\" height=\"500\"/>\n",
    "\n",
    "### Multiheaded attention and Implementation\n",
    "\n",
    "The above steps are usually performed using matrices instead of vectors. This is also where we'll see how and why multihead attention is implemented.\n",
    "1. The first step is to calculate the query, key and value matrices by projecting them using trainable weights. In code, these weights correspond to linear layers. $W^{Q}$ corresponds to `fc_q`, $W^{K}$ to `fc_k` and $W^{V}$ to `fc_v`. Projecting these gives us $Q$, $K$ and $V$ as seen in code too. \n",
    "\n",
    "<img src=\"images/selfattn3.PNG\" width=\"600\" height=\"500\"/>\n",
    "Similar representations for value and key are also calculated. The dimensions of the above matrices will be explained below.\n",
    "2. Calculation of scores can be easily visualized as follows,\n",
    "<img src=\"images/selfattn4.PNG\" width=\"800\" height=\"1000\"/>  \n",
    "In code this is achieved by calculating the `energy` of $K$ and $Q$ using `torch.matmul`.\n",
    "3. The final step is to scale, take softmax of the scores and multiply the matrix by the value matrix.\n",
    "`scale` is calculated by taking the square root of `head_dim`. After scaling the `energy` tensor or the scores at different positions, we apply softmax to this tensor and multiply it with $V$ using `torch.matmul` once again. \n",
    "<img src=\"images/selfattn5.PNG\" width=\"600\" height=\"500\"/>\n",
    "\n",
    "\n",
    "In the original transformer model, the input embedding size is 512. Before projecting these embeddings, we split them into 8 parts which brings us to multihead attention. This paper uses 8 attention heads.  \n",
    " Multiheaded attention expands the model's ability to focus on different positions.\n",
    " > *It gives the attention layer mutiple \"representation subspaces.\"*\n",
    " \n",
    "These subspaces are nothing but different projection matrices. Instead of having just one projection matrix $W^{Q}$ for query, we'll have 8 projection matrices for query, key and value. Weights for each of these \"subspaces\" are learnt via backpropagation during training. An analogy for this can be the use of multiple convolutional filters to learn unique features from the image.  \n",
    "Therefore, now the dimension of key, query and value matrices would be 64 (512/8). In code, splitting weight matrices for multiple attention heads is done right after getting $K$, $Q$ and $V$. This is done by first calculating the `head_dimension` and then splitting the tensors using the `view` function.\n",
    " <img src=\"images/selfattn6.PNG\" width=\"600\" height=\"500\"/>\n",
    "The above image shows projection matrices for 2 attention heads. There are 8 such heads. This would give us 8 $Z$ matrices in the end. The output dimension of the self attention layer should be same as the input dimension. Hence, we need to recombine the results of all the attention heads before passing the output to the next layer. To combine them, in code, we simply use `view` to drop the head dimension and further make a projection using `fc_o` to ensure that the input dimension is same as the output dimension.  \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiheadAttentionLayer(nn.Module):\n",
    "    \n",
    "    def __init__(self, hid_dim, num_heads, device):\n",
    "        \n",
    "        super().__init__()\n",
    "        self.num_heads = num_heads\n",
    "        self.device = device\n",
    "        self.hid_dim = hid_dim\n",
    "        \n",
    "        self.head_dim = self.hid_dim // self.num_heads\n",
    "        \n",
    "        self.fc_q = nn.Linear(hid_dim, hid_dim)\n",
    "        \n",
    "        self.fc_k = nn.Linear(hid_dim, hid_dim)\n",
    "        \n",
    "        self.fc_v = nn.Linear(hid_dim, hid_dim)\n",
    "        \n",
    "        self.fc_o = nn.Linear(hid_dim, hid_dim)\n",
    "        \n",
    "        self.scale = torch.sqrt(torch.FloatTensor([self.head_dim])).to(device)\n",
    "        \n",
    "        \n",
    "    def forward(self, x, mask):\n",
    "        # x = [bs, len_x, hid_dim]\n",
    "        # mask = [bs, len_x]\n",
    "        \n",
    "        batch_size = x.shape[0]\n",
    "        \n",
    "        Q = self.fc_q(x)\n",
    "        K = self.fc_k(x)\n",
    "        V = self.fc_v(x)\n",
    "        # Q = K = V = [bs, len_x, hid_dim]\n",
    "        \n",
    "        Q = Q.view(batch_size, -1, self.num_heads, self.head_dim).permute(0,2,1,3)\n",
    "        K = K.view(batch_size, -1, self.num_heads, self.head_dim).permute(0,2,1,3)\n",
    "        V = V.view(batch_size, -1, self.num_heads, self.head_dim).permute(0,2,1,3)\n",
    "        # [bs, len_x, num_heads, head_dim ]  => [bs, num_heads, len_x, head_dim]\n",
    "        \n",
    "        K = K.permute(0,1,3,2)\n",
    "        # [bs, num_heads, head_dim, len_x]\n",
    "        \n",
    "        energy = torch.matmul(Q, K) / self.scale\n",
    "        # (bs, num_heads){[len_x, head_dim] * [head_dim, len_x]} => [bs, num_heads, len_x, len_x]\n",
    "        \n",
    "        mask = mask.unsqueeze(1).unsqueeze(2)\n",
    "        # [bs, 1, 1, len_x]\n",
    "        \n",
    "        #print(\"Mask: \", mask)\n",
    "        #print(\"Energy: \", energy)\n",
    "        \n",
    "        energy = energy.masked_fill(mask == 1, -1e10)\n",
    "        \n",
    "        #print(\"energy after masking: \", energy)\n",
    "        \n",
    "        alpha = torch.softmax(energy, dim=-1)\n",
    "        #  [bs, num_heads, len_x, len_x]\n",
    "        \n",
    "        #print(\"energy after smax: \", alpha)\n",
    "        alpha = F.dropout(alpha, p=0.1)\n",
    "        \n",
    "        a = torch.matmul(alpha, V)\n",
    "        # [bs, num_heads, len_x, head_dim]\n",
    "        \n",
    "        a = a.permute(0,2,1,3)\n",
    "        # [bs, len_x, num_heads, hid_dim]\n",
    "        \n",
    "        a = a.contiguous().view(batch_size, -1, self.hid_dim)\n",
    "        # [bs, len_x, hid_dim]\n",
    "        \n",
    "        a = self.fc_o(a)\n",
    "        # [bs, len_x, hid_dim]\n",
    "        \n",
    "        #print(\"Multihead output: \", a.shape)\n",
    "        return a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.autograd import Variable\n",
    "import math"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Positional Embedding\n",
    "\n",
    "The model so far does not have any idea about the positioning of words in a sentence. In previous models, this was taken care of because we were using RNNs or LSTMs at some stage to encode this information. RNNs process input in a sequential order and maintain hidden states for each position in the input sequence. However, here we need to come up with a method to inject the positional information of tokens into the model.  \n",
    "One simple method of doing this is to assign a single number to each token between $[0, 1]$, where first word starts with 0 and the last word corresponds to 1. This solution presents some problems. For different sentence lengths, we'll have different intervals over which tokens are distributed. We would not have a consistent meaning of a particular position across all inputs(of varying lengths).   \n",
    "Another method is to use learned position embeddings. This is used in BERT, where, the positional embedding a lookup table of size $[512, 768]$ where 512 is the maximum sequence length that BERT can process. This lookup matrix is randomly intialized and trained along with the model.    \n",
    "Here however, the authors have used another method of encoding position which is same as that proposed in the original transformers paper. The positional embedding can be defined as,\n",
    "<img src=\"images/posemb.PNG\" width=\"500\" height=\"400\"/>\n",
    "\n",
    "where $pos$ is the position, $i$ is the dimension of embedding, and $d_{model}$ is the model dimension.  \n",
    "These embeddings are simply added to the word embeddings of the tokens at their respective positions. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PositionEncoder(nn.Module):\n",
    "    \n",
    "    def __init__(self, model_dim, device, max_length=400):\n",
    "        \n",
    "        super().__init__()\n",
    "        \n",
    "        self.device = device\n",
    "        \n",
    "        self.model_dim = model_dim\n",
    "        \n",
    "        pos_encoding = torch.zeros(max_length, model_dim)\n",
    "        \n",
    "        for pos in range(max_length):\n",
    "            \n",
    "            for i in range(0, model_dim, 2):\n",
    "                \n",
    "                pos_encoding[pos, i] = math.sin(pos / (10000 ** ((2*i)/model_dim)))\n",
    "                pos_encoding[pos, i+1] = math.cos(pos / (10000 ** ((2*(i+1))/model_dim)))\n",
    "            \n",
    "        \n",
    "        pos_encoding = pos_encoding.unsqueeze(0).to(device)\n",
    "        self.register_buffer('pos_encoding', pos_encoding)\n",
    "        \n",
    "    \n",
    "    def forward(self, x):\n",
    "        #print(\"PE shape: \", self.pos_encoding.shape)\n",
    "        #print(\"PE input: \", x.shape)\n",
    "        x = x + Variable(self.pos_encoding[:, :x.shape[1]], requires_grad=False)\n",
    "        #print(\"PE output: \", x.shape)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Encoder Block\n",
    "\n",
    "This layer brings together all the components discussed so far. \n",
    "<img src=\"images/encoderblock.PNG\" width=\"250\" height=\"50\"/>\n",
    "\n",
    "> *We use the same Encoder Block throughout the model, only varying the number of convolutional layers for each block. We use layer norm and residual connection between every layer in the Encoder Block.*\n",
    "\n",
    "The following steps are performed by this layer:  \n",
    "\n",
    "* A positional embedding is injected into the input.\n",
    "* This is then passed through a series of convolutional layers. The number of these layers depend upon the layer of which these encoder blocks are a part of. For embedding encoder layer, this number is 4 and for model encoder layer it is 2. The layers of convolution are defined using `nn.Modulelist`. \n",
    "* The output of this is then passed to a multiheaded self attention layer and finally to a feedforward network which is simply a linear layer.\n",
    "* As can be seen in the figure above, the model involves residual connections, layer normalizations and dropouts too. These too are implemented appropriately. An easy way to understand the residual connections in code would be draw 2-3 iterations of the lower block (that involves convolution) and ensure that everything matches.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "class EncoderBlock(nn.Module):\n",
    "    \n",
    "    def __init__(self, model_dim, num_heads, num_conv_layers, kernel_size, device):\n",
    "        \n",
    "        super().__init__()\n",
    "        \n",
    "        self.num_conv_layers = num_conv_layers\n",
    "        \n",
    "        self.conv_layers = nn.ModuleList([DepthwiseSeparableConvolution(model_dim, model_dim, kernel_size)\n",
    "                                          for _ in range(num_conv_layers)])\n",
    "        \n",
    "        self.multihead_self_attn = MultiheadAttentionLayer(model_dim, num_heads, device)\n",
    "        \n",
    "        self.position_encoder = PositionEncoder(model_dim, device)\n",
    "        \n",
    "        self.pos_norm = nn.LayerNorm(model_dim)\n",
    "        \n",
    "        self.conv_norm = nn.ModuleList([nn.LayerNorm(model_dim) for _ in range(self.num_conv_layers)])\n",
    "        \n",
    "        self.feedfwd_norm = nn.LayerNorm(model_dim)\n",
    "        \n",
    "        self.feed_fwd = nn.Linear(model_dim, model_dim)\n",
    "        \n",
    "    def forward(self, x, mask):\n",
    "        # x = [bs, len_x, model_dim]\n",
    "        # mask = [bs, len_x]\n",
    "        \n",
    "        out = self.position_encoder(x)\n",
    "        # [bs, len_x, model_dim]\n",
    "        \n",
    "        res = out\n",
    "        \n",
    "        out = self.pos_norm(out)\n",
    "        # [bs, len_x, model_dim]\n",
    "        \n",
    "        for i, conv_layer in enumerate(self.conv_layers):\n",
    "            \n",
    "            out = F.relu(conv_layer(out))\n",
    "            out = out + res\n",
    "            if (i+1) % 2 == 0:\n",
    "                out = F.dropout(out, p=0.1)\n",
    "            res = out\n",
    "            out = self.conv_norm[i](out)\n",
    "        \n",
    "        \n",
    "        out = self.multihead_self_attn(out, mask)\n",
    "        # [bs, len_x, model_dim]\n",
    "        \n",
    "        out = F.dropout(out + res, p=0.1)\n",
    "        \n",
    "        res = out\n",
    "        \n",
    "        out = self.feedfwd_norm(out)\n",
    "        \n",
    "        out = F.relu(self.feed_fwd(out))\n",
    "        # [bs, len_x, model_dim]\n",
    "            \n",
    "        out = F.dropout(out + res, p=0.1)\n",
    "        # [bs, len_x, model_dim]\n",
    "        #print(\"Encoder block output: \", out.shape)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Context-Query Attention Layer\n",
    "\n",
    "This layer is very similar to the attention flow layer in BIDAF. It calculates attention in two directions. Context-query attention tells us what query words are the most relevant to each context word.   \n",
    "Let $C$ and $Q$ represent the encoded context and query respectively. Given that the context length is $n$ and query length is $m$, a similarity matrix is calculated first. The similarity matrix captures the similarity between each pair of context and query words. It is denoted by $S$ and is a $n$-by-$m$ matrix. The similarity matrix is calculated as,\n",
    "$$ S = f\\ (Q,\\ C)$$\n",
    "where $f$ is a trilinear similarity function defined as,\n",
    "$$ f(q,c) = W_{0}\\ [q\\ ;\\ c\\ ;\\ q \\odot c] $$,\n",
    "where $W_{0}$ is trainable variable, $;$ denotes concatenation and $\\odot$ denotes element wise multiplication.  \n",
    "Context-to-Query attention can then be calculated as,\n",
    "$$ A = \\overline S\\ .\\ Q^{T} $$,\n",
    "where $\\overline S$ is obtained by normalizing each row of $S$ using softmax. The computations so far are exactly similar to those in BIDAF. You can refer to the previous notebook for a more detailed explanation.  \n",
    "\n",
    "> *Most high performing models additionally use some form of query-to-context attention, such as BiDaF and DCN. Empirically, we ﬁnd that, the DCN attention can provide a little beneﬁt over simply applying context-to-query attention, so we adopt this strategy.*\n",
    "\n",
    "Query-to-Context attention is calculated as,\n",
    "$$B = \\overline S\\ .\\ \\overline{\\overline S}^{T}\\ .\\ C^{T}$$,\n",
    "where $\\overline{\\overline S}^{T}$ is the column-normalized matrix of $S$ by softmax function.  \n",
    "\n",
    "The implementation is fairly straightforward and is just about multiplying the said tensors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ContextQueryAttentionLayer(nn.Module):\n",
    "    \n",
    "    def __init__(self, model_dim):\n",
    "        \n",
    "        super().__init__() \n",
    "        \n",
    "        self.W0 = nn.Linear(3*model_dim, 1, bias=False)\n",
    "        \n",
    "    def forward(self, C, Q, c_mask, q_mask):\n",
    "        # C = [bs, ctx_len, model_dim]\n",
    "        # Q = [bs, qtn_len, model_dim]\n",
    "        # c_mask = [bs, ctx_len]\n",
    "        # q_mask = [bs, qtn_len]\n",
    "        \n",
    "        c_mask = c_mask.unsqueeze(2)\n",
    "        # [bs, ctx_len, 1]\n",
    "        \n",
    "        q_mask = q_mask.unsqueeze(1)\n",
    "        # [bs, 1, qtn_len]\n",
    "        \n",
    "        ctx_len = C.shape[1]\n",
    "        qtn_len = Q.shape[1]\n",
    "        \n",
    "        C_ = C.unsqueeze(2).repeat(1,1,qtn_len,1)\n",
    "        # [bs, ctx_len, qtn_len, model_dim] \n",
    "        \n",
    "        Q_ = Q.unsqueeze(1).repeat(1,ctx_len,1,1)\n",
    "        # [bs, ctx_len, qtn_len, model_dim]\n",
    "        \n",
    "        C_elemwise_Q = torch.mul(C_, Q_)\n",
    "        # [bs, ctx_len, qtn_len, model_dim]\n",
    "        \n",
    "        S = torch.cat([C_, Q_, C_elemwise_Q], dim=3)\n",
    "        # [bs, ctx_len, qtn_len, model_dim*3]\n",
    "        \n",
    "        S = self.W0(S).squeeze()\n",
    "        #print(\"Simi matrix: \", S.shape)\n",
    "        # [bs, ctx_len, qtn_len, 1] => # [bs, ctx_len, qtn_len]\n",
    "        \n",
    "        S_row = S.masked_fill(q_mask==1, -1e10)\n",
    "        S_row = F.softmax(S_row, dim=2)\n",
    "        \n",
    "        S_col = S.masked_fill(c_mask==1, -1e10)\n",
    "        S_col = F.softmax(S_col, dim=1)\n",
    "        \n",
    "        A = torch.bmm(S_row, Q)\n",
    "        # (bs)[ctx_len, qtn_len] X [qtn_len, model_dim] => [bs, ctx_len, model_dim]\n",
    "        \n",
    "        B = torch.bmm(torch.bmm(S_row,S_col.transpose(1,2)), C)\n",
    "        # [ctx_len, qtn_len] X [qtn_len, ctx_len] => [bs, ctx_len, ctx_len]\n",
    "        # [ctx_len, ctx_len] X [ctx_len, model_dim ] => [bs, ctx_len, model_dim]\n",
    "        \n",
    "        model_out = torch.cat([C, A, torch.mul(C,A), torch.mul(C,B)], dim=2)\n",
    "        # [bs, ctx_len, model_dim*4]\n",
    "        \n",
    "        #print(\"C2Q output: \", model_out.shape)\n",
    "        return F.dropout(model_out, p=0.1)\n",
    "        \n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Output Layer\n",
    "\n",
    "The output layer is tasked with predicting the start and end indices of the answer from the context. The input to this layer\n",
    "$M_{1}$, $M_{2}$ and $M_{3}$ are the outputs of 3 model encoders(explained below), from bottom to top. The start index $p_{1}$ is then calculated as,  \n",
    "\n",
    "$$ p_{1} = softmax\\ (\\ W_{1}\\ [M_{1}\\ ;\\ M_{2}])$$\n",
    "and end as,\n",
    "$$ p_{2} = softmax\\ (\\ W_{2}\\ [M_{1}\\ ;\\ M_{3}])$$\n",
    "\n",
    "where $W_{1}$ and $W_{2}$ are trainable variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "class OutputLayer(nn.Module):\n",
    "    \n",
    "    def __init__(self, model_dim):\n",
    "        \n",
    "        super().__init__()\n",
    "        \n",
    "        self.W1 = nn.Linear(2*model_dim, 1, bias=False)\n",
    "        \n",
    "        self.W2 = nn.Linear(2*model_dim, 1, bias=False)\n",
    "        \n",
    "        \n",
    "    def forward(self, M1, M2, M3, c_mask):\n",
    "        \n",
    "        start = torch.cat([M1,M2], dim=2)\n",
    "        \n",
    "        start = self.W1(start).squeeze()\n",
    "        \n",
    "        p1 = start.masked_fill(c_mask==1, -1e10)\n",
    "        \n",
    "        #p1 = F.log_softmax(start.masked_fill(c_mask==1, -1e10), dim=1)\n",
    "        \n",
    "        end = torch.cat([M1, M3], dim=2)\n",
    "        \n",
    "        end = self.W2(end).squeeze()\n",
    "        \n",
    "        p2 = end.masked_fill(c_mask==1, -1e10)\n",
    "        \n",
    "        #p2 = F.log_softmax(end.masked_fill(c_mask==1, -1e10), dim=1)\n",
    "        \n",
    "        #print(\"preds: \", [p1.shape,p2.shape])\n",
    "        return p1, p2\n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## QANet\n",
    "\n",
    "This module wraps up everything. It brings together all the components together that we've seen so far. \n",
    "<img src=\"images/qanet.PNG\" width=\"500\" height=\"600\"/>\n",
    "\n",
    "Going up the flowchart above, the following module does the following end-to-end:\n",
    "* The inputs to the `forward` method are word-level and character-level tokens for both the context and the query. These tokens are passed to the embedding layer.  \n",
    "\n",
    " > *The word embedding is ﬁxed during training and initialized from the p1 = 300 dimensional pre-trained GloVe word vectors, which are ﬁxed during training.*  \n",
    "\n",
    " > *The character embedding is obtained as follows: Each character is represented as a trainable vector of dimension p2 = 200, meaning each word can be viewed as the concatenation of the embedding vectors for each of its characters.*   \n",
    "\n",
    " For each word the concatenation of these two embeddings is passed on to a 2-layer highway network. Highway network does not affect the shape of the input. Hence the output shape from the `EmbeddingLayer` defined above would be `[bs, ctx_len, word_emb_dim + char_emb_dim]` = `[batch_size, ctx_len, 500]`. This is then supposed to be passed to `embedding_encoder` or the Embedding Encoder Layer. This layer however requires the input dimension to be 128 which is the `model_dim`\n",
    " and not 500. As clearly mentioned in the paper,    \n",
    " > *Note that the input of this layer is a vector of dimension p1 + p2 = 500 for each individual word, which is immediately mapped to d = 128 by a one-dimensional convolution. The output of this layer is a also of dimension d = 128.*   \n",
    " \n",
    " We therefore map the output of embedding to 128 in code using `ctx_resizer` and `qtn_resizer`.  \n",
    "\n",
    "* The resized tensors are then passed on to the *Embedding Encoding Layer* which is a single encoder block with 4 conv layers. 8 attention heads are used in the self-attention module which is the same for all the encoder blocks in the model.\n",
    "  \n",
    "* The output of previous layer is then passed on to the *Contex-Query Attention Layer*.  The output dimension of this layer is `4 * model_dim`. This is again resized using `c2q_resizer` to have a dimension of `model_dim`. \n",
    "* Next the encoded representation so far is passed on to the *Model Encoder Layer*. This layer comprises of 7 blocks of encoder, with each block having 2 convolutional layers. \n",
    " > *We share weights between each of the 3 repetitions of the model encoder.*\n",
    " This can be seen in code while calculating $M_{1}$, $M_{2}$ and $M_{3}$.\n",
    "* Finally the shared-weight matrices are passed to the output layer which predicts the start and end index of the answer.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "class QANet(nn.Module):\n",
    "    \n",
    "    def __init__(self, char_vocab_dim, char_emb_dim, word_emb_dim, kernel_size, model_dim, num_heads, device):\n",
    "        \n",
    "        super().__init__()\n",
    "        \n",
    "        self.embedding = EmbeddingLayer(char_vocab_dim, char_emb_dim, kernel_size, device)\n",
    "        \n",
    "        self.ctx_resizer = DepthwiseSeparableConvolution(char_emb_dim+word_emb_dim, model_dim, 5)\n",
    "        \n",
    "        self.qtn_resizer = DepthwiseSeparableConvolution(char_emb_dim+word_emb_dim, model_dim, 5)\n",
    "        \n",
    "        self.embedding_encoder = EncoderBlock(model_dim, num_heads, 4, 5, device)\n",
    "        \n",
    "        self.c2q_attention = ContextQueryAttentionLayer(model_dim)\n",
    "        \n",
    "        self.c2q_resizer = DepthwiseSeparableConvolution(model_dim*4, model_dim, 5)\n",
    "        \n",
    "        self.model_encoder_layers = nn.ModuleList([EncoderBlock(model_dim, num_heads, 2, 5, device)\n",
    "                                                   for _ in range(7)])\n",
    "        \n",
    "        self.output = OutputLayer(model_dim)\n",
    "        \n",
    "        self.device=device\n",
    "    \n",
    "    def forward(self, ctx, qtn, ctx_char, qtn_char):\n",
    "        \n",
    "        c_mask = torch.eq(ctx, 1).float().to(self.device)\n",
    "        q_mask = torch.eq(qtn, 1).float().to(self.device)\n",
    "        \n",
    "        ctx_emb = self.embedding(ctx, ctx_char)\n",
    "        # [bs, ctx_len, ch_emb_dim + word_emb_dim]\n",
    "        \n",
    "        ctx_emb = self.ctx_resizer(ctx_emb)\n",
    "        #  [bs, ctx_len, model_dim]\n",
    "        \n",
    "        qtn_emb = self.embedding(qtn, qtn_char)\n",
    "        # [bs, ctx_len, ch_emb_dim + word_emb_dim]\n",
    "        \n",
    "        qtn_emb = self.qtn_resizer(qtn_emb)\n",
    "        # [bs, qtn_len, model_dim]\n",
    "        \n",
    "        C = self.embedding_encoder(ctx_emb, c_mask)\n",
    "        # [bs, ctx_len, model_dim]\n",
    "        \n",
    "        Q = self.embedding_encoder(qtn_emb, q_mask)\n",
    "        # [bs, qtn_len, model_dim]\n",
    "            \n",
    "        C2Q = self.c2q_attention(C, Q, c_mask, q_mask)\n",
    "        # [bs, ctx_len, model_dim*4]\n",
    "        \n",
    "        M1 = self.c2q_resizer(C2Q)\n",
    "        # [bs, ctx_len, model_dim]\n",
    "    \n",
    "        for layer in self.model_encoder_layers:\n",
    "            M1 = layer(M1, c_mask)\n",
    "        \n",
    "        M2 = M1\n",
    "        # [bs, ctx_len, model_dim]  \n",
    "        \n",
    "        for layer in self.model_encoder_layers:\n",
    "            M2 = layer(M2, c_mask)\n",
    "        \n",
    "        M3 = M2\n",
    "        # [bs, ctx_len, model_dim]\n",
    "        \n",
    "        for layer in self.model_encoder_layers:\n",
    "            M3 = layer(M3, c_mask)\n",
    "            \n",
    "        p1, p2 = self.output(M1, M2, M3, c_mask)\n",
    "        \n",
    "        return p1, p2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "CHAR_VOCAB_DIM = len(char2idx)\n",
    "CHAR_EMB_DIM = 200\n",
    "WORD_EMB_DIM = 300\n",
    "KERNEL_SIZE = 5\n",
    "MODEL_DIM = 128\n",
    "NUM_ATTENTION_HEADS = 8\n",
    "device = torch.device('cuda')\n",
    "\n",
    "\n",
    "model = QANet(CHAR_VOCAB_DIM,\n",
    "              CHAR_EMB_DIM, \n",
    "              WORD_EMB_DIM,\n",
    "              KERNEL_SIZE,\n",
    "              MODEL_DIM,\n",
    "              NUM_ATTENTION_HEADS,\n",
    "              device).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The model has 2,263,312 trainable parameters\n"
     ]
    }
   ],
   "source": [
    "def count_parameters(model):\n",
    "    return sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "print(f'The model has {count_parameters(model):,} trainable parameters')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> *We use the ADAM optimizer (Kingma & Ba, 2014) with β1 = 0.8,β2 = 0.999, $\\epsilon$ = 10−7. We use a learning rate warm-up scheme with an inverse exponential increase from 0.0 to 0.001 in the ﬁrst 1000 steps, and then maintain a constant learning rate for the remainder of training.*\n",
    "\n",
    "Note: I have not used learning-rate warm up scheme to keep things simple for initial training. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch.optim as optim\n",
    "optimizer = optim.Adam(model.parameters(), betas=(0.8,0.999), eps=10e-7, weight_decay=3*10e-7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(model, train_dataset):\n",
    "    print(\"Starting training ........\")\n",
    "   \n",
    "\n",
    "    train_loss = 0.\n",
    "    batch_count = 0\n",
    "\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, char_ctx, char_ques, label, ctx_text, ans, ids = batch\n",
    "        \n",
    "        # place data on GPU\n",
    "        context, question, char_ctx, char_ques, label = context.to(device), question.to(device),\\\n",
    "                                    char_ctx.to(device), char_ques.to(device), label.to(device)\n",
    "        \n",
    "        # forward pass, get predictions\n",
    "        preds = model(context, question, char_ctx, char_ques)\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\n",
    "        loss.backward()\n",
    "        \n",
    "        # update the gradients\n",
    "        optimizer.step()\n",
    "\n",
    "        # zero the gradients so that they do not accumulate\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",
    "    print(\"Starting validation .........\")\n",
    "   \n",
    "    valid_loss = 0.\n",
    "\n",
    "    batch_count = 0\n",
    "    \n",
    "    f1, em = 0., 0.\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, char_ctx, char_ques, label, ctx_text, ans, ids = batch\n",
    "\n",
    "        context, question, char_ctx, char_ques, label = context.to(device), question.to(device),\\\n",
    "                                    char_ctx.to(device), char_ques.to(device), label.to(device)\n",
    "\n",
    "        with torch.no_grad():\n",
    "\n",
    "            preds = model(context, question, char_ctx, char_ques)\n",
    "\n",
    "            p1, p2 = preds\n",
    "\n",
    "            y1, y2 = label[:,0], label[:,1]\n",
    "\n",
    "            loss = F.nll_loss(p1, y1) + F.nll_loss(p2, y2)\n",
    "\n",
    "            valid_loss += loss.item()\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",
    "            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",
    "           \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",
    "    em, f1 = evaluate(predictions)\n",
    "    return valid_loss/len(valid_dataset), em, f1           \n",
    "  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "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": 51,
   "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",
    "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": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "train_losses = []\n",
    "valid_losses = []\n",
    "ems = []\n",
    "f1s = []\n",
    "epochs = 3\n",
    "for epoch in range(epochs):\n",
    "    print(f\"Epoch {epoch+1}\")\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. The QANet paper: https://arxiv.org/abs/1804.09541\n",
    "    2. Attention is All You Need https://papers.nips.cc/paper/7181-attention-is-all-you-need.pdf\n",
    "    3. Convolutional Neural Networks for Sentence Classification: https://arxiv.org/abs/1408.5882\n",
    "    4. Highway Networks: https://arxiv.org/abs/1505.00387\n",
    "* Other helpful links:\n",
    "    1. https://lilianweng.github.io/lil-log/2018/06/24/attention-attention.html\n",
    "    2. The Illustrated Transformer:http://jalammar.github.io/illustrated-transformer/. This is an excellent piece of writing with amazing easy-to-understand visualizations. Must read.\n",
    "    3. https://mccormickml.com/2019/11/11/bert-research-ep-1-key-concepts-and-sources/. Chris McCormick's BERT research series is another great resource to learn about self attention and various other details about BERT. He has a blog as well as youtube video series on the same.\n",
    "    4. https://kazemnejad.com/blog/transformer_architecture_positional_encoding/\n",
    "    5. https://nlp.seas.harvard.edu/2018/04/03/attention.html. The annotated Transformer.\n",
    "    6. https://nlp.seas.harvard.edu/slides/aaai16.pdf. A great resource for character embeddings.\n",
    "    7. https://www.youtube.com/watch?v=T7o3xvJLuHk. Easy explanation of depthwise separable convolutions.\n",
    "    8. https://towardsdatascience.com/a-basic-introduction-to-separable-convolutions-b99ec3102728. Another amazing blog for depthwise separable convolutions.\n",
    "    9. https://github.com/bentrevett/pytorch-seq2seq. A great series of notebooks on Machine Translation using PyTorch.  \n",
    "Some of the repositories below might be out of date. \n",
    "    10. https://github.com/BangLiu/QANet-PyTorch\n",
    "    11. https://github.com/NLPLearn/QANet\n",
    "    12. https://github.com/setoidz/QANet-pytorch\n",
    "    13. https://github.com/hackiey/QAnet-pytorch/tree/master/qanet"
   ]
  },
  {
   "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": 2
}