entity_recognition/entity_recognition_training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d0d2e007",
   "metadata": {},
   "source": [
    "# Training Entity Recognition Model for Token Classification"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5432fca6",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/cleanlab/examples/blob/master/entity_recognition/entity_recognition_training.ipynb) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6676421",
   "metadata": {},
   "source": [
    "This notebook demonstrates how to train a NLP model for entity recognition and use it to produce out-of-sample predicted probabilities for each token. These are required inputs to find label issues in token classification datasets with cleanlab. The specific token classification task we consider here is Named Entity Recognition with the [CoNLL-2003 dataset](https://deepai.org/dataset/conll-2003-english), and we train a Transformer network from [HuggingFace's transformers library](https://github.com/huggingface/transformers). This notebook demonstrates how to produce the `pred_probs`, using them to find label issues is demonstrated in cleanlab's [Token Classification Tutorial](https://docs.cleanlab.ai/stable/tutorials/token_classification.html). \n",
    "\n",
    "Note: running this notebook requires the **.py** files from the **entity_recognition/** parent folder, if running in Colab or locally, make sure you've copied these helper **.py** files to your environment as well. \n",
    "\n",
    "**Overview of what we'll do in this notebook:** \n",
    "- Read and process text datasets with per-token labels in the CoNLL format. \n",
    "- Compute out-of-sample predicted probabilities by training a BERT Transformer network via cross-validation. \n",
    "- Aggregate subword-level tokens\\* into word-level tokens which are more individually meaningful. \n",
    "\n",
    "\\* In NLP, tokens typically refer to words or punctuation marks, but modern tokenizers used with Transformers often break down longer words into smaller subwords. To avoid confusion, we use \"word-level tokens\" to refer to the individual given tokens in the original dataset (with a separate class label provided for each such token). Tokens obtained from processing the raw text with tokenizers are referred as \"subword-level tokens\". "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1da020bc",
   "metadata": {},
   "source": [
    "## 1. Fetch data and load required dependencies"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89f03a66",
   "metadata": {},
   "source": [
    "Please install the dependencies specified in this [requirements.txt](https://github.com/cleanlab/examples/blob/master/entity_recognition/requirements.txt) file before running the notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9018de0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget -nc https://data.deepai.org/conll2003.zip && mkdir -p data \n",
    "!unzip conll2003.zip -d data/ && rm conll2003.zip "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a1349304",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Package versions we used: tqdm==4.64.0 transformers==4.22.0.dev0 numpy==1.23.0 sklearn==1.1.1 \n",
    "import os \n",
    "import warnings \n",
    "import string \n",
    "import numpy as np \n",
    "from tqdm import tqdm \n",
    "from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline \n",
    "from sklearn.metrics import balanced_accuracy_score \n",
    "import nltk \n",
    "from cleanlab.internal.token_classification_utils import get_sentence, filter_sentence, process_token, mapping, merge_probs \n",
    "from bert import Ner \n",
    "from token_classification_tutorial_utils import create_folds, modified, get_pred_probs, to_dict \n",
    "from run_ner import train \n",
    "\n",
    "nltk.download(\"punkt\") \n",
    "warnings.filterwarnings(\"ignore\") \n",
    "\n",
    "# Disable `TOKENIZERS_PARALLELISM` if multiple processors exist. \n",
    "if os.cpu_count() > 1: \n",
    "    os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\" "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99e32604",
   "metadata": {},
   "source": [
    "## 2. Load the CONLL-2003 dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d80cd5a9",
   "metadata": {},
   "source": [
    "CONLL-2003 data are in the following format: \n",
    "\n",
    "`-DOCSTART- -X- -X- O` \n",
    "\n",
    "`[word] [pos_tags] [chunk_tags] [ner_tags]` <- Start of first sentence \n",
    "\n",
    "`...`\n",
    "\n",
    "`[word] [pos_tags] [chunk_tags] [ner_tags]` \n",
    "\n",
    "`[empty line]` \n",
    "\n",
    "`[word] [pos_tags] [chunk_tags] [ner_tags]` <- Start of second sentence \n",
    "\n",
    "`...`\n",
    "\n",
    "`[word] [pos_tags] [chunk_tags] [ner_tags]` \n",
    "\n",
    "Here we consider the `ner_tags` (named-entity recognition tags stored in the IOB2 format), which include the following classes: \n",
    "\n",
    "| `ner_tags` |             Description              |\n",
    "|:----------:|:------------------------------------:|\n",
    "|     `O`    |      Other (not a named entity)      |\n",
    "|   `B-MIS`  | Beginning of a miscellaneous entity  |\n",
    "|   `I-MIS`  |         Miscellaneous entity         |\n",
    "|   `B-PER`  |     Beginning of a person entity     |\n",
    "|   `I-PER`  |            Person entity             |\n",
    "|   `B-ORG`  | Beginning of an organization entity  |\n",
    "|   `I-ORG`  |         Organization entity          |\n",
    "|   `B-LOC`  |    Beginning of a location entity    |\n",
    "|   `I-LOC`  |           Location entity            | \n",
    "\n",
    "For more information, see [here](https://paperswithcode.com/dataset/conll-2003). We cast all-caps words into lowercase except for the first character (eg. `JAPAN` -> `Japan`), to discourage the tokenizer from breaking such words into multiple subwords."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "871730b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "filepaths = ['data/train.txt', 'data/valid.txt', 'data/test.txt'] \n",
    "entities = ['O', 'B-MISC', 'I-MISC', 'B-PER', 'I-PER', 'B-ORG', 'I-ORG', 'B-LOC', 'I-LOC']\n",
    "entity_map = {entity: i for i, entity in enumerate(entities)} \n",
    "\n",
    "# This code is adapted from: https://github.com/kamalkraj/BERT-NER/blob/dev/run_ner.py \n",
    "def readfile(filepath, sep=' '): \n",
    "    lines = open(filepath)\n",
    "    \n",
    "    data, sentence, label = [], [], []\n",
    "    for line in lines:\n",
    "        if len(line) == 0 or line.startswith('-DOCSTART') or line[0] == '\\n':\n",
    "            if len(sentence) > 0:\n",
    "                data.append((sentence, label))\n",
    "                sentence, label = [], []\n",
    "            continue\n",
    "        splits = line.split(sep) \n",
    "        word = splits[0]\n",
    "        if len(word) > 0 and word[0].isalpha() and word.isupper():\n",
    "            word = word[0] + word[1:].lower()\n",
    "        sentence.append(word)\n",
    "        label.append(entity_map[splits[-1][:-1]])\n",
    "\n",
    "    if len(sentence) > 0:\n",
    "        data.append((sentence, label))\n",
    "        \n",
    "    given_words = [d[0] for d in data] \n",
    "    given_labels = [d[1] for d in data] \n",
    "    \n",
    "    return given_words, given_labels \n",
    "\n",
    "given_words, given_labels = [], [] \n",
    "\n",
    "for filepath in filepaths: \n",
    "    words, labels = readfile(filepath) \n",
    "    given_words.extend(words) \n",
    "    given_labels.extend(labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12148277",
   "metadata": {},
   "source": [
    "`given_words` and `given_labels` above are strings/labels corresponding to each word-level token, represented as nested lists."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c1d505f4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Word\t\tLabel\tEntity\n",
      "-------------------------------\n",
      "Eu              5\tB-ORG     \n",
      "rejects         0\tO         \n",
      "German          1\tB-MISC    \n",
      "call            0\tO         \n",
      "to              0\tO         \n",
      "boycott         0\tO         \n",
      "British         1\tB-MISC    \n",
      "lamb            0\tO         \n",
      ".               0\tO         \n"
     ]
    }
   ],
   "source": [
    "i = 0  # change this to view a different example from the dataset \n",
    "\n",
    "print('Word\\t\\tLabel\\tEntity\\n-------------------------------') \n",
    "for word, label in zip(given_words[i], given_labels[i]): \n",
    "    print('{:14s}{:3d}\\t{:10s}'.format(word, label, entities[label])) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16f8639b",
   "metadata": {},
   "source": [
    "We next apply minor pre-processing for readability. Sentences containing the `#` character are removed for simplicity, because this special character is later used to represent subword-tokens by the sentence tokenizers used in HuggingFace (See section 4). We also remove single token sentences with `len(sentence) <= 1`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "8f19eeae",
   "metadata": {},
   "outputs": [],
   "source": [
    "sentences = list(map(get_sentence, given_words)) \n",
    "\n",
    "sentences, mask = filter_sentence(sentences) \n",
    "given_words = [words for m, words in zip(mask, given_words) if m] \n",
    "given_labels = [labels for m, labels in zip(mask, given_labels) if m] "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "50c496b8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of sentences: 20718\n",
      "Eu rejects German call to boycott British lamb.\n"
     ]
    }
   ],
   "source": [
    "print('Number of sentences: %d' % len(sentences)) \n",
    "print(sentences[0])  # display first sentence in the processed dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0fb1b2b0",
   "metadata": {},
   "source": [
    "## 3. Train Token Classification Model using Cross-Validation \n",
    "\n",
    "To later find label issues in the training dataset, we first compute out-of-sample predicted probabilities (`pred_probs`) using cross-validation. We start by partitioning the dataset into `k = 5` disjoint folds: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "797df997",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "'folds/' already exists, skipping...\n"
     ]
    }
   ],
   "source": [
    "lines = [[]] \n",
    "for filepath in filepaths: \n",
    "    for line in open(filepath) : \n",
    "        if len(line) == 0 or line.startswith('-DOCSTART') or line[0] == '\\n':\n",
    "            if len(lines[-1]) > 0: \n",
    "                lines.append([]) \n",
    "        else: \n",
    "            lines[-1].append(line) \n",
    "        \n",
    "lines = lines[:-1] \n",
    "lines = [line for m, line in zip(mask, lines) if m] \n",
    "\n",
    "k = 5 \n",
    "indices = create_folds(lines, k=k) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "304653aa",
   "metadata": {},
   "source": [
    "We train one model for each fold's training/testing pair: \n",
    "\n",
    "- Warning! The following code will take a long time to execute, and is recommended to run on GPU, otherwise it will take forever. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "6c877f7e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model 0 already exists, skipping...\n",
      "Model 1 already exists, skipping...\n",
      "Model 2 already exists, skipping...\n",
      "Model 3 already exists, skipping...\n",
      "Model 4 already exists, skipping...\n"
     ]
    }
   ],
   "source": [
    "for i in range(k): \n",
    "    if os.path.exists('folds/fold%d/model/' % i): \n",
    "        print('Model %d already exists, skipping...' % i) \n",
    "    else: \n",
    "        print('Training model on fold %d (out of %d) of cross-validation...' % (i, k)) \n",
    "        train(data_dir='folds/fold%d' % i, \n",
    "              bert_model='bert-base-cased', \n",
    "              task_name='ner', \n",
    "              output_dir='folds/fold%d/model' % i, \n",
    "              max_seq_length=256, \n",
    "              do_train=True, \n",
    "              num_train_epochs=10, \n",
    "              warmup_proportion=0.1,\n",
    "              train_batch_size=16\n",
    "        ) \n",
    "        print('Model %d saved' % i) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80621481",
   "metadata": {},
   "source": [
    "## 4. Compute Out-of-Sample Predicted Probabilities \n",
    "\n",
    "We obtain predicted class probabilities for each token using the model in which this token was held out during training (i.e. its sentence was part of the validation fold above). Note that most modern tokenizers break sentences into subword-level tokens (smaller units than word-level tokens). For example, the following sentence: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "560f8626",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Eu rejects German call to boycott British lamb.\n"
     ]
    }
   ],
   "source": [
    "i = 0  # change this to view a different example from the dataset \n",
    "print(sentences[i]) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c97d0119",
   "metadata": {},
   "source": [
    "is tokenized into: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08446655",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Ner(\"folds/fold0/model/\") \n",
    "tokens = model.tokenize(sentences[i])[0] \n",
    "print(tokens) \n",
    "tokens = [token.replace('#', '') for token in tokens] "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "493203d3",
   "metadata": {},
   "source": [
    "`##` here is appended by the tokenizer to indicate that the token is a subword. Let's collect both the predicted class probabilities and strings for each token in each sentence. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9fac29e",
   "metadata": {},
   "outputs": [],
   "source": [
    "sentence_tokens, sentence_probs = {}, {} \n",
    "for i in range(k): \n",
    "    model = Ner(\"folds/fold%d/model/\" % i) \n",
    "    for index in tqdm(indices[i]): \n",
    "        sentence_probs[index], sentence_tokens[index] = model.predict(sentences[index]) \n",
    "        \n",
    "sentence_tokens = [sentence_tokens[i] for i in range(len(sentences))] \n",
    "sentence_probs = [sentence_probs[i] for i in range(len(sentences))] "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "754a24f4",
   "metadata": {},
   "source": [
    "Most tokenizers partition sentences into subword-level tokens without altering the characters. However, you should verify whether any characters are modified, particularly for edge cases such as single or double quotations. In this example, the tokenizer has broken down double quotations `\"` into two `'`s or two `` ` ``s. The following code checks if any characters in the sentences were modified. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "237ff9c2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "print(modified(given_words, sentence_tokens)) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b498649",
   "metadata": {},
   "source": [
    "Given that some characters were modified, we should map `sentence_tokens` back to the characters from the original dataset, so that we can better compare predicted labels with the given labels to spot label issues. This mapping differs between different models, and may not be required by many tokenizers. We should not work with modified tokens directly because we lack their given labels. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "a09fe4c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Code to map `sentence_tokens` back to the characters from the original text \n",
    "replace = [('#', ''), ('``', '\"'), (\"''\", '\"')] \n",
    "sentence_tokens = [[process_token(token, replace) for token in sentence_tokens[i]] for i in range(len(sentences))] \n",
    "\n",
    "for i in range(len(sentences)): \n",
    "    short = ''.join(given_words[i]) \n",
    "    if \"''\" in short: \n",
    "        processed_tokens, processed_probs = [], [] \n",
    "        for token, prob in zip(sentence_tokens[i], sentence_probs[i]): \n",
    "            if token != '\"': \n",
    "                processed_tokens.append(token) \n",
    "                processed_probs.append(prob) \n",
    "            else: \n",
    "                for _ in range(2): \n",
    "                    processed_tokens.append(\"'\") \n",
    "                    processed_probs.append(prob) \n",
    "        sentence_tokens[i] = [token for token in processed_tokens] \n",
    "        sentence_probs[i] = np.array([prob for prob in processed_probs]) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00edbe08",
   "metadata": {},
   "source": [
    "Here we are more interested in severe types of mislabels, such as `B-LOC` vs. `B-PER`, instead of `B-LOC` vs. `I-LOC`. Therefore, we discard the `B-` and `I-` prefixes, and compute model predicted probabilities over this reduced set of classes for each subword-level token. The merged entities are `[O, MIS, PER, ORG, LOC]`, which correspond to the classes in our token classification task. In cleanlab's Token Classification Tutorial, we use these probabilistic predictions produced here to identify instances where the class label was incorrectly chosen for particular tokens. As shown below:\n",
    "\n",
    "- `given_maps` is an array of length equal to the original number of entities, such that `given_maps[i]` is the mapped entity of the i'th entity \n",
    "- `model_maps` is an array of length equal to the number of model predicted labels, such that `model_maps[i]` is the mapped entity of the i'th predicted entity (according to our trained model). If `model_maps[i] < 0`, it indicates that the entity does not map to a valid named entity. This usually occurs when the model predicted entities include start/end tags, which we do not consider here. If `np.any(model_maps < 0)`, `pred_probs` will be re-normalized to sum to one over the relevant classes considered here. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "b2cc02e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "given_maps = [0, 1, 1, 2, 2, 3, 3, 4, 4] \n",
    "model_maps = [-1, 0, 1, 1, 2, 2, 3, 3, 4, 4, -1, -1] \n",
    "given_labels = [mapping(labels, maps=given_maps) for labels in given_labels] \n",
    "sentence_probs = [merge_probs(pred_prob, maps=model_maps) for pred_prob in sentence_probs] "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb16f9cf",
   "metadata": {},
   "source": [
    "`merge_probs` above takes in two parameters: \n",
    "\n",
    "- `probs`: `np.array` of shape `(N, L)`, where `N` is the number of tokens in the sentence, and `L` is the number of classes potentially predicted by the model. \n",
    "- `maps`: `list` of length `L`, where `L` is the number of classes potentially predicted by the model, with details specified above in `model_maps`. \n",
    "\n",
    "and returns: \n",
    "\n",
    "- `probs_merged`: `np.array` of shape `(N, K)`, where `N` is the number of tokens in the sentence, and `K` is the number of classes corresponding to our considered set of entities (less than `L`). \n",
    "It ensures that `probs_merged[:, j] == \\sum_{maps[j']=j} probs[:, j']`. If any class in `maps` is negative (does not map to anything in the new set of classes), `probs_merged` is re-normalized such that each row sums to 1. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab0a3baf",
   "metadata": {},
   "source": [
    "## 5. Reducing from subword-level to word-level granularity \n",
    "\n",
    "When a sentence gets tokenized, each word-level token may be broken down into subword-level tokens, each of which has a predicted class probability vector generated by our model. Given that the dataset only provides labels for word-level tokens, we reduce subword-level tokens to word-level tokens. In this notebook, most subwords-to-words reductions are handled internally, but for models the reduction has to be done manually. Below we show our reduction code, which is slightly different from how the `bert` model [does the reduction](https://github.com/kamalkraj/BERT-NER/blob/dev/bert.py#L85). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "c43366dc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sentence:\tEu rejects German call to boycott British lamb.\n",
      "Given words:\t['Eu', 'rejects', 'German', 'call', 'to', 'boycott', 'British', 'lamb', '.']\n",
      "Subwords:\t['E', 'u', 'rejects', 'German', 'call', 'to', 'boycott', 'British', 'la', 'mb', '.']\n"
     ]
    }
   ],
   "source": [
    "print('Sentence:\\t' + sentences[0]) \n",
    "print('Given words:\\t' + str(given_words[0])) \n",
    "print('Subwords:\\t' + str(tokens)) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16eccd57",
   "metadata": {},
   "source": [
    "The word `lamb` is tokenized into two subwords `la` and `mb`. In this case, we average the predicted probabilities of the two subwords to serve as our predicted probability vector for the `lamb` token. Alternatively, we can take a weighted average, with the weight for each predicted probability proportional to the length of its corresponding subword (to emphasize longer subwords). We observed the benefits of this are insignificant for most datasets. \n",
    "\n",
    "Each tokenizer tokenizes sentences differently. In some rare cases, a subword may overlap two word-level tokens, resulting in a misalignment in tokenization. For example, consider the following tokenization: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "d0565e85",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Sentence:\tMassachusetts Institute of Technology (MIT)\n",
      "Given words:\t['Massachusetts', 'Institute', 'of', 'Technology', '(', 'MIT', ')']\n",
      "Subwords:\t['Massachusetts', 'Institute', 'of', 'Technology', '(M', 'IT', ')']\n"
     ]
    }
   ],
   "source": [
    "demo_sentence = 'Massachusetts Institute of Technology (MIT)' \n",
    "demo_given_words = ['Massachusetts', 'Institute', 'of', 'Technology', '(', 'MIT', ')'] \n",
    "demo_subwords = ['Massachusetts', 'Institute', 'of', 'Technology', '(M', 'IT', ')'] \n",
    "\n",
    "print('Sentence:\\t' + demo_sentence) \n",
    "print('Given words:\\t' + str(demo_given_words)) \n",
    "print('Subwords:\\t' + str(demo_subwords)) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a794f3d0",
   "metadata": {},
   "source": [
    "In this case, we assign the predicted probabilities of `(M` to `(`, and the average predicted probabilities of `(M` and `IT` to `MIT`. \n",
    "\n",
    "Let's use this approach to map the predicted probabilities for each subword-level token generated by the tokenizers to each word-level token in the original dataset. `get_pred_probs` below returns a nested list `pred_probs` where `pred_probs[i]` is a `np.ndarray` of shape `(N, K)`, such that `N` is the number of given tokens for sentence `i`, (recall `K` is the number of classes for our considered entities). This is the expected `pred_probs` format for methods in the `cleanlab.token_classification` module. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "f3656f29",
   "metadata": {},
   "outputs": [],
   "source": [
    "pred_probs = [get_pred_probs(sentence_probs[i], sentence_tokens[i], given_words[i]) \n",
    "                         for i in range(len(sentences))] "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a4c583d",
   "metadata": {},
   "source": [
    "Let's look at the tokens of the first sentence in the dataset, their given labels, and the corresponding predicted probabilities for each token. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "a2f2c113",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Token: Eu, given label: ORG\n",
      "Predicted probabilities: [3.041000e-04 2.383000e-04 9.993621e-01 7.010000e-05 2.550000e-05]\n",
      "Predicted label: PER\n",
      "\n",
      "Token: rejects, given label: O\n",
      "Predicted probabilities: [9.99988e-01 4.00000e-06 2.20000e-06 4.50000e-06 1.30000e-06]\n",
      "Predicted label: O\n",
      "\n",
      "Token: German, given label: MISC\n",
      "Predicted probabilities: [7.500000e-06 9.999611e-01 1.370000e-05 8.700000e-06 9.000000e-06]\n",
      "Predicted label: MISC\n",
      "\n",
      "Token: call, given label: O\n",
      "Predicted probabilities: [9.999894e-01 3.800000e-06 1.800000e-06 3.700000e-06 1.400000e-06]\n",
      "Predicted label: O\n",
      "\n",
      "Token: to, given label: O\n",
      "Predicted probabilities: [9.99991e-01 2.70000e-06 1.70000e-06 3.50000e-06 1.10000e-06]\n",
      "Predicted label: O\n",
      "\n",
      "Token: boycott, given label: O\n",
      "Predicted probabilities: [9.999877e-01 4.800000e-06 2.000000e-06 4.400000e-06 1.100000e-06]\n",
      "Predicted label: O\n",
      "\n",
      "Token: British, given label: MISC\n",
      "Predicted probabilities: [4.700000e-06 9.999639e-01 1.100000e-05 1.160000e-05 8.800000e-06]\n",
      "Predicted label: MISC\n",
      "\n",
      "Token: lamb, given label: O\n",
      "Predicted probabilities: [9.999867e-01 3.600000e-06 2.100000e-06 4.700000e-06 2.800000e-06]\n",
      "Predicted label: O\n",
      "\n",
      "Token: ., given label: O\n",
      "Predicted probabilities: [9.999908e-01 2.100000e-06 1.600000e-06 4.400000e-06 1.100000e-06]\n",
      "Predicted label: O\n",
      "\n"
     ]
    }
   ],
   "source": [
    "i = 0  # change this to view different examples \n",
    "entities = ['O', 'MISC', 'PER', 'ORG', 'LOC'] \n",
    "for word, label, prob in zip(given_words[i], given_labels[i], pred_probs[i]): \n",
    "    print('Token: %s, given label: %s' % (word, entities[label])) \n",
    "    print('Predicted probabilities: %s' % str(np.round(prob, 7))) \n",
    "    print('Predicted label: %s\\n' % entities[np.argmax(prob)]) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eeb30126",
   "metadata": {},
   "source": [
    "## 6. Save properly formatted `pred_probs` \n",
    "\n",
    "Now that we have properly formatted `tokens`, `labels`, and `pred_probs` for use with `cleanlab.token_classification`, lets save them to file. Below, `to_dict` converts `pred_probs` into a dictionary `d` where `d[str(i)]==pred_probs[i]`, which can be saved as a `.npz` file. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "3066803a",
   "metadata": {},
   "outputs": [],
   "source": [
    "pred_probs_dict = to_dict(pred_probs) \n",
    "np.savez('pred_probs.npz', **pred_probs_dict) "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2de2d51a",
   "metadata": {},
   "source": [
    "## 7. Model evaluation  \n",
    "\n",
    "For fun, let's evaluate the predictive accuracy of our trained model. We employ [the definition of precision/recall introduced by CoNLL-2003](https://www.davidsbatista.net/blog/2018/05/09/Named_Entity_Evaluation/): \n",
    "\n",
    "> Precision is the percentage of named entities found by the learning system that are correct. Recall is the percentage of named entities present in the corpus that are found by the system. A named entity is correct only if it is an exact match of the corresponding entity in the data file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "d285325a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Precision\t\t0.949\n",
      "Recall\t\t\t0.952\n",
      "f1-score\t\t0.951\n",
      "Accuracy\t\t0.989\n",
      "Balanced Accuracy\t0.955\n"
     ]
    }
   ],
   "source": [
    "predictions = [pred_prob.argmax(axis=1) for pred_prob in pred_probs] \n",
    "predictions_flatten = [pred for prediction in predictions for pred in prediction] \n",
    "given_labels_flatten = [label for given_label in given_labels for label in given_label] \n",
    "\n",
    "counts = [0, 0, 0, 0] \n",
    "correct = 0 \n",
    "\n",
    "for truth, prediction in zip(given_labels_flatten, predictions_flatten): \n",
    "    if truth != 0: \n",
    "        if truth == prediction: \n",
    "            counts[0] += 1 \n",
    "        counts[1] += 1 \n",
    "    if prediction != 0: \n",
    "        if truth == prediction: \n",
    "            counts[2] += 1 \n",
    "        counts[3] += 1 \n",
    "    if truth == prediction: \n",
    "        correct += 1 \n",
    "        \n",
    "precision = counts[2] / counts[3] \n",
    "recall = counts[0] / counts[1] \n",
    "f1 = 2 * precision * recall / (precision + recall) \n",
    "accuracy = correct / len(given_labels_flatten) \n",
    "\n",
    "balanced_accuracy = balanced_accuracy_score(given_labels_flatten, predictions_flatten) \n",
    "\n",
    "print('Precision\\t\\t%.3f\\nRecall\\t\\t\\t%.3f\\nf1-score\\t\\t%.3f\\nAccuracy\\t\\t%.3f\\nBalanced Accuracy\\t%.3f' % \n",
    "     (precision, recall, f1, accuracy, balanced_accuracy)) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "e0ab910c",
   "metadata": {
    "nbsphinx": "hidden"
   },
   "outputs": [],
   "source": [
    "# This cell is only for our internal CI and you can ignore it. \n",
    "expected_words = ['Eu', 'rejects', 'German', 'call', 'to', 'boycott', 'British', 'lamb', '.'] \n",
    "expected_labels = [3, 0, 1, 0, 0, 0, 1, 0, 0] \n",
    "if given_words[0] != expected_words or given_labels[0] != expected_labels: \n",
    "    raise Exception(\"Something wrong with reading file\") "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "00885e89789f58e60dbba52a405dc834aaf92411914fde0d391f9b48289a0610"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}