transformers_ner.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "rFvcIe4qz_2t"
   },
   "source": [
    "# Fine Tuning Transformer for Named Entity Recognition"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Zt90X_Dw0B_T"
   },
   "source": [
    "### Introduction\n",
    "\n",
    "In this tutorial we will be fine tuning a transformer model for the **Named Entity Recognition** problem. \n",
    "This is one of the most common business problems where a given piece of text/sentence/document different entites need to be identified such as: Name, Location, Number, Entity etc.\n",
    "\n",
    "#### Flow of the notebook\n",
    "\n",
    "The notebook will be divided into seperate sections to provide a organized walk through for the process used. This process can be modified for individual use cases. The sections are:\n",
    "\n",
    "1. [Installing packages for preparing the system](#section00)\n",
    "2. [Importing Python Libraries and preparing the environment](#section01)\n",
    "3. [Importing and Pre-Processing the domain data](#section02)\n",
    "4. [Preparing the Dataset and Dataloader](#section03)\n",
    "5. [Creating the Neural Network for Fine Tuning](#section04)\n",
    "6. [Fine Tuning the Model](#section05)\n",
    "7. [Validating the Model Performance](#section06)\n",
    "\n",
    "#### Technical Details\n",
    "\n",
    "This script leverages on multiple tools designed by other teams. Details of the tools used below. Please ensure that these elements are present in your setup to successfully implement this script.\n",
    "\n",
    " - Data:\n",
    "\t- We are working from a dataset available on [Kaggle](https://www.kaggle.com/)\n",
    "    - This NER annotated dataset is available at the following [link](https://www.kaggle.com/abhinavwalia95/entity-annotated-corpus)\n",
    "    - We will be working with the file `ner.csv` from the dataset. \n",
    "    - In the given file we will be looking at the following columns for the purpose of this fine tuning:\n",
    "        - `sentence_idx` : This is the identifier that the word in the row is part of the same sentence\n",
    "        - `word` : Word in the sentence\n",
    "        - `tag` : This is the identifier that is used to identify the entity in the dataset. \n",
    "    - The various entites tagged in this dataset are as per below:\n",
    "        - geo = Geographical Entity\n",
    "        - org = Organization\n",
    "        - per = Person\n",
    "        - gpe = Geopolitical Entity\n",
    "        - tim = Time indicator\n",
    "        - art = Artifact\n",
    "        - eve = Event\n",
    "        - nat = Natural Phenomenon\n",
    "\n",
    "\n",
    " - Language Model Used:\n",
    "\t - We are using BERT for this project. Hugging face team has created a customized model for token classification, called **BertForTokenClassification**. We will be using it in our custommodel class for training. \n",
    "\t - [Blog-Post](https://ai.googleblog.com/2018/11/open-sourcing-bert-state-of-art-pre.html)\n",
    "     - [Documentation for python](https://huggingface.co/transformers/model_doc/bert.html#bertfortokenclassification)\n",
    "\n",
    "\n",
    " - Hardware Requirements:\n",
    "\t - Python 3.6 and above\n",
    "\t - Pytorch, Transformers and All the stock Python ML Libraries\n",
    "\t - TPU enabled setup. This can also be executed over GPU but the code base will need some changes. \n",
    "\n",
    "\n",
    " - Script Objective:\n",
    "\t - The objective of this script is to fine tune **BertForTokenClassification**` to be able to identify the entites as per the given test dataset. The entities labled in the given dataset are as follows:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='section00'></a>\n",
    "### Installing packages for preparing the system\n",
    "\n",
    "We are installing 2 packages for the purposes of TPU execution and f1 metric score calculation respectively\n",
    "*You can skip this step if you already have these libraries installed in your environment*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 773
    },
    "colab_type": "code",
    "id": "pWbkd8Ld8MwL",
    "outputId": "b44f7ea3-2c0a-4e7c-f7ed-19f43d62de28"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  % Total    % Received % Xferd  Average Speed   Time    Time     Time  Current\n",
      "                                 Dload  Upload   Total   Spent    Left  Speed\n",
      "\r",
      "  0     0    0     0    0     0      0      0 --:--:-- --:--:-- --:--:--     0\r",
      "100  3727  100  3727    0     0  41876      0 --:--:-- --:--:-- --:--:-- 41876\n",
      "Updating TPU and VM. This may take around 2 minutes.\n",
      "Updating TPU runtime to pytorch-dev20200325 ...\n",
      "Done updating TPU runtime: <Response [200]>\n",
      "Uninstalling torch-1.5.0a0+d6149a7:\n",
      "  Successfully uninstalled torch-1.5.0a0+d6149a7\n",
      "Uninstalling torchvision-0.6.0a0+3c254fb:\n",
      "  Successfully uninstalled torchvision-0.6.0a0+3c254fb\n",
      "Copying gs://tpu-pytorch/wheels/torch-nightly+20200325-cp36-cp36m-linux_x86_64.whl...\n",
      "- [1 files][ 83.4 MiB/ 83.4 MiB]                                                \n",
      "Operation completed over 1 objects/83.4 MiB.                                     \n",
      "Copying gs://tpu-pytorch/wheels/torch_xla-nightly+20200325-cp36-cp36m-linux_x86_64.whl...\n",
      "\\ [1 files][114.5 MiB/114.5 MiB]                                                \n",
      "Operation completed over 1 objects/114.5 MiB.                                    \n",
      "Copying gs://tpu-pytorch/wheels/torchvision-nightly+20200325-cp36-cp36m-linux_x86_64.whl...\n",
      "/ [1 files][  2.5 MiB/  2.5 MiB]                                                \n",
      "Operation completed over 1 objects/2.5 MiB.                                      \n",
      "Processing ./torch-nightly+20200325-cp36-cp36m-linux_x86_64.whl\n",
      "Requirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from torch==nightly+20200325) (1.18.3)\n",
      "Requirement already satisfied: future in /usr/local/lib/python3.6/dist-packages (from torch==nightly+20200325) (0.16.0)\n",
      "\u001b[31mERROR: fastai 1.0.61 requires torchvision, which is not installed.\u001b[0m\n",
      "Installing collected packages: torch\n",
      "Successfully installed torch-1.5.0a0+d6149a7\n",
      "Processing ./torch_xla-nightly+20200325-cp36-cp36m-linux_x86_64.whl\n",
      "Installing collected packages: torch-xla\n",
      "  Found existing installation: torch-xla 1.6+e788e5b\n",
      "    Uninstalling torch-xla-1.6+e788e5b:\n",
      "      Successfully uninstalled torch-xla-1.6+e788e5b\n",
      "Successfully installed torch-xla-1.6+e788e5b\n",
      "Processing ./torchvision-nightly+20200325-cp36-cp36m-linux_x86_64.whl\n",
      "Requirement already satisfied: torch in /usr/local/lib/python3.6/dist-packages (from torchvision==nightly+20200325) (1.5.0a0+d6149a7)\n",
      "Requirement already satisfied: pillow>=4.1.1 in /usr/local/lib/python3.6/dist-packages (from torchvision==nightly+20200325) (7.0.0)\n",
      "Requirement already satisfied: six in /usr/local/lib/python3.6/dist-packages (from torchvision==nightly+20200325) (1.12.0)\n",
      "Requirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from torchvision==nightly+20200325) (1.18.3)\n",
      "Requirement already satisfied: future in /usr/local/lib/python3.6/dist-packages (from torch->torchvision==nightly+20200325) (0.16.0)\n",
      "Installing collected packages: torchvision\n",
      "Successfully installed torchvision-0.6.0a0+3c254fb\n",
      "Reading package lists... Done\n",
      "Building dependency tree       \n",
      "Reading state information... Done\n",
      "libomp5 is already the newest version (5.0.1-1).\n",
      "libopenblas-dev is already the newest version (0.2.20+ds-4).\n",
      "0 upgraded, 0 newly installed, 0 to remove and 29 not upgraded.\n"
     ]
    }
   ],
   "source": [
    "!curl -q https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py\n",
    "!python pytorch-xla-env-setup.py --apt-packages libomp5 libopenblas-dev\n",
    "!pip -q install seqeval"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "6ts5noAi0a6K"
   },
   "source": [
    "<a id='section01'></a>\n",
    "### Importing Python Libraries and preparing the environment\n",
    "\n",
    "At this step we will be importing the libraries and modules needed to run our script. Libraries are:\n",
    "* Pandas\n",
    "* Pytorch\n",
    "* Pytorch Utils for Dataset and Dataloader\n",
    "* Transformers\n",
    "* BERT Model and Tokenizer\n",
    "\n",
    "Followed by that we will preapre the device for TPU execeution. This configuration is needed if you want to leverage on onboard TPU. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "NYUqKiOZdR1H"
   },
   "outputs": [],
   "source": [
    "# Importing pytorch and the library for TPU execution\n",
    "\n",
    "import torch\n",
    "import torch_xla\n",
    "import torch_xla.core.xla_model as xm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "y3jTWir2cBlN"
   },
   "outputs": [],
   "source": [
    "# Importing stock ml libraries\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import transformers\n",
    "from torch.utils.data import Dataset, DataLoader, RandomSampler, SequentialSampler\n",
    "from transformers import BertForTokenClassification, BertTokenizer, BertConfig, BertModel\n",
    "\n",
    "# Preparing for TPU usage\n",
    "dev = xm.xla_device()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='section02'></a>\n",
    "### Importing and Pre-Processing the domain data\n",
    "\n",
    "We will be working with the data and preparing for fine tuning purposes. \n",
    "*Assuming that the `ner.csv` is already downloaded in your `data` folder*\n",
    "\n",
    "* Import the file in a dataframe and give it the headers as per the documentation.\n",
    "* Cleaning the file to remove the unwanted columns.\n",
    "* We will create a class `SentenceGetter` that will pull the words from the columns and create them into sentences\n",
    "* Followed by that we will create some additional lists and dict to keep the data that will be used for future processing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 212
    },
    "colab_type": "code",
    "id": "81kDZbz2cDn7",
    "outputId": "1312c879-5c24-4964-c6ab-ba1b9b189964"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "b'Skipping line 281837: expected 25 fields, saw 34\\n'\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>pos</th>\n",
       "      <th>sentence_idx</th>\n",
       "      <th>word</th>\n",
       "      <th>tag</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>NNS</td>\n",
       "      <td>1.0</td>\n",
       "      <td>Thousands</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>IN</td>\n",
       "      <td>1.0</td>\n",
       "      <td>of</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>NNS</td>\n",
       "      <td>1.0</td>\n",
       "      <td>demonstrators</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>VBP</td>\n",
       "      <td>1.0</td>\n",
       "      <td>have</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>VBN</td>\n",
       "      <td>1.0</td>\n",
       "      <td>marched</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   pos  sentence_idx           word tag\n",
       "0  NNS           1.0      Thousands   O\n",
       "1   IN           1.0             of   O\n",
       "2  NNS           1.0  demonstrators   O\n",
       "3  VBP           1.0           have   O\n",
       "4  VBN           1.0        marched   O"
      ]
     },
     "execution_count": 5,
     "metadata": {
      "tags": []
     },
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = pd.read_csv(\"./data/ner.csv\", encoding = \"ISO-8859-1\", error_bad_lines=False)\n",
    "dataset=df.drop(['Unnamed: 0', 'lemma', 'next-lemma', 'next-next-lemma', 'next-next-pos',\n",
    "       'next-next-shape', 'next-next-word', 'next-pos', 'next-shape',\n",
    "       'next-word', 'prev-iob', 'prev-lemma', 'prev-pos',\n",
    "       'prev-prev-iob', 'prev-prev-lemma', 'prev-prev-pos', 'prev-prev-shape',\n",
    "       'prev-prev-word', 'prev-shape', 'prev-word','shape'],axis=1)\n",
    "dataset.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "sdqhHeAqcLnO"
   },
   "outputs": [],
   "source": [
    "# Creating a class to pull the words from the columns and create them into sentences\n",
    "\n",
    "class SentenceGetter(object):\n",
    "    \n",
    "    def __init__(self, dataset):\n",
    "        self.n_sent = 1\n",
    "        self.dataset = dataset\n",
    "        self.empty = False\n",
    "        agg_func = lambda s: [(w,p, t) for w,p, t in zip(s[\"word\"].values.tolist(),\n",
    "                                                       s['pos'].values.tolist(),\n",
    "                                                        s[\"tag\"].values.tolist())]\n",
    "        self.grouped = self.dataset.groupby(\"sentence_idx\").apply(agg_func)\n",
    "        self.sentences = [s for s in self.grouped]\n",
    "    \n",
    "    def get_next(self):\n",
    "        try:\n",
    "            s = self.grouped[\"Sentence: {}\".format(self.n_sent)]\n",
    "            self.n_sent += 1\n",
    "            return s\n",
    "        except:\n",
    "            return None\n",
    "\n",
    "getter = SentenceGetter(dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Cqln3QJecNsJ"
   },
   "outputs": [],
   "source": [
    "# Creating new lists and dicts that will be used at a later stage for reference and processing\n",
    "\n",
    "tags_vals = list(set(dataset[\"tag\"].values))\n",
    "tag2idx = {t: i for i, t in enumerate(tags_vals)}\n",
    "sentences = [' '.join([s[0] for s in sent]) for sent in getter.sentences]\n",
    "labels = [[s[2] for s in sent] for sent in getter.sentences]\n",
    "labels = [[tag2idx.get(l) for l in lab] for lab in labels]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='section03'></a>\n",
    "### Preparing the Dataset and Dataloader\n",
    "\n",
    "We will start with defining few key variables that will be used later during the training/fine tuning stage.\n",
    "Followed by creation of Dataset class - This defines how the text is pre-processed before sending it to the neural network. We will also define the Dataloader that will feed  the data in batches to the neural network for suitable training and processing. \n",
    "Dataset and Dataloader are constructs of the PyTorch library for defining and controlling the data pre-processing and its passage to neural network. For further reading into Dataset and Dataloader read the [docs at PyTorch](https://pytorch.org/docs/stable/data.html)\n",
    "\n",
    "#### *CustomDataset* Dataset Class\n",
    "- This class is defined to accept the `tokenizer`, `sentences` and `labels` as input and generate tokenized output and tags that is used by the BERT model for training. \n",
    "- We are using the BERT tokenizer to tokenize the data in the `sentences` list for encoding. \n",
    "- The tokenizer uses the `encode_plus` method to perform tokenization and generate the necessary outputs, namely: `ids`, `attention_mask`\n",
    "- To read further into the tokenizer, [refer to this document](https://huggingface.co/transformers/model_doc/bert.html#berttokenizer)\n",
    "- `tags` is the encoded entity from the annonated dataset. \n",
    "- The *CustomDataset* class is used to create 2 datasets, for training and for validation.\n",
    "- *Training Dataset* is used to fine tune the model: **80% of the original data**\n",
    "- *Validation Dataset* is used to evaluate the performance of the model. The model has not seen this data during training. \n",
    "\n",
    "#### Dataloader\n",
    "- Dataloader is used to for creating training and validation dataloader that load data to the neural network in a defined manner. This is needed because all the data from the dataset cannot be loaded to the memory at once, hence the amount of dataloaded to the memory and then passed to the neural network needs to be controlled.\n",
    "- This control is achieved using the parameters such as `batch_size` and `max_len`.\n",
    "- Training and Validation dataloaders are used in the training and validation part of the flow respectively"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "kL0b1VIQcTAC"
   },
   "outputs": [],
   "source": [
    "# Defining some key variables that will be used later on in the training\n",
    "\n",
    "MAX_LEN = 200\n",
    "TRAIN_BATCH_SIZE = 32\n",
    "VALID_BATCH_SIZE = 16\n",
    "EPOCHS = 5\n",
    "LEARNING_RATE = 2e-05\n",
    "tokenizer = BertTokenizer.from_pretrained('bert-base-cased')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "IV72GFgq_ZYb"
   },
   "outputs": [],
   "source": [
    "class CustomDataset(Dataset):\n",
    "    def __init__(self, tokenizer, sentences, labels, max_len):\n",
    "        self.len = len(sentences)\n",
    "        self.sentences = sentences\n",
    "        self.labels = labels\n",
    "        self.tokenizer = tokenizer\n",
    "        self.max_len = max_len\n",
    "        \n",
    "    def __getitem__(self, index):\n",
    "        sentence = str(self.sentences[index])\n",
    "        inputs = self.tokenizer.encode_plus(\n",
    "            sentence,\n",
    "            None,\n",
    "            add_special_tokens=True,\n",
    "            max_length=self.max_len,\n",
    "            pad_to_max_length=True,\n",
    "            return_token_type_ids=True\n",
    "        )\n",
    "        ids = inputs['input_ids']\n",
    "        mask = inputs['attention_mask']\n",
    "        label = self.labels[index]\n",
    "        label.extend([4]*200)\n",
    "        label=label[:200]\n",
    "\n",
    "        return {\n",
    "            'ids': torch.tensor(ids, dtype=torch.long),\n",
    "            'mask': torch.tensor(mask, dtype=torch.long),\n",
    "            'tags': torch.tensor(label, dtype=torch.long)\n",
    "        } \n",
    "    \n",
    "    def __len__(self):\n",
    "        return self.len"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 67
    },
    "colab_type": "code",
    "id": "7BvniiyvX-rB",
    "outputId": "61fdc431-4544-4a19-a176-46f19ca58f0f"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FULL Dataset: 35177\n",
      "TRAIN Dataset: 28141\n",
      "TEST Dataset: 7036\n"
     ]
    }
   ],
   "source": [
    "# Creating the dataset and dataloader for the neural network\n",
    "\n",
    "train_percent = 0.8\n",
    "train_size = int(train_percent*len(sentences))\n",
    "# train_dataset=df.sample(frac=train_size,random_state=200).reset_index(drop=True)\n",
    "# test_dataset=df.drop(train_dataset.index).reset_index(drop=True)\n",
    "train_sentences = sentences[0:train_size]\n",
    "train_labels = labels[0:train_size]\n",
    "\n",
    "test_sentences = sentences[train_size:]\n",
    "test_labels = labels[train_size:]\n",
    "\n",
    "print(\"FULL Dataset: {}\".format(len(sentences)))\n",
    "print(\"TRAIN Dataset: {}\".format(len(train_sentences)))\n",
    "print(\"TEST Dataset: {}\".format(len(test_sentences)))\n",
    "\n",
    "training_set = CustomDataset(tokenizer, train_sentences, train_labels, MAX_LEN)\n",
    "testing_set = CustomDataset(tokenizer, test_sentences, test_labels, MAX_LEN)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "F05bou5zZYtV"
   },
   "outputs": [],
   "source": [
    "train_params = {'batch_size': TRAIN_BATCH_SIZE,\n",
    "                'shuffle': True,\n",
    "                'num_workers': 0\n",
    "                }\n",
    "\n",
    "test_params = {'batch_size': VALID_BATCH_SIZE,\n",
    "                'shuffle': True,\n",
    "                'num_workers': 0\n",
    "                }\n",
    "\n",
    "training_loader = DataLoader(training_set, **train_params)\n",
    "testing_loader = DataLoader(testing_set, **test_params)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='section04'></a>\n",
    "### Creating the Neural Network for Fine Tuning\n",
    "\n",
    "#### Neural Network\n",
    " - We will be creating a neural network with the `BERTClass`. \n",
    " - This network will have the `BertForTokenClassification` model. \n",
    " - The data will be fed to the `BertForTokenClassification` as defined in the dataset. \n",
    " - Final layer outputs is what will be used to calcuate the loss and to determine the accuracy of models prediction. \n",
    " - We will initiate an instance of the network called `model`. This instance will be used for training and then to save the final trained model for future inference. \n",
    " \n",
    "#### Loss Function and Optimizer\n",
    " - `Optimizer` is defined in the next cell.\n",
    " - We do not define any `Loss function` since the specified model already outputs `Loss` for a given input. \n",
    " - `Optimizer` is used to update the weights of the neural network to improve its performance.\n",
    " \n",
    "#### Further Reading\n",
    "- You can refer to my [Pytorch Tutorials](https://github.com/abhimishra91/pytorch-tutorials) to get an intuition of Loss Function and Optimizer.\n",
    "- [Pytorch Documentation for Loss Function](https://pytorch.org/docs/stable/nn.html#loss-functions)\n",
    "- [Pytorch Documentation for Optimizer](https://pytorch.org/docs/stable/optim.html)\n",
    "- Refer to the links provided on the top of the notebook to read more about `BertForTokenClassification`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "9vuIJrvSZble"
   },
   "outputs": [],
   "source": [
    "# Creating the customized model, by adding a drop out and a dense layer on top of distil bert to get the final output for the model. \n",
    "\n",
    "class BERTClass(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super(BERTClass, self).__init__()\n",
    "        self.l1 = transformers.BertForTokenClassification.from_pretrained('bert-base-cased', num_labels=18)\n",
    "        # self.l2 = torch.nn.Dropout(0.3)\n",
    "        # self.l3 = torch.nn.Linear(768, 200)\n",
    "    \n",
    "    def forward(self, ids, mask, labels):\n",
    "        output_1= self.l1(ids, mask, labels = labels)\n",
    "        # output_2 = self.l2(output_1[0])\n",
    "        # output = self.l3(output_2)\n",
    "        return output_1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "colab_type": "code",
    "collapsed": true,
    "id": "CflOeT2-ZoV6",
    "outputId": "8045068c-9aa0-48b8-cad1-8ac7c3b1e807"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "BertForTokenClassification(\n",
       "  (bert): BertModel(\n",
       "    (embeddings): BertEmbeddings(\n",
       "      (word_embeddings): Embedding(28996, 768, padding_idx=0)\n",
       "      (position_embeddings): Embedding(512, 768)\n",
       "      (token_type_embeddings): Embedding(2, 768)\n",
       "      (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "      (dropout): Dropout(p=0.1, inplace=False)\n",
       "    )\n",
       "    (encoder): BertEncoder(\n",
       "      (layer): ModuleList(\n",
       "        (0): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (1): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (2): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (3): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (4): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (5): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (6): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (7): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (8): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (9): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (10): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (11): BertLayer(\n",
       "          (attention): BertAttention(\n",
       "            (self): BertSelfAttention(\n",
       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "            (output): BertSelfOutput(\n",
       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "              (dropout): Dropout(p=0.1, inplace=False)\n",
       "            )\n",
       "          )\n",
       "          (intermediate): BertIntermediate(\n",
       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          )\n",
       "          (output): BertOutput(\n",
       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (pooler): BertPooler(\n",
       "      (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "      (activation): Tanh()\n",
       "    )\n",
       "  )\n",
       "  (dropout): Dropout(p=0.1, inplace=False)\n",
       "  (classifier): Linear(in_features=768, out_features=18, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 16,
     "metadata": {
      "tags": []
     },
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = BERTClass()\n",
    "model.to(dev)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "DN_u9NC5aaa_"
   },
   "outputs": [],
   "source": [
    "optimizer = torch.optim.Adam(params =  model.parameters(), lr=LEARNING_RATE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='section05'></a>\n",
    "### Fine Tuning the Model\n",
    "\n",
    "After all the effort of loading and preparing the data and datasets, creating the model and defining its loss and optimizer. This is probably the easier steps in the process. \n",
    "\n",
    "Here we define a training function that trains the model on the training dataset created above, specified number of times (EPOCH), An epoch defines how many times the complete data will be passed through the network. \n",
    "\n",
    "Following events happen in this function to fine tune the neural network:\n",
    "- The dataloader passes data to the model based on the batch size. \n",
    "- Subsequent output from the model and the actual category are compared to calculate the loss. \n",
    "- Loss value is used to optimize the weights of the neurons in the network.\n",
    "- After every 500 steps the loss value is printed in the console.\n",
    "\n",
    "As you can see just in 1 epoch by the final step the model was working with a miniscule loss of 0.08503091335296631 i.e. the output is extremely close to the actual output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "8aQ6WCk2a-Vd"
   },
   "outputs": [],
   "source": [
    "def train(epoch):\n",
    "    model.train()\n",
    "    for _,data in enumerate(training_loader, 0):\n",
    "        ids = data['ids'].to(dev, dtype = torch.long)\n",
    "        mask = data['mask'].to(dev, dtype = torch.long)\n",
    "        targets = data['tags'].to(dev, dtype = torch.long)\n",
    "\n",
    "        loss = model(ids, mask, labels = targets)[0]\n",
    "\n",
    "        # optimizer.zero_grad()\n",
    "        if _%500==0:\n",
    "            print(f'Epoch: {epoch}, Loss:  {loss.item()}')\n",
    "        \n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        xm.optimizer_step(optimizer)\n",
    "        xm.mark_step() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 689
    },
    "colab_type": "code",
    "collapsed": true,
    "id": "50oMGTe0bvl0",
    "outputId": "292ede16-6da3-460f-d174-f7ee10c729e0"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 0, Loss:  0.21416641771793365\n",
      "Epoch: 0, Loss:  0.08791390806436539\n",
      "Epoch: 0, Loss:  0.1277497559785843\n",
      "Epoch: 0, Loss:  0.25511449575424194\n",
      "Epoch: 0, Loss:  0.11072967946529388\n",
      "Epoch: 0, Loss:  0.1202322468161583\n",
      "Epoch: 0, Loss:  0.16198261082172394\n",
      "Epoch: 0, Loss:  0.31682807207107544\n",
      "Epoch: 1, Loss:  0.09211093187332153\n",
      "Epoch: 1, Loss:  0.15079179406166077\n",
      "Epoch: 1, Loss:  0.1959223747253418\n",
      "Epoch: 1, Loss:  0.09143798053264618\n",
      "Epoch: 1, Loss:  0.29411888122558594\n",
      "Epoch: 1, Loss:  0.11708520352840424\n",
      "Epoch: 1, Loss:  0.11245028674602509\n",
      "Epoch: 1, Loss:  0.14728033542633057\n",
      "Epoch: 2, Loss:  0.1607980579137802\n",
      "Epoch: 2, Loss:  0.08060580492019653\n",
      "Epoch: 2, Loss:  0.14363577961921692\n",
      "Epoch: 2, Loss:  0.12225533276796341\n",
      "Epoch: 2, Loss:  0.10335233807563782\n",
      "Epoch: 2, Loss:  0.04923604056239128\n",
      "Epoch: 2, Loss:  0.09237729012966156\n",
      "Epoch: 2, Loss:  0.12473192811012268\n",
      "Epoch: 3, Loss:  0.09085617959499359\n",
      "Epoch: 3, Loss:  0.09351193159818649\n",
      "Epoch: 3, Loss:  0.06728512048721313\n",
      "Epoch: 3, Loss:  0.1666068434715271\n",
      "Epoch: 3, Loss:  0.19255675375461578\n",
      "Epoch: 3, Loss:  0.16131675243377686\n",
      "Epoch: 3, Loss:  0.15462705492973328\n",
      "Epoch: 3, Loss:  0.18679684400558472\n",
      "Epoch: 4, Loss:  0.11378277838230133\n",
      "Epoch: 4, Loss:  0.025372153148055077\n",
      "Epoch: 4, Loss:  0.08231651782989502\n",
      "Epoch: 4, Loss:  0.2682102620601654\n",
      "Epoch: 4, Loss:  0.05264609679579735\n",
      "Epoch: 4, Loss:  0.056522976607084274\n",
      "Epoch: 4, Loss:  0.15710100531578064\n",
      "Epoch: 4, Loss:  0.08503091335296631\n"
     ]
    }
   ],
   "source": [
    "for epoch in range(5):\n",
    "    train(epoch)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id='section06'></a>\n",
    "### Validating the Model\n",
    "\n",
    "During the validation stage we pass the unseen data(Testing Dataset) to the model. This step determines how good the model performs on the unseen data. \n",
    "\n",
    "This unseen data is the 30% of `ner.csv` which was seperated during the Dataset creation stage. \n",
    "During the validation stage the weights of the model are not updated. Only the final output is compared to the actual value. This comparison is then used to calcuate the accuracy of the model. \n",
    "\n",
    "The metric used for measuring the performance of model for these problem statements is called F1 score. We will create a helper function for helping us with f1 score calcuation and also import a library for the same. `seqeval`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "6OckC0XNkWWm",
    "outputId": "d27682ee-c34f-4811-d791-0c922afa8b05"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  Building wheel for seqeval (setup.py) ... \u001b[?25l\u001b[?25hdone\n"
     ]
    }
   ],
   "source": [
    "from seqeval.metrics import f1_score\n",
    "\n",
    "def flat_accuracy(preds, labels):\n",
    "    flat_preds = np.argmax(preds, axis=2).flatten()\n",
    "    flat_labels = labels.flatten()\n",
    "    return np.sum(flat_preds == flat_labels)/len(flat_labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "9zwFzzBriLMO"
   },
   "outputs": [],
   "source": [
    "def valid(model, testing_loader):\n",
    "    model.eval()\n",
    "    eval_loss = 0; eval_accuracy = 0\n",
    "    n_correct = 0; n_wrong = 0; total = 0\n",
    "    predictions , true_labels = [], []\n",
    "    nb_eval_steps, nb_eval_examples = 0, 0\n",
    "    with torch.no_grad():\n",
    "        for _, data in enumerate(testing_loader, 0):\n",
    "            ids = data['ids'].to(dev, dtype = torch.long)\n",
    "            mask = data['mask'].to(dev, dtype = torch.long)\n",
    "            targets = data['tags'].to(dev, dtype = torch.long)\n",
    "\n",
    "            output = model(ids, mask, labels=targets)\n",
    "            loss, logits = output[:2]\n",
    "            logits = logits.detach().cpu().numpy()\n",
    "            label_ids = targets.to('cpu').numpy()\n",
    "            predictions.extend([list(p) for p in np.argmax(logits, axis=2)])\n",
    "            true_labels.append(label_ids)\n",
    "            accuracy = flat_accuracy(logits, label_ids)\n",
    "            eval_loss += loss.mean().item()\n",
    "            eval_accuracy += accuracy\n",
    "            nb_eval_examples += ids.size(0)\n",
    "            nb_eval_steps += 1\n",
    "        eval_loss = eval_loss/nb_eval_steps\n",
    "        print(\"Validation loss: {}\".format(eval_loss))\n",
    "        print(\"Validation Accuracy: {}\".format(eval_accuracy/nb_eval_steps))\n",
    "        pred_tags = [tags_vals[p_i] for p in predictions for p_i in p]\n",
    "        valid_tags = [tags_vals[l_ii] for l in true_labels for l_i in l for l_ii in l_i]\n",
    "        print(\"F1-Score: {}\".format(f1_score(pred_tags, valid_tags)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# To get the results on the validation set. This data is not seen by the model\n",
    "\n",
    "valid(model, testing_loader)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "TPU",
  "colab": {
   "collapsed_sections": [],
   "name": "transformers_ner.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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}