F3-P2D-NER-CUAD-v1-Fine-Tune-Transformer.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Paper-2-Data Workflow for Data Extraction - CUADv1 - Fine Tune Transformer"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Sources of information, code and discussions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "1. The foundation workflow is from Hugging Face's Token Classification example hosted on Colab [here][1]\n",
    "2. The models are base models, each using a downstream token clasification task, example [here][2]\n",
    "\n",
    "[1]: https://colab.research.google.com/github/huggingface/notebooks/blob/master/examples/token_classification.ipynb\n",
    "[2]: https://huggingface.co/roberta-base"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Initialize Environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os, re, math, random, json, string\n",
    "\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from IPython.display import display, HTML\n",
    "import wandb\n",
    "\n",
    "import transformers\n",
    "from transformers import AutoTokenizer, AutoModelForTokenClassification, TrainingArguments, Trainer\n",
    "from transformers import TrainerCallback, AdamW, get_cosine_schedule_with_warmup\n",
    "from transformers import DataCollatorForTokenClassification, PreTrainedModel, RobertaTokenizerFast\n",
    "\n",
    "from datasets import load_dataset, ClassLabel, Sequence, load_metric\n",
    "\n",
    "from seqeval.metrics import accuracy_score, precision_score, recall_score, f1_score, classification_report"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33mnative\u001b[0m (use `wandb login --relogin` to force relogin)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Need to log in to weights and biases in the command line using: wandb login\n",
    "wandb.login()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configuration\n",
    "\n",
    "* BATCH_SIZES - These are batch sizes for each fold. For maximum speed, it is best to use the largest batch size your GPU or TPU memory allows.\n",
    "* EPOCHS - These are maximum epochs. Note that each fold, the best epoch model is saved and used. So if epochs is too large, it won't matter. Consider early stopping in Callbacks.\n",
    "* MODEL_CHECKPOINT - The name of the transformer model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hugging Face model references for Transformer library\n",
    "models = dict(\n",
    "    ROBERTA = \"roberta-base\",\n",
    "    DISTILBERT_U = \"distilbert-base-uncased\",\n",
    "    DISTILBERT_C = \"distilbert-base-cased\",\n",
    "    DEBERTA_V2_XL = \"microsoft/deberta-v2-xlarge\",\n",
    "    DEBERTA_V2_XXL = \"microsoft/deberta-v2-xxlarge\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Logging date for w&b\n",
    "from datetime import date\n",
    "today = date.today()\n",
    "log_date = today.strftime(\"%d-%m-%Y\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "_cell_guid": "79c7e3d0-c299-4dcb-8224-4455121ee9b0",
    "_uuid": "d629ff2d2480ee46fbb7e2d37f6b5fab8052498a"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "env: WANDB_PROJECT='P2D-NER-2021'\n",
      "env: WANDB_LOG_MODEL=false\n"
     ]
    }
   ],
   "source": [
    "# LOAD OR TRAIN MODEL\n",
    "TRAIN = 1 # 1 to TRAIN WEIGHTS or 0 to LOAD WEIGHTS\n",
    "\n",
    "# TRAIN/VALIDATION SPLIT\n",
    "TRAIN_SPLIT = 0.90\n",
    "\n",
    "# RANDOM SEED FOR REPRODUCIBILITY\n",
    "RANDOM_SEED = 42\n",
    "\n",
    "# BATCH SIZE\n",
    "# TRY 4, 8, 16, 32, 64, 128, 256. REDUCE IF OOM ERROR, HIGHER FOR TPUS\n",
    "BATCH_SIZES = 1\n",
    "\n",
    "# EPOCHS - TRANSFORMERS ARE TYPICALLY FINE-TUNED BETWEEN 1 AND 3 EPOCHS \n",
    "EPOCHS = 10\n",
    "\n",
    "# WHICH PRE-TRAINED TRANSFORMER TO FINE-TUNE?\n",
    "MODEL_CHECKPOINT = models['DEBERTA_V2_XL']\n",
    "\n",
    "# SPECIFY THE WEIGHTS AND BIASES PROJECT NAME\n",
    "%env WANDB_PROJECT = 'P2D-NER-2021' \n",
    "\n",
    "# DETERMINE WHETHER TO SAVE THE MODEL IN THE 100GB OF FREE W&B STORAGE\n",
    "%env WANDB_LOG_MODEL = false "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step1: File and dataset handling\n",
    "Data cleaning, annotations and  formatting has already been done, tokenized to seperate words, tagged using the IOB format and serialized using the Pandas df.to_json() function using the orient=\"table\" parameter to a JSONL file. \n",
    "\n",
    "Here we load in the dataset with this JSON format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "FEATURE_CLASS_LABELS = \"feature_class_labels.json\"\n",
    "DATA_FILE = 'cuad-v1-annotated.json'\n",
    "TEMP_MODEL_OUTPUT_DIR = 'temp_model_output_dir'\n",
    "SAVED_MODEL = f\"p2d-NER-Fine-Tune-Transformer-{MODEL_CHECKPOINT}\" # Change for notebook version"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using custom data configuration default-075146409b74ff61\n",
      "Reusing dataset json (/home/phil/.cache/huggingface/datasets/json/default-075146409b74ff61/0.0.0/83d5b3a2f62630efc6b5315f00f20209b4ad91a00ac586597caee3a4da0bef02)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DatasetDict({\n",
      "    train: Dataset({\n",
      "        features: ['id', 'ner_tags', 'split_tokens'],\n",
      "        num_rows: 282\n",
      "    })\n",
      "    test: Dataset({\n",
      "        features: ['id', 'ner_tags', 'split_tokens'],\n",
      "        num_rows: 32\n",
      "    })\n",
      "})\n"
     ]
    }
   ],
   "source": [
    "data_files = DATA_FILE\n",
    "datasets = load_dataset('json', data_files=data_files, field='data')\n",
    "\n",
    "# Create train and validation datasets\n",
    "datasets = datasets['train'].train_test_split(test_size=1-TRAIN_SPLIT, seed=RANDOM_SEED)\n",
    "print(datasets)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 B-AGMT_DATE\n",
      "1 B-DOC_NAME\n",
      "2 B-PARTY\n",
      "3 I-AGMT_DATE\n",
      "4 I-DOC_NAME\n",
      "5 I-PARTY\n",
      "6 O\n"
     ]
    }
   ],
   "source": [
    "# Open the label list created in pre-processing corresponding to the ner_tag indices\n",
    "with open(FEATURE_CLASS_LABELS, 'r') as f:\n",
    "    label_list = json.load(f)\n",
    "\n",
    "for n in range(len(label_list)):\n",
    "    print(n, label_list[n])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>id</th>\n",
       "      <th>ner_tags</th>\n",
       "      <th>split_tokens</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>15754</td>\n",
       "      <td>[6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 1, 4, 4, 4, 4, 4, 6, 6, 6, 0, 3, 3, 3, 6, 6, 6, 2, 5, 5, 5, 6, 2, 5, 5, 5, 5, 6, 6, 2, 5, 5, 5, 5, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, ...]</td>\n",
       "      <td>[Exhibit, 10.3, [, *, ], Certain, information, in, this, document, has, been, excluded, pursuant, to, Regulation, S, -, K, ,, Item, 601(b)(10, ), ., Such, excluded, information, is, not, material, and, would, likely, cause, competitive, harm, to, the, registrant, if, publicly, disclosed, ., Execution, Copy, LICENSE, ,, DEVELOPMENT, AND, COMMERCIALIZATION, AGREEMENT, DATED, AS, OF, FEBRUARY, 4, ,, 2020, BY, AND, BETWEEN, XENCOR, ,, INC, ., AND, AIMMUNE, THERAPEUTICS, ,, INC, ., Source, :, AIMMUNE, THERAPEUTICS, ,, INC, ., ,, 8, -, K, ,, 2/5/2020, TABLE, OF, CONTENTS, Page, ARTICLE, 1, Definitions, 1, ARTICLE, 2, Licenses, 13, ARTICLE, 3, Development, 16, ...]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>15721</td>\n",
       "      <td>[6, 6, 6, 1, 4, 2, 5, 5, 5, 6, 2, 5, 5, 5, 6, 2, 5, 5, 5, 5, 6, 6, 6, 6, 6, 6, 6, 6, 2, 5, 5, 5, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 2, 5, 5, 5, 6, 6, 6, 6, 6, 6, 2, 5, 5, 5, 5, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 0, 3, 3, 3, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, ...]</td>\n",
       "      <td>[EXHIBIT, 10(d, ), PROMOTION, AGREEMENT, ASHWORTH, ,, INC, ., ,, JAMES, W., NANTZ, III, AND, NANTZ, COMMUNICATIONS, ,, INC, ., THIS, AGREEMENT, is, entered, into, by, and, among, ASHWORTH, ,, INC, ., (, The, \", Company, \", or, \", Ashworth, \", ), ,, JAMES, W., NANTZ, III, (, \", Nantz, \", ), and, NANTZ, COMMUNICATIONS, ,, INC, ., (, \", Nantz, Communications, \", ), ,, effective, as, of, June, 1, ,, 1998, ., WHEREAS, ,, the, Company, desires, to, retain, Nantz, Communications, and, Nantz, to, provide, certain, promotional, and, other, services, and, Nantz, Communications, and, Nantz, are, willing, to, provide, ...]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>15579</td>\n",
       "      <td>[6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 1, 4, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 0, 3, 3, 3, 3, 3, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 2, 5, 5, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 6, 2, 5, 5, 6, 6, ...]</td>\n",
       "      <td>[CONFIDENTIAL, TREATMENT, HAS, BEEN, REQUESTED, FOR, PORTIONS, OF, THIS, EXHIBIT, ., THE, CONFIDENTIAL, PORTIONS, HAVE, BEEN, REDACTED, AND, ARE, DENOTED, BY, [, *, ], ., THE, CONFIDENTIAL, PORTIONS, HAVE, BEEN, SEPARATELY, FILED, WITH, THE, SECURITIES, AND, EXCHANGE, COMMISSION, ., SPONSORSHIP, AGREEMENT, This, agreement, (, \", Agreement, \", ), is, entered, into, as, of, the, 15th, day, of, December, ,, 1997, (, \", Effective, Date, \", ), ,, by, and, between, Excite, ,, Inc., ,, a, California, corporation, ,, located, at, 555, Broadway, ,, Redwood, City, ,, California, 94063, (, \", Excite, \", ), ,, and, NetGrocer, ,, Inc., a, Delaware, ...]</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Check some random samples to ensure data loaded as expected:\n",
    "def show_random_elements(dataset, num_examples=1):\n",
    "    assert num_examples <= len(dataset), \"Can't pick more elements than there are in the dataset.\"\n",
    "    picks = []\n",
    "    for _ in range(num_examples):\n",
    "        pick = random.randint(0, len(dataset)-1)\n",
    "        while pick in picks:\n",
    "            pick = random.randint(0, len(dataset)-1)\n",
    "        picks.append(pick)\n",
    "    \n",
    "    df = pd.DataFrame(dataset[picks])\n",
    "    for column, typ in dataset.features.items():\n",
    "        if isinstance(typ, ClassLabel):\n",
    "            df[column] = df[column].transform(lambda i: typ.names[i])\n",
    "        elif isinstance(typ, Sequence) and isinstance(typ.feature, ClassLabel):\n",
    "            df[column] = df[column].transform(lambda x: [typ.feature.names[i] for i in x])\n",
    "    display(HTML(df.to_html()))\n",
    "\n",
    "show_random_elements(datasets[\"train\"], num_examples=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2: Preprocessing Data - Tokenization\n",
    "Before we can feed those texts to our model, we need to preprocess them specifically for the pre-trained model that we are using. Even though we have already tokenized and split our words into a list, each model will have it's own further method of tokenization to match the dictionary for each specific model. \n",
    "\n",
    "This is done by a 🤗 Transformers Tokenizer which will (as the name indicates) tokenize the inputs (including converting the tokens to their corresponding IDs in the pretrained vocabulary) and put it in a format the model expects, as well as generate the other inputs that model requires.\n",
    "\n",
    "To do all of this, we instantiate our tokenizer with the AutoTokenizer.from_pretrained method, which will ensure:\n",
    " - we get a tokenizer that corresponds to the model architecture we want to use,\n",
    " - we download the vocabulary used when pretraining this specific checkpoint.\n",
    " \n",
    "The exception here is for Roberta Base where we specifically call the tokenizer for this model, due to inconsistencies in the Hugging Face library for this model.\n",
    "\n",
    "The vocabulary will be cached, so it's not downloaded again the next time we run the cell.\n",
    "\n",
    "If, as is the case here, inputs have already been split into words,we pass the list of words to the tokenzier with the argument \"is_split_into_words=True\"\n",
    "\n",
    "Some tokens will be split into subtokens if the token is not in the model dictionary. This means that we need to do some processing on our labels as the input ids returned by the tokenizer are longer than the lists of labels our dataset contain, first because some special tokens might be added (eg a [CLS] and a [SEP]) and then because of those possible splits of words in multiple tokens.\n",
    "\n",
    "Some tokenizers returns outputs that have a word_ids method which can help us. Otherwise we have to build our own, as is the case for DeBERTa for example.\n",
    "\n",
    "We will return a list with the same number of elements as our processed input ids, mapping special tokens to None and all other tokens to their respective word. This way, we can align the labels with the processed input ids.\n",
    "\n",
    "Here we set the labels of all special tokens to -100 (the index that is ignored by PyTorch) and the labels of all other tokens to the label of the word they come from.\n",
    "\n",
    "Another strategy is to set the label only on the first token obtained from a given word, and give a label of -100 to the other subtokens from the same word. Just change the value of the following flag : \"label_all_tokens = True/False\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    }
   ],
   "source": [
    "# Instantiate the tokenizer\n",
    "#For RoBERTa-base, need to use RobertaTokenizerFast with add_prefix_space=True to use it with pretokenized inputs.\n",
    "# SentencePiece will need to be installed for DeBERTa v2: pip install sentencepiece\n",
    "\n",
    "if MODEL_CHECKPOINT == models['ROBERTA']:\n",
    "    tokenizer = RobertaTokenizerFast.from_pretrained(\"roberta-base\", add_prefix_space=True)\n",
    "else:\n",
    "    tokenizer = AutoTokenizer.from_pretrained(MODEL_CHECKPOINT)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def word_id_func(input_ids, print_labs=False):\n",
    "    tokens = tokenizer.convert_ids_to_tokens(input_ids)\n",
    "    \n",
    "    word_ids = []\n",
    "    i=0\n",
    "    spec_toks = ['[CLS]', '[SEP]', '[PAD]']\n",
    "    for t in tokens:\n",
    "        if t in spec_toks:\n",
    "            word_ids.append(-100)\n",
    "            print(t, i) if print_labs else None\n",
    "        elif t.startswith('▁'):\n",
    "            i += 1\n",
    "            word_ids.append(i)\n",
    "            print(t, i) if print_labs else None\n",
    "        else:\n",
    "            word_ids.append(i)\n",
    "            print(t, i) if print_labs else None\n",
    "        print(\"Total:\", i) if print_labs else None\n",
    "    return word_ids\n",
    "\n",
    "def tokenize_and_align_labels(examples, label_all_tokens=False):\n",
    "    tokenized_inputs = tokenizer(examples[\"split_tokens\"],\n",
    "                                 truncation=True,\n",
    "                                 is_split_into_words=True)\n",
    "    labels = []\n",
    "    for i, label in enumerate(examples[\"ner_tags\"]):\n",
    "        word_ids = tokenized_inputs.word_ids(batch_index=i)\n",
    "        previous_word_idx = None\n",
    "        label_ids = []\n",
    "        for word_idx in word_ids:\n",
    "            # Special tokens have a word id that is None. We set the label to -100 so they are automatically\n",
    "            # ignored in the loss function.\n",
    "            if word_idx is None:\n",
    "                label_ids.append(-100)\n",
    "            # We set the label for the first token of each word.\n",
    "            elif word_idx != previous_word_idx:\n",
    "                label_ids.append(label[word_idx])\n",
    "            # For the other tokens in a word, we set the label to either the current label or -100, depending on\n",
    "            # the label_all_tokens flag.\n",
    "            else:\n",
    "                label_ids.append(label[word_idx] if label_all_tokens else -100)\n",
    "            previous_word_idx = word_idx\n",
    "        labels.append(label_ids)\n",
    "\n",
    "    tokenized_inputs[\"labels\"] = labels\n",
    "    return tokenized_inputs\n",
    "\n",
    "def tokenize_and_align_labels_deberta(examples, label_all_tokens=False):\n",
    "    tokenized_inputs = tokenizer(examples[\"split_tokens\"],\n",
    "                                 truncation=True,\n",
    "                                 is_split_into_words=True)\n",
    "    labels = []\n",
    "    word_ids_list = []\n",
    "    for input_ids in tokenized_inputs[\"input_ids\"]:\n",
    "        wids = word_id_func(input_ids, print_labs=False)\n",
    "        word_ids_list.append(wids)\n",
    "    \n",
    "    for i, label in enumerate(examples[\"ner_tags\"]):\n",
    "        word_ids = word_ids_list[i]\n",
    "        previous_word_idx = None\n",
    "        label_ids = []\n",
    "        for word_idx in word_ids:\n",
    "            # Special tokens have a word id that is None. We set the label to -100 so they are automatically\n",
    "            # ignored in the loss function.\n",
    "            if word_idx == -100:\n",
    "                label_ids.append(-100)\n",
    "            #We set the label for the first token of each word.\n",
    "            elif word_idx != previous_word_idx:\n",
    "                label_ids.append(label[word_idx-1])\n",
    "            # For the other tokens in a word, we set the label to either the current label or -100, depending on\n",
    "            # the label_all_tokens flag.\n",
    "            else:\n",
    "                label_ids.append(label[word_idx-1] if label_all_tokens else -100)\n",
    "            previous_word_idx = word_idx\n",
    "        labels.append(label_ids)\n",
    "\n",
    "    tokenized_inputs[\"labels\"] = labels\n",
    "    return tokenized_inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1e8c9bfafaf54eaabe0b1e9644ad0a61",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, max=1.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "c00deced44ed4622bf0388f25cb786cb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "HBox(children=(FloatProgress(value=0.0, max=1.0), HTML(value='')))"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# To apply this function on all the words and labels in our dataset,\n",
    "# we just use the map method of our dataset object we created earlier.\n",
    "# This will apply the function on all the elements of all the splits in dataset, so our training, \n",
    "# validation and testing data will be preprocessed in one single command.\n",
    "\n",
    "# 🤗 Datasets warns you when it uses cached files, you can pass load_from_cache_file=False in the\n",
    "# call to map to not use the cached files and force the preprocessing to be applied again.\n",
    "if MODEL_CHECKPOINT == models['DEBERTA_V2_XL'] or MODEL_CHECKPOINT == models['DEBERTA_V2_XXL']:\n",
    "    tokenize_and_align_labels = tokenize_and_align_labels_deberta\n",
    "\n",
    "tokenized_datasets = datasets.map(tokenize_and_align_labels, batched=True, load_from_cache_file=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 3: Build Model\n",
    "Since all our tasks are about token classification, we use the AutoModelForTokenClassification class. Like with the tokenizer, the from_pretrained method will download and cache the model for us. The only thing we have to specify is the number of labels for our problem (which we can get from the features, as seen before).\n",
    "\n",
    "The warning is telling us we are throwing away some weights (the vocab_transform and vocab_layer_norm layers) and randomly initializing some other (the pre_classifier and classifier layers). This is absolutely normal in this case, because we are removing the head used to pretrain the model on a masked language modeling objective and replacing it with a new head for which we don't have pretrained weights, so the library warns us we should fine-tune this model before using it for inference, which is exactly what we are going to do."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at microsoft/deberta-v2-xlarge were not used when initializing DebertaV2ForTokenClassification: ['lm_predictions.lm_head.bias', 'deberta.embeddings.position_embeddings.weight', 'lm_predictions.lm_head.LayerNorm.bias', 'lm_predictions.lm_head.LayerNorm.weight', 'lm_predictions.lm_head.dense.weight', 'lm_predictions.lm_head.dense.bias']\n",
      "- This IS expected if you are initializing DebertaV2ForTokenClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing DebertaV2ForTokenClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some weights of DebertaV2ForTokenClassification were not initialized from the model checkpoint at microsoft/deberta-v2-xlarge and are newly initialized: ['classifier.bias', 'classifier.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "model = AutoModelForTokenClassification.from_pretrained(MODEL_CHECKPOINT, num_labels=len(label_list))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Training Schedule\n",
    "This is a common train schedule for transfer learning. The learning rate starts at zero, to initially preserve the pre-trained weights, then increases to a maximum, then reduces using a cosine exponential curve to attempt to find the global optima.\n",
    "\n",
    "Changing the schedule and/or learning rates is a popular way to experiment to find good model performance. Note how the learning rate max is larger with larger batches sizes. This is a good practice to follow.\n",
    "\n",
    "Weight decay is the amount of L2 regularization to force into the model's optimizer to make it work harder and offset any tendancy for the model to overfit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The maximum learning rate is:  7.5e-06\n"
     ]
    }
   ],
   "source": [
    "#Optimizer\n",
    "learning_rate = 0.0000075\n",
    "lr_max = learning_rate * BATCH_SIZES\n",
    "weight_decay = 0.05\n",
    "\n",
    "optimizer = AdamW(\n",
    "    model.parameters(),\n",
    "    lr=lr_max,\n",
    "    weight_decay=weight_decay)\n",
    "\n",
    "print(\"The maximum learning rate is: \",lr_max)\n",
    "\n",
    "# Learning Rate Schedule\n",
    "num_train_samples = len(datasets[\"train\"])\n",
    "warmup_ratio = 0.2 # Percentage of total steps to go from zero to max learning rate\n",
    "num_cycles=0.8 # The cosine exponential rate\n",
    "\n",
    "num_training_steps = num_train_samples*EPOCHS/BATCH_SIZES\n",
    "num_warmup_steps = num_training_steps*warmup_ratio\n",
    "\n",
    "lr_sched = get_cosine_schedule_with_warmup(optimizer=optimizer,\n",
    "                                           num_warmup_steps=num_warmup_steps,\n",
    "                                           num_training_steps = num_training_steps,\n",
    "                                           num_cycles=num_cycles)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "To instantiate a Trainer, we will need to define three more things. The most important is the TrainingArguments, which is a class that contains all the attributes to customize the training. It requires one folder name, which will be used to save the checkpoints of the model, and all other arguments are optional."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "args = TrainingArguments(output_dir = TEMP_MODEL_OUTPUT_DIR,\n",
    "                         evaluation_strategy = \"epoch\",\n",
    "                         learning_rate=lr_max,\n",
    "                         per_device_train_batch_size=BATCH_SIZES,\n",
    "                         per_device_eval_batch_size=BATCH_SIZES,\n",
    "                         num_train_epochs=EPOCHS,\n",
    "                         weight_decay=weight_decay,\n",
    "                         lr_scheduler_type = 'cosine',\n",
    "                         warmup_ratio=warmup_ratio,\n",
    "                         logging_strategy=\"epoch\",\n",
    "                         save_strategy=\"epoch\",\n",
    "                         seed=RANDOM_SEED,\n",
    "                         report_to = 'wandb', # enable logging to W&B\n",
    "                         run_name = MODEL_CHECKPOINT+\"-\"+log_date,\n",
    "                         metric_for_best_model=\"f1\",\n",
    "                         load_best_model_at_end = True)   # name of the W&B run (optional)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we will need a data collator that will batch our processed examples together while applying padding to make them all the same size (each pad will be padded to the length of its longest example). There is a data collator for this task in the Transformers library, that not only pads the inputs, but also the labels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_collator = DataCollatorForTokenClassification(tokenizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 4: Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The last thing to define for our Trainer is how to compute the metrics from the predictions. Here we will load the seqeval metrics (which are commonly used to evaluate results on the benchmark CONLL dataset). https://github.com/chakki-works/seqeval\n",
    "\n",
    "Note - Either BILOU or IOB tags can be used. Whilst BILOU provides for more features, research suggests using the simpler IOB for token classification shouldn't impact accuracy. \n",
    "\n",
    "So we will need to do a bit of post-processing on our predictions:\n",
    " - select the predicted index (with the maximum logit) for each token\n",
    " - convert it to its string label\n",
    " - ignore everywhere we set a label of -100\n",
    "\n",
    "The following function does all this post-processing on the result of Trainer.evaluate (which is a namedtuple containing predictions and labels) before applying the metric:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_metrics(p):\n",
    "    predictions, labels = p\n",
    "    predictions = np.argmax(predictions, axis=2)\n",
    "\n",
    "    # Remove ignored index (special tokens)\n",
    "    true_predictions = [\n",
    "        [label_list[p] for (p, l) in zip(prediction, label) if l != -100]\n",
    "        for prediction, label in zip(predictions, labels)]\n",
    "    true_labels = [\n",
    "        [label_list[l] for (p, l) in zip(prediction, label) if l != -100]\n",
    "        for prediction, label in zip(predictions, labels)]\n",
    "\n",
    "    # Define the metric parameters\n",
    "    overall_precision = precision_score(true_labels, true_predictions, zero_division=1)\n",
    "    overall_recall = recall_score(true_labels, true_predictions, zero_division=1)\n",
    "    overall_f1 = f1_score(true_labels, true_predictions, zero_division=1)\n",
    "    overall_accuracy = accuracy_score(true_labels, true_predictions)\n",
    "    \n",
    "    # Return a dictionary with the calculated metrics\n",
    "    return {\n",
    "        \"precision\": overall_precision,\n",
    "        \"recall\": overall_recall,\n",
    "        \"f1\": overall_f1,\n",
    "        \"accuracy\": overall_accuracy,}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define and instantiate the Trainer...\n",
    "trainer = Trainer(\n",
    "                model=model,\n",
    "                args=args,\n",
    "                train_dataset=tokenized_datasets[\"train\"],\n",
    "                eval_dataset=tokenized_datasets[\"test\"],\n",
    "                data_collator=data_collator,\n",
    "                tokenizer=tokenizer,\n",
    "                compute_metrics=compute_metrics,\n",
    "                optimizers=(optimizer, lr_sched)\n",
    "                )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "                Tracking run with wandb version 0.10.30<br/>\n",
       "                Syncing run <strong style=\"color:#cdcd00\">microsoft/deberta-v2-xlarge-19-05-2021</strong> to <a href=\"https://wandb.ai\" target=\"_blank\">Weights & Biases</a> <a href=\"https://docs.wandb.com/integrations/jupyter.html\" target=\"_blank\">(Documentation)</a>.<br/>\n",
       "                Project page: <a href=\"https://wandb.ai/native/%27P2D-NER-2021%27\" target=\"_blank\">https://wandb.ai/native/%27P2D-NER-2021%27</a><br/>\n",
       "                Run page: <a href=\"https://wandb.ai/native/%27P2D-NER-2021%27/runs/xk12ip6h\" target=\"_blank\">https://wandb.ai/native/%27P2D-NER-2021%27/runs/xk12ip6h</a><br/>\n",
       "                Run data is saved locally in <code>/home/phil/Dropbox/Python/Paper To Data/paper-2-data/wandb/run-20210519_110427-xk12ip6h</code><br/><br/>\n",
       "            "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/phil/anaconda3/envs/p2d/lib/python3.8/site-packages/torch/nn/parallel/_functions.py:65: UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; will instead unsqueeze and return a vector.\n",
      "  warnings.warn('Was asked to gather along dimension 0, but all '\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='1410' max='1410' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [1410/1410 16:51, Epoch 10/10]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: left;\">\n",
       "      <th>Epoch</th>\n",
       "      <th>Training Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "      <th>Precision</th>\n",
       "      <th>Recall</th>\n",
       "      <th>F1</th>\n",
       "      <th>Accuracy</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>0.846900</td>\n",
       "      <td>0.128038</td>\n",
       "      <td>0.472973</td>\n",
       "      <td>0.567568</td>\n",
       "      <td>0.515971</td>\n",
       "      <td>0.964672</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>0.068900</td>\n",
       "      <td>0.037652</td>\n",
       "      <td>0.889474</td>\n",
       "      <td>0.913514</td>\n",
       "      <td>0.901333</td>\n",
       "      <td>0.992257</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3</td>\n",
       "      <td>0.038000</td>\n",
       "      <td>0.039528</td>\n",
       "      <td>0.870192</td>\n",
       "      <td>0.978378</td>\n",
       "      <td>0.921120</td>\n",
       "      <td>0.988385</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4</td>\n",
       "      <td>0.027500</td>\n",
       "      <td>0.040539</td>\n",
       "      <td>0.875000</td>\n",
       "      <td>0.945946</td>\n",
       "      <td>0.909091</td>\n",
       "      <td>0.989837</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5</td>\n",
       "      <td>0.025900</td>\n",
       "      <td>0.033613</td>\n",
       "      <td>0.930481</td>\n",
       "      <td>0.940541</td>\n",
       "      <td>0.935484</td>\n",
       "      <td>0.990644</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>6</td>\n",
       "      <td>0.019100</td>\n",
       "      <td>0.038605</td>\n",
       "      <td>0.874372</td>\n",
       "      <td>0.940541</td>\n",
       "      <td>0.906250</td>\n",
       "      <td>0.989676</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>7</td>\n",
       "      <td>0.013700</td>\n",
       "      <td>0.031372</td>\n",
       "      <td>0.946524</td>\n",
       "      <td>0.956757</td>\n",
       "      <td>0.951613</td>\n",
       "      <td>0.993547</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>8</td>\n",
       "      <td>0.009700</td>\n",
       "      <td>0.028113</td>\n",
       "      <td>0.916667</td>\n",
       "      <td>0.951351</td>\n",
       "      <td>0.933687</td>\n",
       "      <td>0.993063</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>9</td>\n",
       "      <td>0.007200</td>\n",
       "      <td>0.041985</td>\n",
       "      <td>0.922280</td>\n",
       "      <td>0.962162</td>\n",
       "      <td>0.941799</td>\n",
       "      <td>0.992741</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>10</td>\n",
       "      <td>0.004100</td>\n",
       "      <td>0.042639</td>\n",
       "      <td>0.917526</td>\n",
       "      <td>0.962162</td>\n",
       "      <td>0.939314</td>\n",
       "      <td>0.992579</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/phil/anaconda3/envs/p2d/lib/python3.8/site-packages/torch/nn/parallel/_functions.py:65: UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; will instead unsqueeze and return a vector.\n",
      "  warnings.warn('Was asked to gather along dimension 0, but all '\n",
      "/home/phil/anaconda3/envs/p2d/lib/python3.8/site-packages/torch/nn/parallel/_functions.py:65: UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; will instead unsqueeze and return a vector.\n",
      "  warnings.warn('Was asked to gather along dimension 0, but all '\n",
      "/home/phil/anaconda3/envs/p2d/lib/python3.8/site-packages/torch/nn/parallel/_functions.py:65: UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; will instead unsqueeze and return a vector.\n",
      "  warnings.warn('Was asked to gather along dimension 0, but all '\n",
      "/home/phil/anaconda3/envs/p2d/lib/python3.8/site-packages/torch/nn/parallel/_functions.py:65: UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; will instead unsqueeze and return a vector.\n",
      "  warnings.warn('Was asked to gather along dimension 0, but all '\n",
      "/home/phil/anaconda3/envs/p2d/lib/python3.8/site-packages/torch/nn/parallel/_functions.py:65: UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; will instead unsqueeze and return a vector.\n",
      "  warnings.warn('Was asked to gather along dimension 0, but all '\n",
      "/home/phil/anaconda3/envs/p2d/lib/python3.8/site-packages/torch/nn/parallel/_functions.py:65: UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; will instead unsqueeze and return a vector.\n",
      "  warnings.warn('Was asked to gather along dimension 0, but all '\n",
      "/home/phil/anaconda3/envs/p2d/lib/python3.8/site-packages/torch/nn/parallel/_functions.py:65: UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; will instead unsqueeze and return a vector.\n",
      "  warnings.warn('Was asked to gather along dimension 0, but all '\n",
      "/home/phil/anaconda3/envs/p2d/lib/python3.8/site-packages/torch/nn/parallel/_functions.py:65: UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; will instead unsqueeze and return a vector.\n",
      "  warnings.warn('Was asked to gather along dimension 0, but all '\n",
      "/home/phil/anaconda3/envs/p2d/lib/python3.8/site-packages/torch/nn/parallel/_functions.py:65: UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; will instead unsqueeze and return a vector.\n",
      "  warnings.warn('Was asked to gather along dimension 0, but all '\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=1410, training_loss=0.10609516404199262, metrics={'train_runtime': 1016.54, 'train_samples_per_second': 1.387, 'total_flos': 0, 'epoch': 10.0, 'init_mem_cpu_alloc_delta': -763260928, 'init_mem_gpu_alloc_delta': 3539468800, 'init_mem_cpu_peaked_delta': 2627313664, 'init_mem_gpu_peaked_delta': 0, 'train_mem_cpu_alloc_delta': 3049766912, 'train_mem_gpu_alloc_delta': 10618411008, 'train_mem_cpu_peaked_delta': 0, 'train_mem_gpu_peaked_delta': 7092582912})"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Train\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/phil/anaconda3/envs/p2d/lib/python3.8/site-packages/torch/nn/parallel/_functions.py:65: UserWarning: Was asked to gather along dimension 0, but all input tensors were scalars; will instead unsqueeze and return a vector.\n",
      "  warnings.warn('Was asked to gather along dimension 0, but all '\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='16' max='16' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [16/16 00:03]\n",
       "    </div>\n",
       "    "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "{'eval_loss': 0.03137160837650299,\n",
       " 'eval_precision': 0.946524064171123,\n",
       " 'eval_recall': 0.9567567567567568,\n",
       " 'eval_f1': 0.9516129032258065,\n",
       " 'eval_accuracy': 0.9935473463461849,\n",
       " 'eval_runtime': 3.6338,\n",
       " 'eval_samples_per_second': 8.806,\n",
       " 'epoch': 10.0,\n",
       " 'eval_mem_cpu_alloc_delta': 0,\n",
       " 'eval_mem_gpu_alloc_delta': 0,\n",
       " 'eval_mem_cpu_peaked_delta': 151552,\n",
       " 'eval_mem_gpu_peaked_delta': 103408640}"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Evaluate based on the chosen epoch (usually best or last)\n",
    "trainer.evaluate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<br/>Waiting for W&B process to finish, PID 10722<br/>Program ended successfully."
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(Label(value=' 0.00MB of 0.00MB uploaded (0.00MB deduped)\\r'), FloatProgress(value=1.0, max=1.0)…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Find user logs for this run at: <code>/home/phil/Dropbox/Python/Paper To Data/paper-2-data/wandb/run-20210519_110427-xk12ip6h/logs/debug.log</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Find internal logs for this run at: <code>/home/phil/Dropbox/Python/Paper To Data/paper-2-data/wandb/run-20210519_110427-xk12ip6h/logs/debug-internal.log</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<h3>Run summary:</h3><br/><style>\n",
       "    table.wandb td:nth-child(1) { padding: 0 10px; text-align: right }\n",
       "    </style><table class=\"wandb\">\n",
       "<tr><td>train/loss</td><td>0.0041</td></tr><tr><td>train/learning_rate</td><td>0.0</td></tr><tr><td>train/epoch</td><td>10.0</td></tr><tr><td>train/global_step</td><td>1410</td></tr><tr><td>_runtime</td><td>1020</td></tr><tr><td>_timestamp</td><td>1621419687</td></tr><tr><td>_step</td><td>21</td></tr><tr><td>eval/loss</td><td>0.03137</td></tr><tr><td>eval/precision</td><td>0.94652</td></tr><tr><td>eval/recall</td><td>0.95676</td></tr><tr><td>eval/f1</td><td>0.95161</td></tr><tr><td>eval/accuracy</td><td>0.99355</td></tr><tr><td>eval/runtime</td><td>3.6338</td></tr><tr><td>eval/samples_per_second</td><td>8.806</td></tr><tr><td>train/train_runtime</td><td>1016.54</td></tr><tr><td>train/train_samples_per_second</td><td>1.387</td></tr><tr><td>train/total_flos</td><td>0</td></tr></table>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<h3>Run history:</h3><br/><style>\n",
       "    table.wandb td:nth-child(1) { padding: 0 10px; text-align: right }\n",
       "    </style><table class=\"wandb\">\n",
       "<tr><td>train/loss</td><td>█▂▁▁▁▁▁▁▁▁</td></tr><tr><td>train/learning_rate</td><td>▁▃▆██▇▆▅▃▂</td></tr><tr><td>train/epoch</td><td>▁▁▂▂▃▃▃▃▄▄▅▅▆▆▆▆▇▇████</td></tr><tr><td>train/global_step</td><td>▁▁▂▂▃▃▃▃▄▄▅▅▆▆▆▆▇▇████</td></tr><tr><td>_runtime</td><td>▁▁▂▂▂▃▃▃▄▄▅▅▆▆▆▆▇▇████</td></tr><tr><td>_timestamp</td><td>▁▁▂▂▂▃▃▃▄▄▅▅▆▆▆▆▇▇████</td></tr><tr><td>_step</td><td>▁▁▂▂▂▃▃▃▄▄▄▅▅▅▆▆▆▇▇▇██</td></tr><tr><td>eval/loss</td><td>█▂▂▂▁▂▁▁▂▂▁</td></tr><tr><td>eval/precision</td><td>▁▇▇▇█▇█████</td></tr><tr><td>eval/recall</td><td>▁▇█▇▇▇█████</td></tr><tr><td>eval/f1</td><td>▁▇█▇█▇█████</td></tr><tr><td>eval/accuracy</td><td>▁█▇▇▇▇█████</td></tr><tr><td>eval/runtime</td><td>▂▄█▅▃▁▄▃▄▆▅</td></tr><tr><td>eval/samples_per_second</td><td>▇▅▁▄▆█▅▆▅▃▄</td></tr><tr><td>train/train_runtime</td><td>▁</td></tr><tr><td>train/train_samples_per_second</td><td>▁</td></tr><tr><td>train/total_flos</td><td>▁</td></tr></table><br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Synced 5 W&B file(s), 0 media file(s), 0 artifact file(s) and 0 other file(s)"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "\n",
       "                    <br/>Synced <strong style=\"color:#cdcd00\">microsoft/deberta-v2-xlarge-19-05-2021</strong>: <a href=\"https://wandb.ai/native/%27P2D-NER-2021%27/runs/xk12ip6h\" target=\"_blank\">https://wandb.ai/native/%27P2D-NER-2021%27/runs/xk12ip6h</a><br/>\n",
       "                "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Finish Weighs & Biases logging for this run\n",
    "wandb.finish() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save the model, good practice given the work required to train a model and  \n",
    "# also can be used just for inference on new data\n",
    "trainer.save_model(SAVED_MODEL)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 5: Predictions and Inference\n",
    "\n",
    "Now to run the trained and serialised model on the evaluation set again, NOT the data used for training.\n",
    "\n",
    "Always take care to ensure that there isn't any data leakage here, eg the same agreements, different agreements from the same set of agreement or potentially different agreements from the same parties. \n",
    "\n",
    "The objective is to ensure that the model is able to generalize well to new agreements never seen before.\n",
    "\n",
    "To match the number of predictions to the original numnber of tokens, need to use: \"label_all_tokens=False\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the model and instantiate\n",
    "loaded_model = AutoModelForTokenClassification.from_pretrained(SAVED_MODEL)\n",
    "\n",
    "pred_trainer = Trainer(\n",
    "    loaded_model,\n",
    "    args,\n",
    "    data_collator=data_collator,\n",
    "    tokenizer=tokenizer,\n",
    "    compute_metrics=compute_metrics\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='16' max='16' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [16/16 00:03]\n",
       "    </div>\n",
       "    "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "              precision    recall  f1-score   support\n",
      "\n",
      "   AGMT_DATE       0.92      0.92      0.92        37\n",
      "    DOC_NAME       0.94      0.98      0.96        60\n",
      "       PARTY       0.97      0.95      0.96        88\n",
      "\n",
      "   micro avg       0.95      0.96      0.95       185\n",
      "   macro avg       0.94      0.95      0.95       185\n",
      "weighted avg       0.95      0.96      0.95       185\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Extract the predictions and produce a classification report\n",
    "predictions, labels, _ = pred_trainer.predict(tokenized_datasets[\"test\"])\n",
    "predictions = np.argmax(predictions, axis=2)\n",
    "\n",
    "# Remove ignored index (special tokens)\n",
    "true_predictions = [\n",
    "    [label_list[p] for (p, l) in zip(prediction, label) if l != -100]\n",
    "    for prediction, label in zip(predictions, labels)\n",
    "]\n",
    "true_labels = [\n",
    "    [label_list[l] for (p, l) in zip(prediction, label) if l != -100]\n",
    "    for prediction, label in zip(predictions, labels)\n",
    "    ]\n",
    "\n",
    "# Generate the metrics and display\n",
    "results = classification_report(true_labels, true_predictions, zero_division=1)\n",
    "print(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "195\n",
      "195\n",
      "195\n"
     ]
    }
   ],
   "source": [
    "# Check a sample for the evaluation set\n",
    "check = 3\n",
    "\n",
    "print(len(datasets[\"test\"][check]['split_tokens']))\n",
    "print(len(true_predictions[check]))\n",
    "print(len(true_labels[check]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('COLLABORATION', 'B-DOC_NAME')\n",
      "('AGREEMENT', 'I-DOC_NAME')\n",
      "('Collaboration', 'B-DOC_NAME')\n",
      "('Agreement', 'I-DOC_NAME')\n",
      "('April', 'B-AGMT_DATE')\n",
      "('14th', 'I-AGMT_DATE')\n",
      "(',', 'I-AGMT_DATE')\n",
      "('2020', 'I-AGMT_DATE')\n",
      "('Anixa', 'B-PARTY')\n",
      "('Biosciences', 'I-PARTY')\n",
      "(',', 'I-PARTY')\n",
      "('Inc.', 'I-PARTY')\n",
      "('OntoChem', 'B-PARTY')\n",
      "('GmbH', 'I-PARTY')\n"
     ]
    }
   ],
   "source": [
    "# Have a look at the predicted extracted data\n",
    "check_pred = zip(datasets[\"test\"][check]['split_tokens'], true_predictions[check])\n",
    "for tup in check_pred:\n",
    "    if tup[1] != 'O':\n",
    "        print(tup)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "('COLLABORATION', 'B-DOC_NAME')\n",
      "('AGREEMENT', 'I-DOC_NAME')\n",
      "('Collaboration', 'B-DOC_NAME')\n",
      "('Agreement', 'I-DOC_NAME')\n",
      "('April', 'B-AGMT_DATE')\n",
      "('14th', 'I-AGMT_DATE')\n",
      "(',', 'I-AGMT_DATE')\n",
      "('2020', 'I-AGMT_DATE')\n",
      "('Anixa', 'B-PARTY')\n",
      "('Biosciences', 'I-PARTY')\n",
      "(',', 'I-PARTY')\n",
      "('Inc.', 'I-PARTY')\n",
      "('OntoChem', 'B-PARTY')\n",
      "('GmbH', 'I-PARTY')\n"
     ]
    }
   ],
   "source": [
    "# Compare to the actual labels\n",
    "check_true = zip(datasets[\"test\"][check]['split_tokens'], true_labels[check])\n",
    "for tup in check_true:\n",
    "    if tup[1] != 'O':\n",
    "        print(tup)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}