Medical Entity Recognition with Pretrained Transformers.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Medical Entity Recognition with Pretrained Transformers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook we explore how we can use pretrained transformer models, such as BERT, to identify medical entities in text."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We're going to work with the [NCBI Disease](https://www.ncbi.nlm.nih.gov/CBBresearch/Dogan/DISEASE/) corpus, a corpus of 793 PubMed abstracts with 6,892 annotated disease mentions. This dataset can be downloaded with the `datasets` package from Huggingface, which gives us easy access to hundreds of interesting datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Reusing dataset ncbi_disease (/home/yves/.cache/huggingface/datasets/ncbi_disease/ncbi_disease/1.0.0/88b09d40ea6f23141af303fa5f80b07d7de87612b2e41f494f537c4eff97372c)\n"
     ]
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "dataset = load_dataset('ncbi_disease')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see what this dataset looks like. The abstracts are split into sentences, which already have been tokenized for us. There are 5433 sentences in the training data, 924 in the validation data and another 941 in the test data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['id', 'tokens', 'ner_tags'],\n",
       "        num_rows: 5433\n",
       "    })\n",
       "    validation: Dataset({\n",
       "        features: ['id', 'tokens', 'ner_tags'],\n",
       "        num_rows: 924\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['id', 'tokens', 'ner_tags'],\n",
       "        num_rows: 941\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first training example is the sentence 'Identification of APC2, a homologue of the adenomatous polyposis coli tumour suppressor.' The phrase 'adenomatous polyposis coli tumour' has been labeled as a disease. The first token has a different label than the other three, because it has been labeled as the start of the disease mention, following the BIO labeling scheme."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'id': '0',\n",
       " 'ner_tags': [0, 0, 0, 0, 0, 0, 0, 0, 1, 2, 2, 2, 0, 0],\n",
       " 'tokens': ['Identification',\n",
       "  'of',\n",
       "  'APC2',\n",
       "  ',',\n",
       "  'a',\n",
       "  'homologue',\n",
       "  'of',\n",
       "  'the',\n",
       "  'adenomatous',\n",
       "  'polyposis',\n",
       "  'coli',\n",
       "  'tumour',\n",
       "  'suppressor',\n",
       "  '.']}"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset[\"train\"][0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preprocessing the texts"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we have our texts, we need to give them the correct preprocessing treatment. As our model, we choose one of the available PubMedBERTs — BERT models that have been pretrained on abstracts (and in this case, also full texts) from PubMed and therefore look perfect for the type of texts we're working with. We start by getting the tokenizer that was used for pretraining this model, because our texts need to be tokenized in exactly the same manner."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer\n",
    "\n",
    "MODEL = \"microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext\"\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(MODEL)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's now use this tokenizer to tokenize our texts. Note that every sentence in our corpus is a list of words, so we need to tell the tokenizer the text has already been split into words. In addition, we'll also ask the tokenizer to pad and/or truncate the texts. Sentences that are longer than 256 tokens will be truncated, and all sentences will be padded to the length of the (resulting) longest one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_texts = [item[\"tokens\"] for item in dataset[\"train\"]]\n",
    "dev_texts = [item[\"tokens\"] for item in dataset[\"validation\"]]\n",
    "test_texts = [item[\"tokens\"] for item in dataset[\"test\"]]\n",
    "\n",
    "train_texts_encoded = tokenizer(train_texts, padding=True, truncation=True, max_length=256, is_split_into_words=True)\n",
    "dev_texts_encoded = tokenizer(dev_texts, padding=True, truncation=True, max_length=256, is_split_into_words=True)\n",
    "test_texts_encoded = tokenizer(test_texts, padding=True, truncation=True, max_length=256, is_split_into_words=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now have three lists of `Encoding`s, which contain all information that our model needs, in particular the ids of the tokens, their type id, and their attention mask. The mask is used to make sure that the model ignores padding tokens. The type id of the tokens is always `0`, because our input consists of single sentences, and not sentence pairs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Encoding(num_tokens=138, attributes=[ids, type_ids, tokens, offsets, attention_mask, special_tokens_mask, overflowing])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_texts_encoded[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we look at the actual tokens, we see the tokenizer has applied a type of tokenization different from traditional tokenization. To keep the size of the vocabulary manageable, unknown words have been split up into known subword parts, such as `apc2`, which has been split up into `apc` and `##2`, where the `##` indicates this is a continuation. \n",
    "\n",
    "At the same time, the tokens also display one of the main benefits of using a BERT model that was pretrained on data from PubMed. While generic BERT would split up complex words such as `adenomatous` or `polyposis`, they occur frequently enough in PubMed data for PubMedBERT to treat them as one single token."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['[CLS]',\n",
       " 'identification',\n",
       " 'of',\n",
       " 'apc',\n",
       " '##2',\n",
       " ',',\n",
       " 'a',\n",
       " 'homologue',\n",
       " 'of',\n",
       " 'the',\n",
       " 'adenomatous',\n",
       " 'polyposis',\n",
       " 'coli',\n",
       " 'tumour',\n",
       " 'suppressor',\n",
       " '.',\n",
       " '[SEP]',\n",
       " '[PAD]',\n",
       " '[PAD]',\n",
       " '[PAD]']"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_texts_encoded[0].tokens[:20]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preprocessing the labels"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There's one remaining challenge. Because our new tokens are different from the original tokens in the corpus, we can't just train the model on the original labels: we need to align the labels with the new tokens. Luckily the tokenizer also provides us with a list of offsets for every new token, where we can easily identify tokens that do not correspond to the original words. \n",
    "\n",
    "For example, the offsets of the first training sentence tell us that `apc2` has been split up into two tokens, one for the first three characters of the word (indices 0 to, but not including, 3) and one for the last character of the word (indices 3 to, but not including, 4). \n",
    "\n",
    "Additionally, we can also identify special tokens, such as `[CLS]` and `[PAD]` by the offset pair `[(0,0)]`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(0, 0),\n",
       " (0, 14),\n",
       " (0, 2),\n",
       " (0, 3),\n",
       " (3, 4),\n",
       " (0, 1),\n",
       " (0, 1),\n",
       " (0, 9),\n",
       " (0, 2),\n",
       " (0, 3),\n",
       " (0, 11),\n",
       " (0, 9),\n",
       " (0, 4),\n",
       " (0, 6),\n",
       " (0, 10),\n",
       " (0, 1),\n",
       " (0, 0),\n",
       " (0, 0),\n",
       " (0, 0),\n",
       " (0, 0)]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_texts_encoded[0].offsets[:20]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There are only three labels in the corpus — `O`, `B-disease` and `I-disease` — which have already been mapped to their index for us."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[0, 1, 2]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "all_labels = list(set([label for item in dataset[\"train\"] for label in item[\"ner_tags\"]]))\n",
    "all_labels"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we have sufficient information to align the entity labels with the new tokens. For each sentence, we first create a numpy array filled with the label `-100`, a special label in the `transformers` library that will be ignored during training. Then we copy the original labels to the tokens at the start of every word. These have zero as their first offset position, and another number as their second position. This means the remaining tokens of the word will still have the label `-100`. This comes in handy during evaluation, as the subword tokenization will not lead to a higher number of entity labels. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "def map_entities_to_tokens(items, encodings):\n",
    "    \n",
    "    labels = [item[\"ner_tags\"] for item in items]\n",
    "    offsets = [encoding.offsets for encoding in encodings]\n",
    "    encoded_labels = []\n",
    "    for doc_labels, doc_offset in zip(labels, offsets):\n",
    "        # create an empty array of -100\n",
    "        doc_enc_labels = np.ones(len(doc_offset),dtype=int) * -100\n",
    "        arr_offset = np.array(doc_offset)\n",
    "\n",
    "        # set labels whose first offset position is 0 and the second is not 0\n",
    "        doc_enc_labels[(arr_offset[:,0] == 0) & (arr_offset[:,1] != 0)] = doc_labels\n",
    "        encoded_labels.append(doc_enc_labels.tolist())\n",
    "\n",
    "    return encoded_labels\n",
    "\n",
    "train_labels = map_entities_to_tokens(dataset[\"train\"], train_texts_encoded.encodings)\n",
    "dev_labels = map_entities_to_tokens(dataset[\"validation\"], dev_texts_encoded.encodings)\n",
    "test_labels = map_entities_to_tokens(dataset[\"test\"], test_texts_encoded.encodings)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is the result for our first training example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('[CLS]', -100),\n",
       " ('identification', 0),\n",
       " ('of', 0),\n",
       " ('apc', 0),\n",
       " ('##2', -100),\n",
       " (',', 0),\n",
       " ('a', 0),\n",
       " ('homologue', 0),\n",
       " ('of', 0),\n",
       " ('the', 0),\n",
       " ('adenomatous', 1),\n",
       " ('polyposis', 2),\n",
       " ('coli', 2),\n",
       " ('tumour', 2),\n",
       " ('suppressor', 0),\n",
       " ('.', 0),\n",
       " ('[SEP]', -100),\n",
       " ('[PAD]', -100),\n",
       " ('[PAD]', -100),\n",
       " ('[PAD]', -100)]"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "list(zip(train_texts_encoded[0].tokens[:20], train_labels[0][:20]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting up the dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We bring the encoded texts and their labels together in an NERDataset. This dataset returns for every item all the information in the encodings as a dictionary, and adds an additional key with the labels. All lists are converted to PyTorch tensors. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train items: 5433\n",
      "Dev items: 924\n",
      "Test items: 941\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "class NERDataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, encodings, labels):\n",
    "        self.encodings = encodings\n",
    "        self.labels = labels\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}\n",
    "        item['labels'] = torch.tensor(self.labels[idx])\n",
    "        return item\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.labels)\n",
    "    \n",
    "\n",
    "train_dataset = NERDataset(train_texts_encoded, train_labels)\n",
    "dev_dataset = NERDataset(dev_texts_encoded, dev_labels)\n",
    "test_dataset = NERDataset(test_texts_encoded, test_labels)\n",
    "\n",
    "print(f\"Train items: {len(train_dataset)}\")\n",
    "print(f\"Dev items: {len(dev_dataset)}\")\n",
    "print(f\"Test items: {len(test_dataset)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we set up the evaluation of the results. We compute an accuracy score on all labels, excluding `-100`. In named entity recognition, accuracy tends to be very high, because most tokens are not part of an entity mention. Even models that do not recognize a single token, will achieve a high accuracy on most datasets. Therefore we also compute precision, recall and F-score on the entity labels only (excluding label `0`). This is a much better measure of the model's success at identifying entities."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import accuracy_score, precision_recall_fscore_support\n",
    "\n",
    "def compute_metrics(pred):\n",
    "    labels = pred.label_ids\n",
    "    preds = pred.predictions.argmax(-1)\n",
    "\n",
    "    flat_labels, flat_preds = [], []\n",
    "    flat_ent_labels, flat_ent_preds = [], []\n",
    "    for label_row, pred_row in zip(labels, preds):\n",
    "        for label, pred_label in zip(label_row, pred_row):\n",
    "            if label != -100:\n",
    "                flat_labels.append(label)\n",
    "                flat_preds.append(pred_label)\n",
    "                if label != 0 or pred_label != 0:\n",
    "                    flat_ent_labels.append(label)\n",
    "                    flat_ent_preds.append(pred_label)\n",
    "                    \n",
    "        \n",
    "    precision, recall, f1, _ = precision_recall_fscore_support(flat_ent_labels, flat_ent_preds, average='micro')\n",
    "    acc = accuracy_score(flat_labels, flat_preds)\n",
    "    return {\n",
    "        'accuracy': acc,\n",
    "        'f1': f1,\n",
    "        'precision': precision,\n",
    "        'recall': recall\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training the model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we're ready to train the model. This is easy to do with the `Trainer` class of the `transformers` package, which we feed with the model, the training and development dataset, the evaluation metrics, along with all training arguments. \n",
    "\n",
    "We'll train the model for 3 epochs, with a batch size of 8, and evaluate and save a checkpoint at every 200 training steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext were not used when initializing BertForTokenClassification: ['cls.predictions.bias', 'cls.predictions.transform.dense.weight', 'cls.predictions.transform.dense.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.transform.LayerNorm.bias', 'cls.predictions.decoder.weight', 'cls.predictions.decoder.bias', 'cls.seq_relationship.weight', 'cls.seq_relationship.bias']\n",
      "- This IS expected if you are initializing BertForTokenClassification 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 BertForTokenClassification 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 BertForTokenClassification were not initialized from the model checkpoint at microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext and are newly initialized: ['classifier.weight', 'classifier.bias']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "        <style>\n",
       "            /* Turns off some styling */\n",
       "            progress {\n",
       "                /* gets rid of default border in Firefox and Opera. */\n",
       "                border: none;\n",
       "                /* Needs to be in here for Safari polyfill so background images work as expected. */\n",
       "                background-size: auto;\n",
       "            }\n",
       "        </style>\n",
       "      \n",
       "      <progress value='2040' max='2040' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [2040/2040 06:19, Epoch 3/3]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "      <th>Validation Loss</th>\n",
       "      <th>Accuracy</th>\n",
       "      <th>F1</th>\n",
       "      <th>Precision</th>\n",
       "      <th>Recall</th>\n",
       "      <th>Runtime</th>\n",
       "      <th>Samples Per Second</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>200</td>\n",
       "      <td>No log</td>\n",
       "      <td>0.087361</td>\n",
       "      <td>0.975427</td>\n",
       "      <td>0.698104</td>\n",
       "      <td>0.698104</td>\n",
       "      <td>0.698104</td>\n",
       "      <td>3.687000</td>\n",
       "      <td>250.609000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>400</td>\n",
       "      <td>No log</td>\n",
       "      <td>0.062153</td>\n",
       "      <td>0.982644</td>\n",
       "      <td>0.804051</td>\n",
       "      <td>0.804051</td>\n",
       "      <td>0.804051</td>\n",
       "      <td>3.685700</td>\n",
       "      <td>250.696000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>600</td>\n",
       "      <td>0.170500</td>\n",
       "      <td>0.053284</td>\n",
       "      <td>0.985231</td>\n",
       "      <td>0.818648</td>\n",
       "      <td>0.818648</td>\n",
       "      <td>0.818648</td>\n",
       "      <td>3.688500</td>\n",
       "      <td>250.508000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>800</td>\n",
       "      <td>0.170500</td>\n",
       "      <td>0.051288</td>\n",
       "      <td>0.984855</td>\n",
       "      <td>0.827635</td>\n",
       "      <td>0.827635</td>\n",
       "      <td>0.827635</td>\n",
       "      <td>3.685100</td>\n",
       "      <td>250.742000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1000</td>\n",
       "      <td>0.041500</td>\n",
       "      <td>0.045846</td>\n",
       "      <td>0.985898</td>\n",
       "      <td>0.834961</td>\n",
       "      <td>0.834961</td>\n",
       "      <td>0.834961</td>\n",
       "      <td>3.699100</td>\n",
       "      <td>249.791000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1200</td>\n",
       "      <td>0.041500</td>\n",
       "      <td>0.044199</td>\n",
       "      <td>0.986733</td>\n",
       "      <td>0.840920</td>\n",
       "      <td>0.840920</td>\n",
       "      <td>0.840920</td>\n",
       "      <td>3.699700</td>\n",
       "      <td>249.747000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1400</td>\n",
       "      <td>0.041500</td>\n",
       "      <td>0.045546</td>\n",
       "      <td>0.987025</td>\n",
       "      <td>0.841650</td>\n",
       "      <td>0.841650</td>\n",
       "      <td>0.841650</td>\n",
       "      <td>3.695600</td>\n",
       "      <td>250.029000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1600</td>\n",
       "      <td>0.027400</td>\n",
       "      <td>0.051103</td>\n",
       "      <td>0.987150</td>\n",
       "      <td>0.845846</td>\n",
       "      <td>0.845846</td>\n",
       "      <td>0.845846</td>\n",
       "      <td>3.686800</td>\n",
       "      <td>250.623000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1800</td>\n",
       "      <td>0.027400</td>\n",
       "      <td>0.056813</td>\n",
       "      <td>0.987734</td>\n",
       "      <td>0.849693</td>\n",
       "      <td>0.849693</td>\n",
       "      <td>0.849693</td>\n",
       "      <td>3.694900</td>\n",
       "      <td>250.074000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2000</td>\n",
       "      <td>0.012200</td>\n",
       "      <td>0.056784</td>\n",
       "      <td>0.988026</td>\n",
       "      <td>0.853646</td>\n",
       "      <td>0.853646</td>\n",
       "      <td>0.853646</td>\n",
       "      <td>3.692000</td>\n",
       "      <td>250.269000</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=2040, training_loss=0.061846355451088325, metrics={'train_runtime': 379.4347, 'train_samples_per_second': 5.376, 'total_flos': 1469586210067260.0, 'epoch': 3.0, 'init_mem_cpu_alloc_delta': 4504356, 'init_mem_gpu_alloc_delta': 436712960, 'init_mem_cpu_peaked_delta': 19471, 'init_mem_gpu_peaked_delta': 0, 'train_mem_cpu_alloc_delta': 3453172, 'train_mem_gpu_alloc_delta': 1789238784, 'train_mem_cpu_peaked_delta': 188073233, 'train_mem_gpu_peaked_delta': 553228800})"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers import Trainer, TrainingArguments, AutoModelForTokenClassification, BertForTokenClassification\n",
    "\n",
    "model = AutoModelForTokenClassification.from_pretrained(MODEL, num_labels=len(all_labels))\n",
    "\n",
    "\n",
    "#device = torch.device(\"cpu\")\n",
    "#model.to(device)\n",
    "\n",
    "training_args = TrainingArguments(\n",
    "    output_dir='./results',          # output directory\n",
    "    num_train_epochs=3,              # total # of training epochs\n",
    "    per_device_train_batch_size=8,  # batch size per device during training\n",
    "    per_device_eval_batch_size=8,   # batch size for evaluation\n",
    "    warmup_steps=int(len(train_dataset)/8),  # number of warmup steps for learning rate scheduler\n",
    "    weight_decay=0.01,               # strength of weight decay\n",
    "    logging_dir='./logs',            # directory for storing logs\n",
    "    evaluation_strategy=\"steps\",\n",
    "    eval_steps=200,\n",
    "    save_steps=200,\n",
    "    save_total_limit=10,\n",
    "    load_best_model_at_end=True,\n",
    "    no_cuda=False\n",
    ")\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=model,                         # the instantiated 🤗 Transformers model to be trained\n",
    "    args=training_args,                  # training arguments, defined above\n",
    "    compute_metrics=compute_metrics,\n",
    "    train_dataset=train_dataset,         # training dataset\n",
    "    eval_dataset=dev_dataset,            # evaluation dataset\n",
    ")\n",
    "\n",
    "trainer.train()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluating the results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we evaluate the model on the test dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "        <style>\n",
       "            /* Turns off some styling */\n",
       "            progress {\n",
       "                /* gets rid of default border in Firefox and Opera. */\n",
       "                border: none;\n",
       "                /* Needs to be in here for Safari polyfill so background images work as expected. */\n",
       "                background-size: auto;\n",
       "            }\n",
       "        </style>\n",
       "      \n",
       "      <progress value='118' max='118' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [118/118 00:04]\n",
       "    </div>\n",
       "    "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "{'eval_loss': 0.04503929987549782,\n",
       " 'eval_accuracy': 0.9859166428542271,\n",
       " 'eval_f1': 0.8441734417344173,\n",
       " 'eval_precision': 0.8441734417344173,\n",
       " 'eval_recall': 0.8441734417344173,\n",
       " 'eval_runtime': 4.5188,\n",
       " 'eval_samples_per_second': 208.241,\n",
       " 'epoch': 3.0,\n",
       " 'eval_mem_cpu_alloc_delta': 2450645,\n",
       " 'eval_mem_gpu_alloc_delta': 0,\n",
       " 'eval_mem_cpu_peaked_delta': 3083163,\n",
       " 'eval_mem_gpu_peaked_delta': 37012480}"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.evaluate(test_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For inference, we can load the model and combine it with the tokenizer in an `ner` pipeline. Now we can easily label new texts and inspect the results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import pipeline\n",
    "\n",
    "model = AutoModelForTokenClassification.from_pretrained(\"results/checkpoint-2000\")\n",
    "nlp = pipeline(\"ner\", tokenizer=tokenizer, model=model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'id': '1', 'ner_tags': [1, 2, 2, 0, 1, 2, 2, 0, 0, 0, 1, 2, 2, 2, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 'tokens': ['Ataxia', '-', 'telangiectasia', '(', 'A', '-', 'T', ')', 'is', 'a', 'recessive', 'multi', '-', 'system', 'disorder', 'caused', 'by', 'mutations', 'in', 'the', 'ATM', 'gene', 'at', '11q22', '-', 'q23', '(', 'ref', '.', '3', ')', '.']}\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[[{'word': 'ataxia',\n",
       "   'score': 0.9266947507858276,\n",
       "   'entity': 'LABEL_1',\n",
       "   'index': 1,\n",
       "   'start': 0,\n",
       "   'end': 6}],\n",
       " [{'word': '-',\n",
       "   'score': 0.9998806715011597,\n",
       "   'entity': 'LABEL_0',\n",
       "   'index': 1,\n",
       "   'start': 0,\n",
       "   'end': 1}],\n",
       " [{'word': 'telangiect',\n",
       "   'score': 0.9920430779457092,\n",
       "   'entity': 'LABEL_1',\n",
       "   'index': 1,\n",
       "   'start': 0,\n",
       "   'end': 10},\n",
       "  {'word': '##asia',\n",
       "   'score': 0.9963495135307312,\n",
       "   'entity': 'LABEL_2',\n",
       "   'index': 2,\n",
       "   'start': 10,\n",
       "   'end': 14}],\n",
       " [{'word': '(',\n",
       "   'score': 0.9998924732208252,\n",
       "   'entity': 'LABEL_0',\n",
       "   'index': 1,\n",
       "   'start': 0,\n",
       "   'end': 1}],\n",
       " [{'word': 'a',\n",
       "   'score': 0.9997237324714661,\n",
       "   'entity': 'LABEL_0',\n",
       "   'index': 1,\n",
       "   'start': 0,\n",
       "   'end': 1}],\n",
       " [{'word': '-',\n",
       "   'score': 0.9998806715011597,\n",
       "   'entity': 'LABEL_0',\n",
       "   'index': 1,\n",
       "   'start': 0,\n",
       "   'end': 1}],\n",
       " [{'word': 't',\n",
       "   'score': 0.9996972680091858,\n",
       "   'entity': 'LABEL_0',\n",
       "   'index': 1,\n",
       "   'start': 0,\n",
       "   'end': 1}],\n",
       " [{'word': ')',\n",
       "   'score': 0.99981689453125,\n",
       "   'entity': 'LABEL_0',\n",
       "   'index': 1,\n",
       "   'start': 0,\n",
       "   'end': 1}],\n",
       " [{'word': 'is',\n",
       "   'score': 0.9998725056648254,\n",
       "   'entity': 'LABEL_0',\n",
       "   'index': 1,\n",
       "   'start': 0,\n",
       "   'end': 2}],\n",
       " [{'word': 'a',\n",
       "   'score': 0.9997237324714661,\n",
       "   'entity': 'LABEL_0',\n",
       "   'index': 1,\n",
       "   'start': 0,\n",
       "   'end': 1}],\n",
       " [{'word': 'recessive',\n",
       "   'score': 0.9997004270553589,\n",
       "   'entity': 'LABEL_0',\n",
       "   'index': 1,\n",
       "   'start': 0,\n",
       "   'end': 9}],\n",
       " [{'word': 'multi',\n",
       "   'score': 0.9998838901519775,\n",
       "   'entity': 'LABEL_0',\n",
       "   'index': 1,\n",
       "   'start': 0,\n",
       "   'end': 5}],\n",
       " [{'word': '-',\n",
       "   'score': 0.9998806715011597,\n",
       "   'entity': 'LABEL_0',\n",
       "   'index': 1,\n",
       "   'start': 0,\n",
       "   'end': 1}],\n",
       " [{'word': 'system',\n",
       "   'score': 0.9998276829719543,\n",
       "   'entity': 'LABEL_0',\n",
       "   'index': 1,\n",
       "   'start': 0,\n",
       "   'end': 6}],\n",
       " [{'word': 'disorder',\n",
       "   'score': 0.9770255088806152,\n",
       "   'entity': 'LABEL_0',\n",
       "   'index': 1,\n",
       "   'start': 0,\n",
       "   'end': 8}],\n",
       " [{'word': 'caused',\n",
       "   'score': 0.9998867511749268,\n",
       "   'entity': 'LABEL_0',\n",
       "   'index': 1,\n",
       "   'start': 0,\n",
       "   'end': 6}],\n",
       " [{'word': 'by',\n",
       "   'score': 0.9997820258140564,\n",
       "   'entity': 'LABEL_0',\n",
       "   'index': 1,\n",
       "   'start': 0,\n",
       "   'end': 2}],\n",
       " [{'word': 'mutations',\n",
       "   'score': 0.9998729825019836,\n",
       "   'entity': 'LABEL_0',\n",
       "   'index': 1,\n",
       "   'start': 0,\n",
       "   'end': 9}],\n",
       " [{'word': 'in',\n",
       "   'score': 0.9997439384460449,\n",
       "   'entity': 'LABEL_0',\n",
       "   'index': 1,\n",
       "   'start': 0,\n",
       "   'end': 2}],\n",
       " [{'word': 'the',\n",
       "   'score': 0.999903678894043,\n",
       "   'entity': 'LABEL_0',\n",
       "   'index': 1,\n",
       "   'start': 0,\n",
       "   'end': 3}],\n",
       " [{'word': 'atm',\n",
       "   'score': 0.9996559023857117,\n",
       "   'entity': 'LABEL_0',\n",
       "   'index': 1,\n",
       "   'start': 0,\n",
       "   'end': 3}],\n",
       " [{'word': 'gene',\n",
       "   'score': 0.9998477697372437,\n",
       "   'entity': 'LABEL_0',\n",
       "   'index': 1,\n",
       "   'start': 0,\n",
       "   'end': 4}],\n",
       " [{'word': 'at',\n",
       "   'score': 0.9990025758743286,\n",
       "   'entity': 'LABEL_0',\n",
       "   'index': 1,\n",
       "   'start': 0,\n",
       "   'end': 2}],\n",
       " [{'word': '11',\n",
       "   'score': 0.9999526143074036,\n",
       "   'entity': 'LABEL_0',\n",
       "   'index': 1,\n",
       "   'start': 0,\n",
       "   'end': 2},\n",
       "  {'word': '##q2',\n",
       "   'score': 0.999955415725708,\n",
       "   'entity': 'LABEL_0',\n",
       "   'index': 2,\n",
       "   'start': 2,\n",
       "   'end': 4},\n",
       "  {'word': '##2',\n",
       "   'score': 0.9999492764472961,\n",
       "   'entity': 'LABEL_0',\n",
       "   'index': 3,\n",
       "   'start': 4,\n",
       "   'end': 5}],\n",
       " [{'word': '-',\n",
       "   'score': 0.9998806715011597,\n",
       "   'entity': 'LABEL_0',\n",
       "   'index': 1,\n",
       "   'start': 0,\n",
       "   'end': 1}],\n",
       " [{'word': 'q2',\n",
       "   'score': 0.999927282333374,\n",
       "   'entity': 'LABEL_0',\n",
       "   'index': 1,\n",
       "   'start': 0,\n",
       "   'end': 2},\n",
       "  {'word': '##3',\n",
       "   'score': 0.9999404549598694,\n",
       "   'entity': 'LABEL_0',\n",
       "   'index': 2,\n",
       "   'start': 2,\n",
       "   'end': 3}],\n",
       " [{'word': '(',\n",
       "   'score': 0.9998924732208252,\n",
       "   'entity': 'LABEL_0',\n",
       "   'index': 1,\n",
       "   'start': 0,\n",
       "   'end': 1}],\n",
       " [{'word': 'ref',\n",
       "   'score': 0.9997204542160034,\n",
       "   'entity': 'LABEL_0',\n",
       "   'index': 1,\n",
       "   'start': 0,\n",
       "   'end': 3}],\n",
       " [{'word': '.',\n",
       "   'score': 0.9999040961265564,\n",
       "   'entity': 'LABEL_0',\n",
       "   'index': 1,\n",
       "   'start': 0,\n",
       "   'end': 1}],\n",
       " [{'word': '3',\n",
       "   'score': 0.999706506729126,\n",
       "   'entity': 'LABEL_0',\n",
       "   'index': 1,\n",
       "   'start': 0,\n",
       "   'end': 1}],\n",
       " [{'word': ')',\n",
       "   'score': 0.99981689453125,\n",
       "   'entity': 'LABEL_0',\n",
       "   'index': 1,\n",
       "   'start': 0,\n",
       "   'end': 1}],\n",
       " [{'word': '.',\n",
       "   'score': 0.9999040961265564,\n",
       "   'entity': 'LABEL_0',\n",
       "   'index': 1,\n",
       "   'start': 0,\n",
       "   'end': 1}]]"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "print(dataset[\"test\"][1])\n",
    "\n",
    "nlp(dataset[\"test\"][1][\"tokens\"])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "notebooks",
   "language": "python",
   "name": "notebooks"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}