GILBERT.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": [],
      "collapsed_sections": []
    },
    "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.6.4"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "d4_YJqjR_Gjw"
      },
      "source": [
        "!pip install --quiet transformers seqeval[gpu]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IEnlUbgm8z3B"
      },
      "source": [
        "import pandas as pd\n",
        "import numpy as np\n",
        "from sklearn.metrics import accuracy_score\n",
        "import torch\n",
        "from torch.utils.data import Dataset, DataLoader\n",
        "from transformers import BertTokenizer, BertConfig, BertForTokenClassification"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Sm1krxJtKxpx",
        "tags": []
      },
      "source": [
        "from torch import cuda\n",
        "device = 'cuda' if cuda.is_available() else 'cpu'\n",
        "print(device)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')\n",
        "%cd /content/drive/My Drive/"
      ],
      "metadata": {
        "id": "f9W_OvYVDTXf"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "deLB9HVX5I6F"
      },
      "source": [
        "data = pd.read_csv(\"bert_data.csv\", encoding='unicode_escape')\n",
        "data.head()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6gMibEJXTKDw"
      },
      "source": [
        "data.count()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "tags": [],
        "id": "s4Jn1fVT_GkO"
      },
      "source": [
        "print(\"Number of tags: {}\".format(len(data.Tag.unique())))\n",
        "frequencies = data.Tag.value_counts()\n",
        "frequencies"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Hmd-ow389k6Y"
      },
      "source": [
        "# let's create a new column called \"sentence\" which groups the words by sentence \n",
        "data['sentence'] = data[['Sentence #','Word','Tag']].groupby(['Sentence #'])['Word'].transform(lambda x: ' '.join(x))\n",
        "# let's also create a new column called \"word_labels\" which groups the tags by sentence \n",
        "data['word_labels'] = data[['Sentence #','Word','Tag']].groupby(['Sentence #'])['Tag'].transform(lambda x: ','.join(x))\n",
        "data.head()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CFRDM8WsQXvL"
      },
      "source": [
        "label2id = {k: v for v, k in enumerate(np.flip(data.Tag.unique()))}\n",
        "id2label = {v: k for v, k in enumerate(np.flip(data.Tag.unique()))}\n",
        "label2id"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SrEgd4PZUgmF"
      },
      "source": [
        "data = data[[\"sentence\", \"word_labels\"]].drop_duplicates().reset_index(drop=True)\n",
        "data.head()"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "r3ArUiVRqw0C"
      },
      "source": [
        "len(data)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "U8obZumRTBrT"
      },
      "source": [
        "Let's verify that a random sentence and its corresponding tags are correct:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eUvupomW_fbe"
      },
      "source": [
        "data.iloc[0].sentence"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0dLyY3Oi_lvp"
      },
      "source": [
        "data.iloc[0].word_labels"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f5EHpuB78pIa"
      },
      "source": [
        "#### **Preparing the dataset and dataloader**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "15x7zmZnTgFx"
      },
      "source": [
        "Now that our data is preprocessed, we can turn it into PyTorch tensors such that we can provide it to the model. Let's start by defining some key variables that will be used later on in the training/evaluation process:"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "model_name = 'dmis-lab/biobert-base-cased-v1.2'"
      ],
      "metadata": {
        "id": "xnWzu1Vvm2HL"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lgNSM8Xz79Mg",
        "tags": []
      },
      "source": [
        "MAX_LEN = 128\n",
        "TRAIN_BATCH_SIZE = 4\n",
        "VALID_BATCH_SIZE = 2\n",
        "EPOCHS = 10\n",
        "LEARNING_RATE = 1e-05\n",
        "MAX_GRAD_NORM = 10\n",
        "tokenizer = BertTokenizer.from_pretrained(model_name)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wPYV2Ld6Tr5I"
      },
      "source": [
        "A tricky part of NER with BERT is that BERT relies on **wordpiece tokenization**, rather than word tokenization. This means that we should also define the labels at the wordpiece-level, rather than the word-level! \n",
        "\n",
        "For example, if you have word like \"Washington\" which is labeled as \"b-gpe\", but it gets tokenized to \"Wash\", \"##ing\", \"##ton\", then we will have to propagate the word’s original label to all of its wordpieces: \"b-gpe\", \"b-gpe\", \"b-gpe\". The model should be able to produce the correct labels for each individual wordpiece. The function below (taken from [here](https://github.com/chambliss/Multilingual_NER/blob/master/python/utils/main_utils.py#L118)) implements this.\n",
        "\n",
        "\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RNzSgZTfGUd8"
      },
      "source": [
        "def tokenize_and_preserve_labels(sentence, text_labels, tokenizer):\n",
        "    \"\"\"\n",
        "    Word piece tokenization makes it difficult to match word labels\n",
        "    back up with individual word pieces. This function tokenizes each\n",
        "    word one at a time so that it is easier to preserve the correct\n",
        "    label for each subword. It is, of course, a bit slower in processing\n",
        "    time, but it will help our model achieve higher accuracy.\n",
        "    \"\"\"\n",
        "\n",
        "    tokenized_sentence = []\n",
        "    labels = []\n",
        "\n",
        "    sentence = sentence.strip()\n",
        "\n",
        "    for word, label in zip(sentence.split(), text_labels.split(\",\")):\n",
        "\n",
        "        # Tokenize the word and count # of subwords the word is broken into\n",
        "        tokenized_word = tokenizer.tokenize(word)\n",
        "        n_subwords = len(tokenized_word)\n",
        "\n",
        "        # Add the tokenized word to the final tokenized word list\n",
        "        tokenized_sentence.extend(tokenized_word)\n",
        "\n",
        "        # Add the same label to the new list of labels `n_subwords` times\n",
        "        labels.extend([label] * n_subwords)\n",
        "\n",
        "    return tokenized_sentence, labels"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ez7qlFHl56ZW"
      },
      "source": [
        "Note that this is a **design decision**. You could also decide to only label the first wordpiece of each word and let the model only learn this (this is what was done in the original BERT paper, see Github discussion [here](https://github.com/huggingface/transformers/issues/64#issuecomment-443703063)). Another design decision could be to give the first wordpiece of each word the original word label, and then use the label “X” for all subsequent subwords of that word.\n",
        "\n",
        "All of them lead to good performance.\n",
        "\n",
        "Next, we define a regular PyTorch [dataset class](https://pytorch.org/docs/stable/data.html) (which transforms examples of a dataframe to PyTorch tensors). Here, each sentence gets tokenized, the special tokens that BERT expects are added, the tokens are padded or truncated based on the max length of the model, the attention mask is created and the labels are created based on the dictionary which we defined above. \n",
        "\n",
        "For more information about BERT's inputs, see [here](https://huggingface.co/transformers/glossary.html).  "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "aJty_Abw8_xK"
      },
      "source": [
        "class dataset(Dataset):\n",
        "    def __init__(self, dataframe, tokenizer, max_len):\n",
        "        self.len = len(dataframe)\n",
        "        self.data = dataframe\n",
        "        self.tokenizer = tokenizer\n",
        "        self.max_len = max_len\n",
        "        \n",
        "    def __getitem__(self, index):\n",
        "        # step 1: tokenize (and adapt corresponding labels)\n",
        "        sentence = self.data.sentence[index]  \n",
        "        word_labels = self.data.word_labels[index]  \n",
        "        tokenized_sentence, labels = tokenize_and_preserve_labels(sentence, word_labels, self.tokenizer)\n",
        "        \n",
        "        # step 2: add special tokens (and corresponding labels)\n",
        "        tokenized_sentence = [\"[CLS]\"] + tokenized_sentence + [\"[SEP]\"] # add special tokens\n",
        "        labels.insert(0, \"KEEP\") # add outside label for [CLS] token\n",
        "        labels.insert(-1, \"KEEP\") # add outside label for [SEP] token\n",
        "\n",
        "        # step 3: truncating/padding\n",
        "        maxlen = self.max_len\n",
        "\n",
        "        if (len(tokenized_sentence) > maxlen):\n",
        "          # truncate\n",
        "          tokenized_sentence = tokenized_sentence[:maxlen]\n",
        "          labels = labels[:maxlen]\n",
        "        else:\n",
        "          # pad\n",
        "          tokenized_sentence = tokenized_sentence + ['[PAD]'for _ in range(maxlen - len(tokenized_sentence))]\n",
        "          labels = labels + [\"KEEP\" for _ in range(maxlen - len(labels))]\n",
        "\n",
        "        # step 4: obtain the attention mask\n",
        "        attn_mask = [1 if tok != '[PAD]' else 0 for tok in tokenized_sentence]\n",
        "        \n",
        "        # step 5: convert tokens to input ids\n",
        "        ids = self.tokenizer.convert_tokens_to_ids(tokenized_sentence)\n",
        "\n",
        "        label_ids = [label2id[label] for label in labels]\n",
        "        # the following line is deprecated\n",
        "        #label_ids = [label if label != 0 else -100 for label in label_ids]\n",
        "        \n",
        "        return {\n",
        "              'ids': torch.tensor(ids, dtype=torch.long),\n",
        "              'mask': torch.tensor(attn_mask, dtype=torch.long),\n",
        "              #'token_type_ids': torch.tensor(token_ids, dtype=torch.long),\n",
        "              'targets': torch.tensor(label_ids, dtype=torch.long)\n",
        "        } \n",
        "    \n",
        "    def __len__(self):\n",
        "        return self.len"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hTP7zuWGWGUd"
      },
      "source": [
        "Now, based on the class we defined above, we can create 2 datasets, one for training and one for testing. Let's use a 80/20 split:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jrkdZBLYHVcB",
        "tags": []
      },
      "source": [
        "train_size = 0.8\n",
        "train_dataset = data.sample(frac=train_size,random_state=200)\n",
        "test_dataset = data.drop(train_dataset.index).reset_index(drop=True)\n",
        "train_dataset = train_dataset.reset_index(drop=True)\n",
        "\n",
        "print(\"FULL Dataset: {}\".format(data.shape))\n",
        "print(\"TRAIN Dataset: {}\".format(train_dataset.shape))\n",
        "print(\"TEST Dataset: {}\".format(test_dataset.shape))\n",
        "\n",
        "training_set = dataset(train_dataset, tokenizer, MAX_LEN)\n",
        "testing_set = dataset(test_dataset, tokenizer, MAX_LEN)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ptv5AT_iTb7W"
      },
      "source": [
        "Let's have a look at the first training example:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "phmPylgAm8Xy"
      },
      "source": [
        "training_set[0]"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VvU4nzL2W2Xo"
      },
      "source": [
        "Let's verify that the input ids and corresponding targets are correct:"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "training_set[0][\"ids\"]"
      ],
      "metadata": {
        "id": "ZHoufyakY18x"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "DWgnNJrYW2GP",
        "tags": []
      },
      "source": [
        "# print the first 30 tokens and corresponding labels\n",
        "for token, label in zip(tokenizer.convert_ids_to_tokens(training_set[0][\"ids\"][:30]), training_set[0][\"targets\"][:30]):\n",
        "  print('{0:10}  {1}'.format(token, id2label[label.item()]))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Ky68FcTgWnfN"
      },
      "source": [
        "Now, let's define the corresponding PyTorch dataloaders:"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "KIw793myWOmi"
      },
      "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)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "73OzU7oXRxR8"
      },
      "source": [
        "#### **Defining the model**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "T-iGhnhdLNdP"
      },
      "source": [
        "Here we define the model, BertForTokenClassification, and load it with the pretrained weights of \"bert-base-uncased\". The only thing we need to additionally specify is the number of labels (as this will determine the architecture of the classification head).\n",
        "\n",
        "Note that only the base layers are initialized with the pretrained weights. The token classification head of top has just randomly initialized weights, which we will train, together with the pretrained weights, using our labelled dataset. This is also printed as a warning when you run the code cell below.\n",
        "\n",
        "Then, we move the model to the GPU."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cB9MR3KcWXUs",
        "tags": []
      },
      "source": [
        "model = BertForTokenClassification.from_pretrained(model_name,\n",
        "                                                   num_labels=len(id2label),\n",
        "                                                   id2label=id2label,\n",
        "                                                   label2id=label2id)\n",
        "model.to(device)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Pp7Yl4JyWhDj"
      },
      "source": [
        "#### **Training the model**\n",
        "\n",
        "Before training the model, let's perform a sanity check, which I learned thanks to Andrej Karpathy's wonderful [cs231n course](http://cs231n.stanford.edu/) at Stanford (see also his [blog post about debugging neural networks](http://karpathy.github.io/2019/04/25/recipe/)). The initial loss of your model should be close to -ln(1/number of classes) = -ln(1/17) = 2.83. \n",
        "\n",
        "Why? Because we are using cross entropy loss. The cross entropy loss is defined as -ln(probability score of the model for the correct class). In the beginning, the weights are random, so the probability distribution for all of the classes for a given token will be uniform, meaning that the probability for the correct class will be near 1/17. The loss for a given token will thus be -ln(1/17). As PyTorch's [CrossEntropyLoss](https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html) (which is used by `BertForTokenClassification`) uses *mean reduction* by default, it will compute the mean loss for each of the tokens in the sequence (in other words, for all of the 512 tokens). The mean of 512 times -log(1/17) is, you guessed it, -log(1/17).  \n",
        "\n",
        "Let's verify this:\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "eqAN7YVIjKTr"
      },
      "source": [
        "ids = training_set[0][\"ids\"].unsqueeze(0)\n",
        "mask = training_set[0][\"mask\"].unsqueeze(0)\n",
        "targets = training_set[0][\"targets\"].unsqueeze(0)\n",
        "ids = ids.to(device)\n",
        "mask = mask.to(device)\n",
        "targets = targets.to(device)\n",
        "outputs = model(input_ids=ids, attention_mask=mask, labels=targets)\n",
        "initial_loss = outputs[0]\n",
        "initial_loss"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yLdwsru9Mh7U"
      },
      "source": [
        "This looks good. Let's also verify that the logits of the neural network have a shape of (batch_size, sequence_length, num_labels):"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "X-z6YCpGnvfj"
      },
      "source": [
        "tr_logits = outputs[1]\n",
        "tr_logits.shape"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kwDLXxOVOCvD"
      },
      "source": [
        "Next, we define the optimizer. Here, we are just going to use Adam with a default learning rate. One can also decide to use more advanced ones such as AdamW (Adam with weight decay fix), which is [included](https://huggingface.co/transformers/main_classes/optimizer_schedules.html) in the Transformers repository, and a learning rate scheduler, but we are not going to do that here."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "kznSQfGIWdU4"
      },
      "source": [
        "optimizer = torch.optim.Adam(params=model.parameters(), lr=LEARNING_RATE)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vZQ8JMF0NOe1"
      },
      "source": [
        "Now let's define a regular PyTorch training function. It is partly based on [a really good repository about multilingual NER](https://github.com/chambliss/Multilingual_NER/blob/master/python/utils/main_utils.py#L344)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GLFivpkwW1HY"
      },
      "source": [
        "# Defining the training function on the 80% of the dataset for tuning the bert model\n",
        "def train(epoch):\n",
        "    tr_loss, tr_accuracy = 0, 0\n",
        "    nb_tr_examples, nb_tr_steps = 0, 0\n",
        "    tr_preds, tr_labels = [], []\n",
        "    # put model in training mode\n",
        "    model.train()\n",
        "    \n",
        "    for idx, batch in enumerate(training_loader):\n",
        "        \n",
        "        ids = batch['ids'].to(device, dtype = torch.long)\n",
        "        mask = batch['mask'].to(device, dtype = torch.long)\n",
        "        targets = batch['targets'].to(device, dtype = torch.long)\n",
        "\n",
        "        outputs = model(input_ids=ids, attention_mask=mask, labels=targets)\n",
        "        loss, tr_logits = outputs.loss, outputs.logits\n",
        "        tr_loss += loss.item()\n",
        "\n",
        "        nb_tr_steps += 1\n",
        "        nb_tr_examples += targets.size(0)\n",
        "        \n",
        "        if idx % 100==0:\n",
        "            loss_step = tr_loss/nb_tr_steps\n",
        "            print(f\"Training loss per 100 training steps: {loss_step}\")\n",
        "           \n",
        "        # compute training accuracy\n",
        "        flattened_targets = targets.view(-1) # shape (batch_size * seq_len,)\n",
        "        active_logits = tr_logits.view(-1, model.num_labels) # shape (batch_size * seq_len, num_labels)\n",
        "        flattened_predictions = torch.argmax(active_logits, axis=1) # shape (batch_size * seq_len,)\n",
        "        # now, use mask to determine where we should compare predictions with targets (includes [CLS] and [SEP] token predictions)\n",
        "        active_accuracy = mask.view(-1) == 1 # active accuracy is also of shape (batch_size * seq_len,)\n",
        "        targets = torch.masked_select(flattened_targets, active_accuracy)\n",
        "        predictions = torch.masked_select(flattened_predictions, active_accuracy)\n",
        "        \n",
        "        tr_preds.extend(predictions)\n",
        "        tr_labels.extend(targets)\n",
        "        \n",
        "        tmp_tr_accuracy = accuracy_score(targets.cpu().numpy(), predictions.cpu().numpy())\n",
        "        tr_accuracy += tmp_tr_accuracy\n",
        "    \n",
        "        # gradient clipping\n",
        "        torch.nn.utils.clip_grad_norm_(\n",
        "            parameters=model.parameters(), max_norm=MAX_GRAD_NORM\n",
        "        )\n",
        "        \n",
        "        # backward pass\n",
        "        optimizer.zero_grad()\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "\n",
        "    epoch_loss = tr_loss / nb_tr_steps\n",
        "    tr_accuracy = tr_accuracy / nb_tr_steps\n",
        "    print(f\"Training loss epoch: {epoch_loss}\")\n",
        "    print(f\"Training accuracy epoch: {tr_accuracy}\")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "k2dsCyP7dcF3"
      },
      "source": [
        "And let's train the model!"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "y07Ybw8rZeZ7",
        "tags": []
      },
      "source": [
        "for epoch in range(EPOCHS):\n",
        "    print(f\"Training epoch: {epoch + 1}\")\n",
        "    train(epoch)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r4jcSOJr680a"
      },
      "source": [
        "#### **Evaluating the model**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rYUTuOEUdfFJ"
      },
      "source": [
        "Now that we've trained our model, we can evaluate its performance on the held-out test set (which is 20% of the data). Note that here, no gradient updates are performed, the model just outputs its logits. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "RIVVfFHi7Aw7"
      },
      "source": [
        "def valid(model, testing_loader):\n",
        "    # put model in evaluation mode\n",
        "    model.eval()\n",
        "    \n",
        "    eval_loss, eval_accuracy = 0, 0\n",
        "    nb_eval_examples, nb_eval_steps = 0, 0\n",
        "    eval_preds, eval_labels = [], []\n",
        "    \n",
        "    with torch.no_grad():\n",
        "        for idx, batch in enumerate(testing_loader):\n",
        "            \n",
        "            ids = batch['ids'].to(device, dtype = torch.long)\n",
        "            mask = batch['mask'].to(device, dtype = torch.long)\n",
        "            targets = batch['targets'].to(device, dtype = torch.long)\n",
        "            \n",
        "            outputs = model(input_ids=ids, attention_mask=mask, labels=targets)\n",
        "            loss, eval_logits = outputs.loss, outputs.logits\n",
        "            \n",
        "            eval_loss += loss.item()\n",
        "\n",
        "            nb_eval_steps += 1\n",
        "            nb_eval_examples += targets.size(0)\n",
        "        \n",
        "            if idx % 100==0:\n",
        "                loss_step = eval_loss/nb_eval_steps\n",
        "                print(f\"Validation loss per 100 evaluation steps: {loss_step}\")\n",
        "              \n",
        "            # compute evaluation accuracy\n",
        "            flattened_targets = targets.view(-1) # shape (batch_size * seq_len,)\n",
        "            active_logits = eval_logits.view(-1, model.num_labels) # shape (batch_size * seq_len, num_labels)\n",
        "            flattened_predictions = torch.argmax(active_logits, axis=1) # shape (batch_size * seq_len,)\n",
        "            # now, use mask to determine where we should compare predictions with targets (includes [CLS] and [SEP] token predictions)\n",
        "            active_accuracy = mask.view(-1) == 1 # active accuracy is also of shape (batch_size * seq_len,)\n",
        "            targets = torch.masked_select(flattened_targets, active_accuracy)\n",
        "            predictions = torch.masked_select(flattened_predictions, active_accuracy)\n",
        "            \n",
        "            eval_labels.extend(targets)\n",
        "            eval_preds.extend(predictions)\n",
        "            \n",
        "            tmp_eval_accuracy = accuracy_score(targets.cpu().numpy(), predictions.cpu().numpy())\n",
        "            eval_accuracy += tmp_eval_accuracy\n",
        "    \n",
        "    #print(eval_labels)\n",
        "    #print(eval_preds)\n",
        "\n",
        "    labels = [id2label[id.item()] for id in eval_labels]\n",
        "    predictions = [id2label[id.item()] for id in eval_preds]\n",
        "\n",
        "    #print(labels)\n",
        "    #print(predictions)\n",
        "    \n",
        "    eval_loss = eval_loss / nb_eval_steps\n",
        "    eval_accuracy = eval_accuracy / nb_eval_steps\n",
        "    print(f\"Validation Loss: {eval_loss}\")\n",
        "    print(f\"Validation Accuracy: {eval_accuracy}\")\n",
        "\n",
        "    return labels, predictions"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zJaONluRdq-e"
      },
      "source": [
        "As we can see below, performance is quite good! Accuracy on the test test is > 93%."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2BrxRjvxApY8",
        "tags": []
      },
      "source": [
        "labels, predictions = valid(model, testing_loader)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SAznLDwx_U2X"
      },
      "source": [
        "However, the accuracy metric is misleading, as a lot of labels are \"outside\" (O), even after omitting predictions on the [PAD] tokens. What is important is looking at the precision, recall and f1-score of the individual tags. For this, we use the seqeval Python library: "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0jDNXrjr-6BW",
        "tags": []
      },
      "source": [
        "from seqeval.metrics import classification_report\n",
        "\n",
        "print(classification_report([labels], [predictions]))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4Gz-wHAw3xMk"
      },
      "source": [
        "#### **Inference**\n",
        "\n",
        "The fun part is when we can quickly test the model on new, unseen sentences. \n",
        "Here, we use the prediction of the **first word piece of every word**. Note that the function we used to train our model (`tokenze_and_preserve_labels`) propagated the label to all subsequent word pieces (so you could for example also perform a majority vote on the predicted labels of all word pieces of a word).\n",
        "\n",
        "*In other words, the code below does not take into account when predictions of different word pieces that belong to the same word do not match.*"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def get_gt(sentence, wl):\n",
        "  sentence = sentence.split()\n",
        "  wl = wl.split(\",\")\n",
        "  i = 0\n",
        "  string = \"\"\n",
        "  for tag in wl:\n",
        "    if tag == \"KEEP\":\n",
        "      string += sentence[i] + \" \"\n",
        "    i+= 1\n",
        "\n",
        "  return string.strip()"
      ],
      "metadata": {
        "id": "DE0ljYggb98Q"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from transformers import pipeline\n",
        "\n",
        "results = []\n",
        "\n",
        "pipe = pipeline(task=\"token-classification\", model=model.to(\"cpu\"), tokenizer=tokenizer, aggregation_strategy=\"simple\")\n",
        "for i in range(len(test_dataset['sentence'])):\n",
        "    sentence = test_dataset['sentence'][i]\n",
        "    # print(\"SENTENCE: \" + sentence)\n",
        "    gt = get_gt(sentence, test_dataset['word_labels'][i])\n",
        "    # print(\"GROUND TRUTH: \" + gt)\n",
        "    p = pipe(sentence)\n",
        "    string = \"\"\n",
        "    for item in p:\n",
        "        # print(type(item), item)\n",
        "        if item['entity_group'] == 'KEEP':\n",
        "          string += item['word'] + \" \"\n",
        "    # print(\"PREDICTION : \" + string.strip())\n",
        "    results.append([sentence, gt, string.strip().replace(\"redemonstrate\", \"demonstrate\")])\n",
        "    # print(\"============================\")"
      ],
      "metadata": {
        "id": "D5KB5TKRcdRT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import string, difflib, regex as re\n",
        "def get_annot(orig,gt):\n",
        "    remove = string.punctuation\n",
        "    remove = remove.replace(\"_\", \"\")\n",
        "    pattern = r\"[{}]\".format(remove)\n",
        "\n",
        "    orig = re.sub(pattern, \"\", orig)\n",
        "    gt = re.sub(pattern, \"\", gt)\n",
        "\n",
        "    orig_gt = difflib.unified_diff(orig.lower().split(), gt.lower().split())\n",
        "    orig_gt_diff = []\n",
        "\n",
        "    for line in orig_gt:\n",
        "        if (line[0] == \"+\" or line[0] == \"-\") and (\"+++\" not in line and \"---\" not in line) and (not line.isspace()):\n",
        "            orig_gt_diff.append(line)\n",
        "    \n",
        "    ret = \"\"\n",
        "\n",
        "    s = set()\n",
        "    \n",
        "    for item in orig.split():\n",
        "        if \"-\" + item in orig_gt_diff and item not in s:\n",
        "            ret += \"REMOVE \"\n",
        "            s.add(item)\n",
        "        else:\n",
        "            ret += \"KEEP \"\n",
        "    \n",
        "    return ret.strip().split()"
      ],
      "metadata": {
        "id": "zvpOLPx2gL1s"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "true_y = []\n",
        "pred_y = []\n",
        "\n",
        "for arr in results:\n",
        "    orig = arr[0]\n",
        "    gt = arr[1]\n",
        "    mod = arr[2]\n",
        "    print(\"ORIGINAL    : \" + orig)\n",
        "    print(\"GROUND TRUTH: \" + gt)\n",
        "    print(\"BERT PREDICT: \" + mod)\n",
        "    print(\"\\n===========\\n\")\n",
        "    orig_gt = get_annot(orig,gt)\n",
        "    orig_mod = get_annot(orig,mod)\n",
        "    for i in range(len(orig_gt)):\n",
        "      if orig_gt[i] == \"KEEP\":\n",
        "        true_y.append(0)\n",
        "      else:\n",
        "        true_y.append(1)\n",
        "      if orig_mod[i] == \"KEEP\":\n",
        "        pred_y.append(0)\n",
        "      else:\n",
        "        pred_y.append(1)"
      ],
      "metadata": {
        "id": "TebRjqvWf-ap"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay"
      ],
      "metadata": {
        "id": "7q5Mu0h7fRAg"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from matplotlib import pyplot as plt\n",
        "cm = confusion_matrix(true_y, pred_y)\n",
        "# print(cm)\n",
        "disp = ConfusionMatrixDisplay(confusion_matrix=cm)\n",
        "disp.plot()\n",
        "plt.show()\n",
        "precision = cm[1][1] / (cm[1][1] + cm[0][1])\n",
        "recall = cm[1][1] / (cm[1][1] + cm[1][0])\n",
        "F1 = 2 * precision * recall / (precision + recall)\n",
        "print(\"F1: \" + str(F1))"
      ],
      "metadata": {
        "id": "qPQ4JnrVfSlh"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sqDklprSqB5d"
      },
      "source": [
        "#### **Saving the model for future use**"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VuUdX_fImswO"
      },
      "source": [
        "Finally, let's save the model and tokenizer files such that we can easily re-use them later on. There are 2 options:\n",
        "\n",
        "* you can save everything locally, simply by calling `model.save_pretrained()` and `tokenizer.save_pretrained()`, providing a directory path as argument. \n",
        "* you can push the files to the [HuggingFace hub](https://huggingface.co/). This way, you can share your model with the community/your colleagues. All files will be tracked by git, as each model on the hub has its own git repo.\n",
        "\n",
        "Both options allow to re-use the model/tokenizer using the `from_pretrained()` method. Here we'll do the latter. \n",
        "\n",
        "To upload a model to the hub, 2 things need to be setup:\n",
        "* install git-LFS, which is used by the hub\n",
        "* set up authentication token\n"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "!sudo apt-get install git-lfs"
      ],
      "metadata": {
        "id": "ZVw5wjcwfGHd"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from huggingface_hub import notebook_login\n",
        "\n",
        "notebook_login()"
      ],
      "metadata": {
        "id": "yxx9U4VgfJ7L"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sDZtSsKKntuI",
        "tags": []
      },
      "source": [
        "mname = \"gilbert\"\n",
        "\n",
        "tokenizer.push_to_hub(\n",
        "    repo_path_or_name=mname,\n",
        "    organization=\"rajpurkarlab\",\n",
        "    commit_message=\"Add tokenizer\",\n",
        "    use_temp_dir=True,\n",
        ")\n",
        "model.push_to_hub(\n",
        "    repo_path_or_name=mname,\n",
        "    organization=\"rajpurkarlab\",\n",
        "    commit_message=\"Add model\",\n",
        "    use_temp_dir=True,\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "You can check the model on the hub: https://huggingface.co/rajpurkarlab/gilbert. Awesome, isn't it? Check the \"files and versions\" tab, it includes all our files :) we can now load it back as follows (here I'm using the [Auto API](https://huggingface.co/docs/transformers/model_doc/auto) - which will load the appropriate tokenizer + model for us):"
      ],
      "metadata": {
        "id": "ZlPINMg3hQQX"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from transformers import AutoTokenizer, AutoModelForTokenClassification\n",
        "\n",
        "modelname = \"rajpurkarlab/gilbert\"\n",
        "tokenizer = AutoTokenizer.from_pretrained(modelname)\n",
        "model = AutoModelForTokenClassification.from_pretrained(modelname)"
      ],
      "metadata": {
        "id": "lGTAfo8bhXrW"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}