NaturalQuestionQA_with_Seq2Seq.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lxMY2HmHn0Ka"
   },
   "source": [
    "Natural Question Dataset:\n",
    "Official Link:\n",
    "This dataset contains really long context/passages and a single question based on the passage. The answer to the questions are subparts of the passage given in the form of <start_token> and <end_token>.\n",
    "What makes this dataset really challenging is the passage and answer length.\n",
    "\n",
    "I chose this dataset because it will be interesting to work with transformers with such long sequnces. When I tried this dataset with a seq2seq I kept getting OOM. The reason being exterme lenghts of contexts and their answers. The context are entire wikipedia pages in HTML format so there length is typically over 2000.\n",
    "\n",
    "I started experimenting with original dataset. I kept watering down the task. First reducing the number of examples, then answer length, batch_size but I kept getting OOM. Finally I restricted context lengths to 1000.\n",
    "\n",
    "The problem was in the sequence length of context. If I restrict them to 1000, the model has good enough memory to work with. No of samples has not much effect. And batch size does. I will remove this 1000 length restriction with transformers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "k7_hoAmMLduI"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import json\n",
    "from torchtext import data\n",
    "from itertools import chain\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import time\n",
    "from torch.nn import Embedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "id": "dgMi39JRLvQY"
   },
   "outputs": [],
   "source": [
    "%%bash\n",
    "file=/content/train.jsonl\n",
    "if [ ! -f \"$file\" ]; then\n",
    "  # Since the dataset itself is huge, we will make train and test set from the original train file itself\n",
    "  # A simple wget wont be able to get the natural questions dataset present at https://ai.google.com/research/NaturalQuestions/download\n",
    "  # Use the advice given in https://www.kaggle.com/c/deepfake-detection-challenge/discussion/121194 \n",
    "  # the download link is https://storage.cloud.google.com/natural_questions/v1.0-simplified/simplified-nq-train.jsonl.gz \n",
    "  \n",
    "  # put curl command here\n",
    "\n",
    "  zcat /content/train.jsonl.gz > /content/train.jsonl\n",
    "  \n",
    "fi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "id": "fe9fVmeXsoS2"
   },
   "outputs": [],
   "source": [
    "# In the beginning restrciting total number of examples to 10k. \n",
    "total_limit = 10000\n",
    "\n",
    "file = \"/content/train.jsonl\"\n",
    "f = open(file)\n",
    "\n",
    "total = 0\n",
    "examples = []\n",
    "for i,line in enumerate(f):\n",
    "    if total >total_limit:\n",
    "        break\n",
    "    ak = json.loads(line)\n",
    "    context = ak['document_text']\n",
    "    question = ak['question_text']\n",
    "    start = ak['annotations'][0]['long_answer']['start_token']\n",
    "    end = ak['annotations'][0]['long_answer']['end_token']\n",
    "    try:\n",
    "        assert start < end and end-start > 200 and len(context.split(\" \")) <1000\n",
    "    except AssertionError:\n",
    "        continue\n",
    "    answer = \" \".join(context.split(\" \")[start:end])\n",
    "    examples.append([context,question,answer])\n",
    "    total += 1\n",
    "f.close()\n",
    "\n",
    "# Will do the architecture with seperate encoders for context and question and a decoder for answer\n",
    "context = data.Field(sequential=True, tokenize='spacy', init_token='<sos>', eos_token='<eos>')\n",
    "question = data.Field(sequential=True, tokenize='spacy', init_token='<sos>', eos_token='<eos>')\n",
    "answer = data.Field(sequential=True, tokenize='spacy', init_token='<sos>', eos_token='<eos>')\n",
    "\n",
    "fields = [('context', context), ('question', question), ('answer', answer)]\n",
    "\n",
    "Examples = [data.Example.fromlist([i[0], i[1], i[2]], fields) for i in examples]\n",
    "Dataset = data.Dataset(Examples, fields)\n",
    "\n",
    "train_dataset,valid_dataset = Dataset.split(split_ratio=[0.85,0.15])\n",
    "\n",
    "context.build_vocab(train_dataset,min_freq=2,max_size = 20000,vectors = \"glove.6B.100d\", \n",
    "                 unk_init = torch.Tensor.normal_)\n",
    "question.build_vocab(train_dataset,min_freq=2,max_size = 5000,vectors = \"glove.6B.100d\", \n",
    "                 unk_init = torch.Tensor.normal_)\n",
    "answer.build_vocab(train_dataset,min_freq=2,max_size = 10000,vectors = \"glove.6B.100d\", \n",
    "                 unk_init = torch.Tensor.normal_)\n",
    "\n",
    "\n",
    "BATCH_SIZE = 32\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "train_iterator, valid_iterator = data.BucketIterator.splits((train_dataset, valid_dataset), batch_size=32,\n",
    "                                                            sort_key=lambda x: len(x.context),\n",
    "                                                            sort_within_batch=True,device=device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "F6NRhae8tmy_",
    "outputId": "c414cbf4-f98c-423a-b1ed-5948f149300e"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<torchtext.data.example.Example object at 0x7f84afaf1048>\n",
      "247\n"
     ]
    }
   ],
   "source": [
    "print(Examples[0])\n",
    "print(len(Examples))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "id": "Os7EnyUltw2_"
   },
   "outputs": [],
   "source": [
    "class Encoder(nn.Module):\n",
    "  def __init__(self,input_dim,emb_dim,hid_dim,n_layers,dropout,bidirectional):\n",
    "    super().__init__()\n",
    "    self.hid_dim = hid_dim\n",
    "    \n",
    "    self.embedding = nn.Embedding(input_dim,emb_dim)\n",
    "\n",
    "    self.rnn = nn.LSTM(input_size=emb_dim,hidden_size = hid_dim,num_layers= n_layers,dropout= dropout,bidirectional = bidirectional)\n",
    "\n",
    "    self.dropout = nn.Dropout(dropout)\n",
    "  \n",
    "  def forward(self,input,hidden=None,cell_state=None):\n",
    "    \n",
    "    embedded = self.dropout(self.embedding(input))\n",
    "\n",
    "    if not hidden == None:\n",
    "      outputs, (hidden,cell_state) = self.rnn(embedded,(hidden,cell_state))\n",
    "    else:\n",
    "      outputs, (hidden,cell_state) = self.rnn(embedded)\n",
    "\n",
    "    return hidden,cell_state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "id": "G2oFrolIuuMF"
   },
   "outputs": [],
   "source": [
    "class Decoder(nn.Module):\n",
    "  def __init__(self,output_dim,emd_dim,hid_dim,n_layers,bidirectional,dropout):\n",
    "    super().__init__()\n",
    "\n",
    "    self.embedded = nn.Embedding(output_dim,emb_dim)\n",
    "\n",
    "    self.rnn = nn.LSTM(input_size=emb_dim,hidden_size=hid_dim,num_layers=n_layers,bidirectional=bidirectional,dropout=dropout)\n",
    "\n",
    "    self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    no_of_directions = 2 if bidirectional else 1\n",
    "\n",
    "    self.fc_out = nn.Linear(no_of_directions*hid_dim,output_dim)\n",
    "\n",
    "  def forward(self,input,hidden,cell_state):\n",
    "    input = input.unsqueeze(0)\n",
    "\n",
    "    input = self.dropout(self.embedded(input))\n",
    "\n",
    "    output, (hidden,cell_state) = self.rnn(input,(hidden,cell_state))\n",
    "\n",
    "    output = output.squeeze(0)\n",
    "    \n",
    "    prediction = self.fc_out(output)\n",
    "\n",
    "    return prediction, hidden , cell_state\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "id": "pnp6p4uTxJFE"
   },
   "outputs": [],
   "source": [
    "import random\n",
    "class Seq2seq(nn.Module):\n",
    "  def __init__(self,context_dim,question_dim,answer_dim,emd_dim,hid_dim,n_layers,bidirectional,dropout):\n",
    "    super().__init__()\n",
    "\n",
    "    self.context_encoder = Encoder(context_dim,emd_dim,hid_dim,n_layers,dropout,bidirectional)\n",
    "    self.question_encoder = Encoder(question_dim,emd_dim,hid_dim,n_layers,dropout,bidirectional)\n",
    "    self.answer_decoder = Decoder(answer_dim,emd_dim,hid_dim,n_layers,bidirectional,dropout)\n",
    "\n",
    "    self.answer_dim = answer_dim\n",
    "\n",
    "  def forward(self,context,question,answer,teacher_forcing =0.5):\n",
    "\n",
    "    hidden,cell_state = self.context_encoder(context)\n",
    "\n",
    "    hidden,cell_state = self.question_encoder(question,hidden,cell_state)\n",
    "\n",
    "    answer_len = len(answer)\n",
    "    batch_size = answer.shape[1]\n",
    "\n",
    "    outputs = torch.zeros(answer_len,batch_size,self.answer_dim).to(device)\n",
    "\n",
    "    for i,j in enumerate(range(answer_len)):\n",
    "      k = answer[j]\n",
    "      if i != 0:\n",
    "        k = prediction.argmax(1) if random.random() < teacher_forcing else k\n",
    "      prediction, hidden, cell_state = self.answer_decoder(k,hidden,cell_state)\n",
    "      outputs[j] = prediction\n",
    "\n",
    "    return outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "7a64ckaL1r80",
    "outputId": "f282f6ca-dcc3-4b4f-8032-7353d95cdaf2"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/rnn.py:61: UserWarning: dropout option adds dropout after all but last recurrent layer, so non-zero dropout expects num_layers greater than 1, but got dropout=0.5 and num_layers=1\n",
      "  \"num_layers={}\".format(dropout, num_layers))\n"
     ]
    }
   ],
   "source": [
    "context_dim = len(context.vocab)\n",
    "question_dim = len(question.vocab)\n",
    "answer_dim = len(answer.vocab)\n",
    "emb_dim = 100\n",
    "hid_dim = 100\n",
    "n_layers = 1\n",
    "bidirectional = False\n",
    "dropout = 0.5\n",
    "\n",
    "model = Seq2seq(context_dim,question_dim,answer_dim,emb_dim,hid_dim,n_layers,bidirectional,dropout).to(device)\n",
    "\n",
    "def init_weights(m):\n",
    "    for name, param in m.named_parameters():\n",
    "      if not isinstance(m, Embedding):\n",
    "        nn.init.normal_(param.data, mean=0, std=0.01)\n",
    "        \n",
    "model.apply(init_weights)\n",
    "\n",
    "optimizer = optim.Adam(model.parameters())\n",
    "\n",
    "\n",
    "TRG_PAD_IDX = answer.vocab.stoi[answer.pad_token]\n",
    "\n",
    "criterion = nn.CrossEntropyLoss(ignore_index = TRG_PAD_IDX)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "uM3lPXZDeKFp",
    "outputId": "1074eb34-a423-4d11-dbe1-f1746ac395c7"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The model has 1,835,590 trainable parameters\n"
     ]
    }
   ],
   "source": [
    "def count_parameters(model):\n",
    "    return sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "\n",
    "print(f'The model has {count_parameters(model):,} trainable parameters')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "id": "SqDxwXf_2D0i"
   },
   "outputs": [],
   "source": [
    "def train(model, iterator, optimizer, criterion, clip):\n",
    "    \n",
    "    model.train()\n",
    "    \n",
    "    epoch_loss = 0\n",
    "    \n",
    "    for i, batch in enumerate(iterator):\n",
    "        \n",
    "        context_ = batch.context\n",
    "        question_ = batch.question\n",
    "        answer_ = batch.answer\n",
    "        \n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        output = model(context_, question_,answer_)\n",
    "        \n",
    "        trg = answer_\n",
    "        #trg = [trg len, batch size]\n",
    "        #output = [trg len, batch size, output dim]\n",
    "        \n",
    "        output_dim = output.shape[-1]\n",
    "        \n",
    "        output = output[1:].view(-1, output_dim)\n",
    "        trg = trg[1:].view(-1)\n",
    "        \n",
    "        #trg = [(trg len - 1) * batch size]\n",
    "        #output = [(trg len - 1) * batch size, output dim]\n",
    "        \n",
    "        loss = criterion(output, trg)\n",
    "        \n",
    "        loss.backward()\n",
    "        \n",
    "        torch.nn.utils.clip_grad_norm_(model.parameters(), clip)\n",
    "        \n",
    "        optimizer.step()\n",
    "        \n",
    "        epoch_loss += loss.item()\n",
    "        \n",
    "    return epoch_loss / len(iterator)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "id": "WmRFlHIW2QQ2"
   },
   "outputs": [],
   "source": [
    "def evaluate(model, iterator, criterion):\n",
    "    \n",
    "    model.eval()\n",
    "    \n",
    "    epoch_loss = 0\n",
    "    \n",
    "    with torch.no_grad():\n",
    "    \n",
    "        for i, batch in enumerate(iterator):\n",
    "\n",
    "            context_ = batch.context\n",
    "            question_ = batch.question\n",
    "            answer_ = batch.answer\n",
    "        \n",
    "            output = model(context_, question_,answer_,0) #turn off teacher forcing\n",
    "\n",
    "            trg = answer_\n",
    "            #trg = [trg len, batch size]\n",
    "            #output = [trg len, batch size, output dim]\n",
    "\n",
    "            output_dim = output.shape[-1]\n",
    "            \n",
    "            output = output[1:].view(-1, output_dim)\n",
    "            trg = trg[1:].view(-1)\n",
    "\n",
    "            #trg = [(trg len - 1) * batch size]\n",
    "            #output = [(trg len - 1) * batch size, output dim]\n",
    "\n",
    "            loss = criterion(output, trg)\n",
    "\n",
    "            epoch_loss += loss.item()\n",
    "        \n",
    "    return epoch_loss / len(iterator)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "id": "iHsk7wYu2Tdj"
   },
   "outputs": [],
   "source": [
    "def epoch_time(start_time, end_time):\n",
    "    elapsed_time = end_time - start_time\n",
    "    elapsed_mins = int(elapsed_time / 60)\n",
    "    elapsed_secs = int(elapsed_time - (elapsed_mins * 60))\n",
    "    return elapsed_mins, elapsed_secs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "kPvtYPOw2VqN",
    "outputId": "e42842c9-b85e-4056-a072-eceb962552f3"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 01 | Time: 0m 34s\n",
      "\tTrain Loss: 8.356 | Train PPL: 4256.789\n",
      "\t Val. Loss: 8.340 |  Val. PPL: 4188.036\n",
      "Epoch: 02 | Time: 0m 34s\n",
      "\tTrain Loss: 8.319 | Train PPL: 4100.958\n",
      "\t Val. Loss: 8.248 |  Val. PPL: 3821.181\n",
      "Epoch: 03 | Time: 0m 34s\n",
      "\tTrain Loss: 7.641 | Train PPL: 2082.310\n",
      "\t Val. Loss: 6.339 |  Val. PPL: 566.443\n",
      "Epoch: 04 | Time: 0m 33s\n",
      "\tTrain Loss: 6.034 | Train PPL: 417.516\n",
      "\t Val. Loss: 5.093 |  Val. PPL: 162.831\n",
      "Epoch: 05 | Time: 0m 34s\n",
      "\tTrain Loss: 5.121 | Train PPL: 167.522\n",
      "\t Val. Loss: 4.200 |  Val. PPL:  66.661\n",
      "Epoch: 06 | Time: 0m 34s\n",
      "\tTrain Loss: 4.584 | Train PPL:  97.898\n",
      "\t Val. Loss: 3.716 |  Val. PPL:  41.115\n",
      "Epoch: 07 | Time: 0m 34s\n",
      "\tTrain Loss: 4.385 | Train PPL:  80.252\n",
      "\t Val. Loss: 3.560 |  Val. PPL:  35.163\n",
      "Epoch: 08 | Time: 0m 34s\n",
      "\tTrain Loss: 4.373 | Train PPL:  79.251\n",
      "\t Val. Loss: 3.541 |  Val. PPL:  34.491\n",
      "Epoch: 09 | Time: 0m 34s\n",
      "\tTrain Loss: 4.376 | Train PPL:  79.549\n",
      "\t Val. Loss: 3.526 |  Val. PPL:  33.981\n",
      "Epoch: 10 | Time: 0m 33s\n",
      "\tTrain Loss: 4.363 | Train PPL:  78.469\n",
      "\t Val. Loss: 3.525 |  Val. PPL:  33.947\n"
     ]
    }
   ],
   "source": [
    "import math\n",
    "N_EPOCHS = 10\n",
    "CLIP = 1\n",
    "\n",
    "best_valid_loss = float('inf')\n",
    "\n",
    "for epoch in range(N_EPOCHS):\n",
    "    \n",
    "    start_time = time.time()\n",
    "    \n",
    "    train_loss = train(model, train_iterator, optimizer, criterion, CLIP)\n",
    "    valid_loss = evaluate(model, valid_iterator, criterion)\n",
    "    \n",
    "    end_time = time.time()\n",
    "    \n",
    "    epoch_mins, epoch_secs = epoch_time(start_time, end_time)\n",
    "    \n",
    "    if valid_loss < best_valid_loss:\n",
    "        best_valid_loss = valid_loss\n",
    "        torch.save(model.state_dict(), 'tut2-model.pt')\n",
    "    \n",
    "    print(f'Epoch: {epoch+1:02} | Time: {epoch_mins}m {epoch_secs}s')\n",
    "    print(f'\\tTrain Loss: {train_loss:.3f} | Train PPL: {math.exp(train_loss):7.3f}')\n",
    "    print(f'\\t Val. Loss: {valid_loss:.3f} |  Val. PPL: {math.exp(valid_loss):7.3f}')"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "name": "NaturalQuestionQA with Seq2Seq.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}