Train_model_TPU_BERT.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/alexpod1000/SQuAD-QA/blob/main/Train_model_TPU_BERT.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "oTkVfFrJ-pzG",
    "outputId": "bff4a152-b667-4f03-f6ab-9429273a4a82"
   },
   "outputs": [],
   "source": [
    "# Run the following cells only if using Colab\n",
    "if 'google.colab' in str(get_ipython()):\n",
    "    # Clone repository\n",
    "    !git clone https://github.com/alexpod1000/SQuAD-QA.git\n",
    "    # Change current working directory to match project\n",
    "    � SQuAD-QA/\n",
    "    !pwd\n",
    "\n",
    "    !pip install transformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "eGf4IyEMuKSe",
    "outputId": "108a392d-f7e8-4883-a85f-4239f1d877ff"
   },
   "outputs": [],
   "source": [
    "!pip install cloud-tpu-client==0.10 https://storage.googleapis.com/tpu-pytorch/wheels/torch_xla-1.8-cp37-cp37m-linux_x86_64.whl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "SdrMxuA4FdDr"
   },
   "outputs": [],
   "source": [
    "!export XLA_USE_BF16=1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "UO4BuDyxuCVi",
    "outputId": "82b46c78-c7a7-4ce7-8499-60a05693484d"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:root:TPU has started up successfully with version pytorch-1.7\n"
     ]
    }
   ],
   "source": [
    "\n",
    "# imports the torch_xla package\n",
    "import torch_xla\n",
    "import torch_xla.core.xla_model as xm\n",
    "import torch_xla.distributed.parallel_loader as pl\n",
    "import torch_xla.distributed.xla_multiprocessing as xmp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "id": "rsBVuJu6_5qN"
   },
   "outputs": [],
   "source": [
    "# External imports\n",
    "import copy\n",
    "import nltk\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import string\n",
    "import torch\n",
    "import json\n",
    "\n",
    "from functools import partial\n",
    "from nltk.tokenize import TreebankWordTokenizer, SpaceTokenizer\n",
    "from transformers import AutoTokenizer\n",
    "from typing import Tuple, List, Dict, Any, Union\n",
    "\n",
    "# Project imports\n",
    "from squad_data.parser import SquadFileParser\n",
    "from squad_data.utils import build_mappers_and_dataframe_bert\n",
    "from models import possible_models_dict\n",
    "from evaluation.evaluate import evaluate_predictions\n",
    "from evaluation.utils import build_evaluation_dict_bert\n",
    "from utils import split_dataframe\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "lIxxj4eptf2A"
   },
   "outputs": [],
   "source": [
    "#current_selected_experiment = \"distilbert\"\n",
    "#current_selected_experiment = \"bert\"\n",
    "#current_selected_experiment = \"albert\"\n",
    "#current_selected_experiment = \"distilroberta\"\n",
    "current_selected_experiment = \"distilroberta_extra_linear\"\n",
    "params_dict = possible_models_dict[current_selected_experiment]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9Rh4dSW-9tYm"
   },
   "source": [
    "### Parse the json and get the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "id": "FAEEYoypAOKA"
   },
   "outputs": [],
   "source": [
    "train_file_json = \"squad_data/data/training_set.json\"\n",
    "test_file_json = \"squad_data/data/dev-v1.1.json\"\n",
    "\n",
    "train_parser = SquadFileParser(train_file_json)\n",
    "test_parser = SquadFileParser(test_file_json)\n",
    "\n",
    "train_data = train_parser.parse_documents()\n",
    "test_data = test_parser.parse_documents()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SKK-4d1_93QE"
   },
   "source": [
    "### Prepare the mappers and datafram"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "id": "cwmRU7yftf2E"
   },
   "outputs": [],
   "source": [
    "def bert_tokenizer_fn(question, paragraph, tokenizer, max_length=384, doc_stride=128):\n",
    "    pad_on_right = tokenizer.padding_side == \"right\"\n",
    "    # Process the sample\n",
    "    tokenized_input_pair = tokenizer(\n",
    "        question,\n",
    "        paragraph,\n",
    "        truncation=\"only_second\" if pad_on_right else \"only_first\",\n",
    "        max_length=max_length,\n",
    "        stride=doc_stride,\n",
    "        return_overflowing_tokens=True,\n",
    "        return_offsets_mapping=True,\n",
    "        padding=\"max_length\",\n",
    "    )\n",
    "    return tokenized_input_pair"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "id": "vd_r9nZLtf2E"
   },
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(params_dict[\"tokenizer_url\"])\n",
    "tokenizer_fn_preprocess = partial(bert_tokenizer_fn, tokenizer=tokenizer, max_length=params_dict[\"tokenizer_max_length\"]-3)\n",
    "tokenizer_fn_train = partial(bert_tokenizer_fn, tokenizer=tokenizer, max_length=params_dict[\"tokenizer_max_length\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 243
    },
    "id": "VbmMXc0Ptf2G",
    "outputId": "f2bb7d77-cf98-4a15-c90e-be68fde7ccc6"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "architecturally, the school has a catholic character. atop the main building's gold dome is a golden statue of the virgin mary. immediately in front of the main building and facing it, is a copper statue of christ with arms upraised with the legend \"venite ad me omnes\". next to the main building is the basilica of the sacred heart. immediately behind the basilica is the grotto, a marian place of prayer and reflection. it is a replica of the grotto at lourdes, france where the virgin mary reputedly appeared to saint bernadette soubirous in 1858. at the end of the main drive (and in a direct line that connects through 3 statues and the gold dome), is a simple, modern stone statue of mary.\n"
     ]
    },
    {
     "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>doc_id</th>\n",
       "      <th>paragraph_id</th>\n",
       "      <th>question_id</th>\n",
       "      <th>answer_id</th>\n",
       "      <th>answer_start</th>\n",
       "      <th>answer_text</th>\n",
       "      <th>question_text</th>\n",
       "      <th>tokenizer_answer_start</th>\n",
       "      <th>tokenizer_answer_end</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>0_0</td>\n",
       "      <td>5733be284776f41900661182</td>\n",
       "      <td>0</td>\n",
       "      <td>515</td>\n",
       "      <td>Saint Bernadette Soubirous</td>\n",
       "      <td>To whom did the Virgin Mary allegedly appear i...</td>\n",
       "      <td>128</td>\n",
       "      <td>135</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0</td>\n",
       "      <td>0_0</td>\n",
       "      <td>5733be284776f4190066117f</td>\n",
       "      <td>0</td>\n",
       "      <td>188</td>\n",
       "      <td>a copper statue of Christ</td>\n",
       "      <td>What is in front of the Notre Dame Main Building?</td>\n",
       "      <td>52</td>\n",
       "      <td>57</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>0</td>\n",
       "      <td>0_0</td>\n",
       "      <td>5733be284776f41900661180</td>\n",
       "      <td>0</td>\n",
       "      <td>279</td>\n",
       "      <td>the Main Building</td>\n",
       "      <td>The Basilica of the Sacred heart at Notre Dame...</td>\n",
       "      <td>83</td>\n",
       "      <td>86</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>0</td>\n",
       "      <td>0_0</td>\n",
       "      <td>5733be284776f41900661181</td>\n",
       "      <td>0</td>\n",
       "      <td>381</td>\n",
       "      <td>a Marian place of prayer and reflection</td>\n",
       "      <td>What is the Grotto at Notre Dame?</td>\n",
       "      <td>95</td>\n",
       "      <td>102</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>0</td>\n",
       "      <td>0_0</td>\n",
       "      <td>5733be284776f4190066117e</td>\n",
       "      <td>0</td>\n",
       "      <td>92</td>\n",
       "      <td>a golden statue of the Virgin Mary</td>\n",
       "      <td>What sits on top of the Main Building at Notre...</td>\n",
       "      <td>33</td>\n",
       "      <td>40</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   doc_id paragraph_id  ... tokenizer_answer_start  tokenizer_answer_end\n",
       "0       0          0_0  ...                    128                   135\n",
       "1       0          0_0  ...                     52                    57\n",
       "2       0          0_0  ...                     83                    86\n",
       "3       0          0_0  ...                     95                   102\n",
       "4       0          0_0  ...                     33                    40\n",
       "\n",
       "[5 rows x 9 columns]"
      ]
     },
     "execution_count": 11,
     "metadata": {
      "tags": []
     },
     "output_type": "execute_result"
    }
   ],
   "source": [
    "paragraphs_mapper, df = build_mappers_and_dataframe_bert(tokenizer, tokenizer_fn_preprocess, train_data, limit_answers=1)\n",
    "print(paragraphs_mapper[next(iter(paragraphs_mapper))])\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "id": "m-shyowvtf2G"
   },
   "outputs": [],
   "source": [
    "df_train, df_val = split_dataframe(df, train_ratio=0.9)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "30y2vgMmtf2H",
    "outputId": "de22177a-5d9f-47ba-a866-90e02990d781"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total samples: 88649, Train samples: 61618, Validation samples: 27031\n"
     ]
    }
   ],
   "source": [
    "print(f\"Total samples: {len(df)}, Train samples: {len(df_train)}, Validation samples: {len(df_val)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Xb4YK_Qa95zK"
   },
   "source": [
    "### DataConverter and CustomQADataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "KDl4CIW-mj_D",
    "outputId": "f5b6b23f-c250-45ea-d62d-d81558a410d5"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([10, 384])\n",
      "torch.Size([10, 2])\n"
     ]
    }
   ],
   "source": [
    "from data_loading.utils import bert_padder_collate_fn\n",
    "from data_loading.qa_dataset import CustomQADatasetBERT\n",
    "\n",
    "datasetQA = CustomQADatasetBERT(tokenizer_fn_train, df_train, paragraphs_mapper)\n",
    "data_loader = torch.utils.data.DataLoader(datasetQA, collate_fn = bert_padder_collate_fn, batch_size=10, shuffle=True)\n",
    "\n",
    "test_batch = next(iter(data_loader))\n",
    "print(test_batch[\"input_ids\"].shape)\n",
    "print(test_batch[\"y_gt\"].shape)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "g1gz_Kq5tf2K"
   },
   "source": [
    "# Model train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "id": "DGoCBTs7tf2K"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "import transformers\n",
    "\n",
    "from timeit import default_timer as timer\n",
    "from tqdm import tqdm\n",
    "from transformers import get_linear_schedule_with_warmup\n",
    "from transformers.optimization import AdamW\n",
    "\n",
    "from models.utils import SpanExtractor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Jd2c5n7Jtf2K",
    "outputId": "042dee0f-391a-4a53-e8d6-dfd8b6da1ef2"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Automatic Mixed Precision Enabled: True\n"
     ]
    }
   ],
   "source": [
    "use_amp = True\n",
    "print(f\"Automatic Mixed Precision Enabled: {use_amp}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nqRp-tr7tf2K"
   },
   "source": [
    "Model:\n",
    "\n",
    "(input_ids, attention_mask) -> (answer_start, answer_end) // for each token in input_ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "id": "0KTdFFVytf2K"
   },
   "outputs": [],
   "source": [
    "def train_step(model, scaler, optimizer, loss_function, dataloader, scheduler=None, device=\"cpu\", show_progress=False):\n",
    "    acc_loss = 0\n",
    "    acc_start_accuracy = 0\n",
    "    acc_end_accuracy = 0\n",
    "    count = 0\n",
    "\n",
    "    time_start = timer()\n",
    "    \n",
    "    model.train()\n",
    "    wrapped_dataloader = tqdm(dataloader) if show_progress else dataloader\n",
    "    for batch in wrapped_dataloader:\n",
    "        # NOTE: we'll pass directly the batch dict to the model for inputs.\n",
    "        answer_spans_start = batch[\"y_gt\"][:, 0]\n",
    "        answer_spans_end = batch[\"y_gt\"][:, 1]\n",
    "        # Clear gradients\n",
    "        model.zero_grad()\n",
    "        # Place to right device\n",
    "        answer_spans_start = answer_spans_start.to(device)\n",
    "        answer_spans_end = answer_spans_end.to(device)\n",
    "        # Use Automatic Mixed Precision if enabled\n",
    "        #with torch.cuda.amp.autocast(enabled=scaler.is_enabled()):\n",
    "        # Run forward pass\n",
    "        pred_answer_start_scores, pred_answer_end_scores = model(batch)\n",
    "        # Compute the CrossEntropyLoss\n",
    "        loss = (loss_function(pred_answer_start_scores, answer_spans_start)   loss_function(pred_answer_end_scores, answer_spans_end))/2.0\n",
    "        scaler.scale(loss).backward()\n",
    "        # Optimizer step (via scaler)\n",
    "        xm.optimizer_step(optimizer)\n",
    "        if scheduler is not None:\n",
    "            scheduler.step()\n",
    "        #scaler.step(optimizer)\n",
    "        #scaler.update()\n",
    "        # --- Compute metrics ---\n",
    "        # Get span indexes\n",
    "        pred_span_start_idxs, pred_span_end_idxs = SpanExtractor.extract_most_probable(pred_answer_start_scores, pred_answer_end_scores)\n",
    "        gt_start_idxs = answer_spans_start.cpu().detach()\n",
    "        gt_end_idxs = answer_spans_end.cpu().detach()\n",
    "        # two accs\n",
    "        start_accuracy = torch.sum(gt_start_idxs == pred_span_start_idxs) / len(pred_span_start_idxs)\n",
    "        end_accuracy = torch.sum(gt_end_idxs == pred_span_end_idxs) / len(pred_span_end_idxs)\n",
    "        # Gather stats\n",
    "        acc_loss  = loss.item()\n",
    "        acc_start_accuracy  = start_accuracy.item()\n",
    "        acc_end_accuracy  = end_accuracy.item()\n",
    "        count  = 1\n",
    "    time_end = timer()\n",
    "    return {\n",
    "        \"loss\": acc_loss / count, \n",
    "        \"accuracy_start\": acc_start_accuracy / count, \n",
    "        \"accuracy_end\": acc_end_accuracy / count,\n",
    "        \"time\": time_end - time_start\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "id": "cQdTcxgktf2K"
   },
   "outputs": [],
   "source": [
    "@torch.no_grad()\n",
    "def validation_step(model, scaler, loss_function, dataloader, device=\"cpu\", show_progress=False):\n",
    "    acc_loss = 0\n",
    "    acc_start_accuracy = 0\n",
    "    acc_end_accuracy = 0\n",
    "    count = 0\n",
    "\n",
    "    time_start = timer()\n",
    "    wrapped_dataloader = tqdm(dataloader) if show_progress else dataloader\n",
    "    \n",
    "    model.eval()\n",
    "    for batch in wrapped_dataloader:\n",
    "        answer_spans_start = batch[\"y_gt\"][:, 0]\n",
    "        answer_spans_end = batch[\"y_gt\"][:, 1]\n",
    "        # Place to right device\n",
    "        answer_spans_start = answer_spans_start.to(device)\n",
    "        answer_spans_end = answer_spans_end.to(device)\n",
    "        # Use Automatic Mixed Precision if enabled\n",
    "        # Run forward pass\n",
    "        pred_answer_start_scores, pred_answer_end_scores = model(batch)\n",
    "        # Compute the CrossEntropyLoss\n",
    "        loss = (loss_function(pred_answer_start_scores, answer_spans_start)   loss_function(pred_answer_end_scores, answer_spans_end))/2.0\n",
    "        # --- Compute metrics ---\n",
    "        # Get span indexes\n",
    "        pred_span_start_idxs, pred_span_end_idxs = SpanExtractor.extract_most_probable(pred_answer_start_scores, pred_answer_end_scores)\n",
    "        gt_start_idxs = answer_spans_start.cpu().detach()\n",
    "        gt_end_idxs = answer_spans_end.cpu().detach()\n",
    "        # two accs\n",
    "        start_accuracy = torch.sum(gt_start_idxs == pred_span_start_idxs) / len(pred_span_start_idxs)\n",
    "        end_accuracy = torch.sum(gt_end_idxs == pred_span_end_idxs) / len(pred_span_end_idxs)\n",
    "        # Gather stats\n",
    "        acc_loss  = loss.item()\n",
    "        acc_start_accuracy  = start_accuracy.item()\n",
    "        acc_end_accuracy  = end_accuracy.item()\n",
    "        count  = 1\n",
    "    time_end = timer()\n",
    "    return {\n",
    "        \"loss\": acc_loss / count, \n",
    "        \"accuracy_start\": acc_start_accuracy / count, \n",
    "        \"accuracy_end\": acc_end_accuracy / count,\n",
    "        \"time\": time_end - time_start\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "id": "tEYK2FuLtf2L"
   },
   "outputs": [],
   "source": [
    "class ParametricBertModelQA(torch.nn.Module):\n",
    "\n",
    "    def __init__(self, hidden_size, num_labels, config_dict, dropout_rate=0.3):\n",
    "        super(ParametricBertModelQA, self).__init__()\n",
    "        self.hidden_size = hidden_size\n",
    "        self.num_labels = num_labels\n",
    "        self.bert = transformers.AutoModel.from_pretrained(config_dict[\"model_url\"])#(bert_config)\n",
    "        self.bert_drop = torch.nn.Dropout(dropout_rate)\n",
    "        self.qa_outputs = torch.nn.Linear(self.hidden_size, self.num_labels)\n",
    "        self.prepare_input_fn = config_dict[\"prepare_model_input_fn\"]\n",
    "\n",
    "    #@torch.cuda.amp.autocast() # goes OOM for whatever reason, don't use.\n",
    "    def forward(self, inputs):\n",
    "        # --- 1) Extract data from inputs dictionary and put it on right device\n",
    "        curr_device = self.bert.device\n",
    "        # --- 2) Run BERT backbone to produce final representation\n",
    "        input_dict_for_bert = self.prepare_input_fn(inputs, curr_device)\n",
    "        output = self.bert(**input_dict_for_bert)\n",
    "        # --- 3) On top of the final representation, run a mapper to get scores for each position.\n",
    "        sequence_output = output[0]   #(None, seq_len, hidden_size)\n",
    "        # do dropout\n",
    "        sequence_output = self.bert_drop(sequence_output)\n",
    "        logits = self.qa_outputs(sequence_output) #(None, seq_len, hidden_size)*(hidden_size, 2)=(None, seq_len, 2)\n",
    "        start_logits, end_logits = logits.split(1, dim=-1)    #(None, seq_len, 1), (None, seq_len, 1)\n",
    "        start_logits = start_logits.squeeze(-1)  #(None, seq_len)\n",
    "        end_logits = end_logits.squeeze(-1)    #(None, seq_len)\n",
    "        # --- 4) Prepare output tuple\n",
    "        outputs = (start_logits, end_logits,) \n",
    "        return outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_params_for_optimizer(model, no_decay, weight_decay=0.0001):\n",
    "    param_optimizer = list(model.named_parameters())\n",
    "    optimizer_parameters = [\n",
    "        {\n",
    "            'params': [\n",
    "                p for n, p in param_optimizer if not any(\n",
    "                    nd in n for nd in no_decay\n",
    "                )\n",
    "            ], \n",
    "            'weight_decay': weight_decay\n",
    "        },\n",
    "        {\n",
    "            'params': [\n",
    "                p for n, p in param_optimizer if any(\n",
    "                    nd in n for nd in no_decay\n",
    "                )\n",
    "            ],\n",
    "            'weight_decay': 0.0\n",
    "        },\n",
    "    ]\n",
    "    return optimizer_parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "id": "KcKLemArzJK7"
   },
   "outputs": [],
   "source": [
    "def train_single_model(index, flags):\n",
    "    torch.set_default_tensor_type('torch.FloatTensor')\n",
    "    #torch.manual_seed(flags['seed'])\n",
    "\n",
    "    # BIG WARNING: THE xm.xla_device() line should be run only inside this function, otherwise the device identifier becomes wrong, and the function throws exceptions thinking the TPU has only 1 core\n",
    "    # Acquires the (unique) Cloud TPU core corresponding to this process's index\n",
    "    device = xm.xla_device() \n",
    "\n",
    "    # will show progress only if master\n",
    "    show_progress = xm.is_master_ordinal()\n",
    "    \n",
    "    dataset_train_QA = CustomQADatasetBERT(tokenizer_fn_train, df_train, paragraphs_mapper)\n",
    "    dataset_val_QA = CustomQADatasetBERT(tokenizer_fn_train, df_val, paragraphs_mapper)\n",
    "\n",
    "    train_sampler = torch.utils.data.distributed.DistributedSampler(\n",
    "        dataset_train_QA,\n",
    "        num_replicas=xm.xrt_world_size(),\n",
    "        rank=xm.get_ordinal(),\n",
    "        shuffle=True)\n",
    "\n",
    "    val_sampler = torch.utils.data.distributed.DistributedSampler(\n",
    "        dataset_val_QA,\n",
    "        num_replicas=xm.xrt_world_size(),\n",
    "        rank=xm.get_ordinal(),\n",
    "        shuffle=True)\n",
    "    \n",
    "    train_data_loader = torch.utils.data.DataLoader(\n",
    "        dataset_train_QA, collate_fn = bert_padder_collate_fn, \n",
    "        batch_size=params_dict[\"train_params\"][\"batch_size_train\"],\n",
    "        sampler=train_sampler\n",
    "    )\n",
    "    val_data_loader = torch.utils.data.DataLoader(\n",
    "        dataset_val_QA, collate_fn = bert_padder_collate_fn, \n",
    "        batch_size=params_dict[\"train_params\"][\"batch_size_val\"],\n",
    "        sampler=val_sampler\n",
    "    )\n",
    "\n",
    "    # Define baseline model\n",
    "    model = params_dict[\"span_model\"](768, 2, params_dict, dropout_rate=params_dict[\"train_params\"][\"dropout_rate\"])\n",
    "    #model = ExtraLayer_ParametricBertModelQA(768, 2, params_dict)\n",
    "    model = xmp.MpModelWrapper(model)\n",
    "    model = model.to(device)\n",
    "\n",
    "    # Loss function\n",
    "    loss_function = nn.CrossEntropyLoss()\n",
    "    # Produce a list of parameters for optimization, with relative weight decay\n",
    "    no_decay = [\"bias\", \"LayerNorm.bias\", \"LayerNorm.weight\"]\n",
    "    #no_decay = []\n",
    "    if params_dict[\"train_params\"][\"weight_decay\"] > 0.0:\n",
    "        model_params_optimizer = get_params_for_optimizer(model, no_decay, weight_decay=params_dict[\"train_params\"][\"weight_decay\"])\n",
    "    else:\n",
    "        model_params_optimizer = model.parameters()\n",
    "    # Optimization\n",
    "    #\"\"\"\n",
    "    optimizer = AdamW(\n",
    "        model_params_optimizer, \n",
    "        lr=params_dict[\"train_params\"][\"initial_lr\"],# * xm.xrt_world_size(), \n",
    "        correct_bias=False\n",
    "    )\n",
    "    #\"\"\"\n",
    "    #optimizer = torch.optim.AdamW(model_params_optimizer, lr=params_dict[\"train_params\"][\"initial_lr\"] * xm.xrt_world_size(), weight_decay=params_dict[\"train_params\"][\"weight_decay\"])\n",
    "\n",
    "    # Estimate the number of train steps for LR scheduler\n",
    "    num_train_steps = int(\n",
    "        len(df_train) / params_dict[\"train_params\"][\"batch_size_train\"] / xm.xrt_world_size() * params_dict[\"train_params\"][\"epochs\"]\n",
    "    )\n",
    "\n",
    "    num_warmup_steps = int(num_train_steps * 0.1) # 10% of warmup steps\n",
    "\n",
    "    # LR scheduler\n",
    "    scheduler = get_linear_schedule_with_warmup(\n",
    "        optimizer,\n",
    "        num_warmup_steps=num_warmup_steps,\n",
    "        num_training_steps=num_train_steps\n",
    "    )\n",
    "\n",
    "    #scheduler = None\n",
    "\n",
    "    scaler = torch.cuda.amp.GradScaler(enabled=False)\n",
    "\n",
    "    history = {\n",
    "        \"train_loss\": [], \"train_acc_start\": [], \"train_acc_end\": [],\n",
    "        \"val_loss\": [], \"val_acc_start\": [], \"val_acc_end\": []\n",
    "    }\n",
    "    # Do the training\n",
    "    loop_start = timer()\n",
    "    for epoch in range(params_dict[\"train_params\"][\"epochs\"]):\n",
    "        # Train step\n",
    "        para_loader = pl.ParallelLoader(train_data_loader, [device])\n",
    "        train_dict = train_step(\n",
    "            model, scaler, optimizer, loss_function, \n",
    "            para_loader.per_device_loader(device), \n",
    "            scheduler=scheduler, \n",
    "            device=device, \n",
    "            show_progress=show_progress\n",
    "        )\n",
    "        # Validation step\n",
    "        para_loader = pl.ParallelLoader(val_data_loader, [device])\n",
    "        val_dict = validation_step(model, scaler, loss_function, para_loader.per_device_loader(device), device=device, show_progress=show_progress)\n",
    "        cur_lr = optimizer.param_groups[0]['lr']\n",
    "\n",
    "        # mesh reduce all the metrics\n",
    "        train_dict[\"loss\"] = xm.mesh_reduce(\"train_loss\", train_dict[\"loss\"], np.mean)\n",
    "        train_dict[\"accuracy_start\"] = xm.mesh_reduce(\"train_accuracy_start\", train_dict[\"accuracy_start\"], np.mean)\n",
    "        train_dict[\"accuracy_end\"] = xm.mesh_reduce(\"train_accuracy_end\", train_dict[\"accuracy_end\"], np.mean)\n",
    "        val_dict[\"loss\"] = xm.mesh_reduce(\"val_loss\", val_dict[\"loss\"], np.mean)\n",
    "        val_dict[\"accuracy_start\"] = xm.mesh_reduce(\"val_accuracy_start\", val_dict[\"accuracy_start\"], np.mean)\n",
    "        val_dict[\"accuracy_end\"] = xm.mesh_reduce(\"val_accuracy_end\", val_dict[\"accuracy_end\"], np.mean)\n",
    "        # end mesh reduce\n",
    "\n",
    "        xm.master_print(f'[xla:{xm.get_ordinal()}], Epoch: {epoch}, '\n",
    "            f'lr: {cur_lr}, '\n",
    "            f'Train loss: {train_dict[\"loss\"]:.4f}, '\n",
    "            f'Train acc start: {train_dict[\"accuracy_start\"]:.4f}, '\n",
    "            f'Train acc end: {train_dict[\"accuracy_end\"]:.4f}, '\n",
    "            f'Val loss: {val_dict[\"loss\"]:.4f}, '\n",
    "            f'Val acc start: {val_dict[\"accuracy_start\"]:.4f}, '\n",
    "            f'Val acc end: {val_dict[\"accuracy_end\"]:.4f}, '\n",
    "            f'Time: {train_dict[\"time\"]:.4f}')\n",
    "        \n",
    "        xm.save(\n",
    "            model.state_dict(),\n",
    "            f'{current_selected_experiment}_{epoch}.pt'\n",
    "        )\n",
    "        \n",
    "        if xm.is_master_ordinal():\n",
    "            history[\"train_loss\"].append(train_dict[\"loss\"]);history[\"train_acc_start\"].append(train_dict[\"accuracy_start\"]);history[\"train_acc_end\"].append(train_dict[\"accuracy_end\"]);\n",
    "            history[\"val_loss\"].append(val_dict[\"loss\"]);history[\"val_acc_start\"].append(val_dict[\"accuracy_start\"]);history[\"val_acc_end\"].append(val_dict[\"accuracy_end\"]);\n",
    "        #scheduler.step(val_dict[\"loss\"])\n",
    "        #print(f\"Evaluation Results: {eval_results}\")\n",
    "    loop_end = timer()\n",
    "    xm.master_print(f\"[xla:{xm.get_ordinal()}] Elapsed time: {(loop_end - loop_start):.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kjRhv034FK9x"
   },
   "source": [
    "XLA_USE_BF16 flag is used to improve performance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Kzj-ybB2zv6v",
    "outputId": "74eb75f1-3be6-41e8-ce4b-b96381e8906d"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 22%|██▏       | 208/963 [01:37<05:11,  2.43it/s]"
     ]
    }
   ],
   "source": [
    "# Configures training (and evaluation) parameters\n",
    "flags = {}\n",
    "flags['num_workers'] = 8\n",
    "flags['num_epochs'] = 1\n",
    "#flags['seed'] = 1234\n",
    "\n",
    "xmp.spawn(train_single_model, args=(flags,), nprocs=8, start_method='fork')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "nmpiBMUftf2M"
   },
   "outputs": [],
   "source": [
    "# Uncomment below line to save model to disk\n",
    "#torch.save(model.state_dict(), \"distilbert_mdl.pt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "M0hTgklltf2M"
   },
   "source": [
    "# Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "WOkB5SAztf2M"
   },
   "outputs": [],
   "source": [
    "# Uncomment below line to load model from disk\n",
    "#model.load_state_dict(torch.load(\"distilbert_mdl.pt\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QAwLi80qtf2M"
   },
   "source": [
    "## Quantitative evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "O7PDBIzIO-wu"
   },
   "outputs": [],
   "source": [
    "scaler = torch.cuda.amp.GradScaler(enabled=False) # for compatibility with non-TPU pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "AVkWot7vtf2M"
   },
   "outputs": [],
   "source": [
    "test_paragraphs_mapper, test_df = build_mappers_and_dataframe_bert(tokenizer, tokenizer_fn_preprocess, test_data, limit_answers=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "VWrVOUCZtf2M"
   },
   "outputs": [],
   "source": [
    "dataset_test_QA = CustomQADatasetBERT(tokenizer_fn_train, test_df, test_paragraphs_mapper)\n",
    "test_data_loader = torch.utils.data.DataLoader(\n",
    "    dataset_test_QA, collate_fn = bert_padder_collate_fn, \n",
    "    batch_size=params_dict[\"train_params\"][\"batch_size_test\"], shuffle=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "VVe-CCIXtf2N"
   },
   "outputs": [],
   "source": [
    "with open(test_file_json, \"r\") as f:\n",
    "    dataset_json = json.load(f)\n",
    "pred_dict = build_evaluation_dict_bert(model, scaler, test_data_loader, test_paragraphs_mapper, tokenizer, device, show_progress=True)\n",
    "eval_results = evaluate_predictions(dataset_json, pred_dict)\n",
    "print(eval_results)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FsM9l3PPtf2N"
   },
   "source": [
    "## Simple qualitative evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tu9Pkn_ztf2O"
   },
   "outputs": [],
   "source": [
    "def get_answer_span_helper(context, question, model, tokenizer_fn, tokenizer, device=\"cpu\"):\n",
    "    tokenized_input = tokenizer_fn(question, context)\n",
    "    output_span = model({\n",
    "        \"input_ids\": torch.tensor(tokenized_input[\"input_ids\"]).to(device), \n",
    "        \"attention_mask\": torch.tensor(tokenized_input[\"attention_mask\"]).to(device)\n",
    "    })\n",
    "    start, end = SpanExtractor.extract_most_probable(output_span[0], output_span[1])\n",
    "    start = start.item()\n",
    "    end = end.item()\n",
    "    return tokenizer.decode(tokenized_input[\"input_ids\"][0][start:end], skip_special_tokens=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "FVnbTPMxtf2O"
   },
   "outputs": [],
   "source": [
    "context = \"This is a test message, written to see if our model can correctly predict its outputs.\"\n",
    "question = \"Who needs to predict its outputs?\"\n",
    "pred_answer = get_answer_span_helper(context, question, model, tokenizer_fn_train, tokenizer, device=\"cuda\")\n",
    "print(pred_answer)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "TPU",
  "colab": {
   "collapsed_sections": [],
   "name": "Playground_NLP_Project.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "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.8.5"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": false,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}