Finetune-LegalBert-NER.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "h-aw3Qx529pj",
    "outputId": "cd3e1f3f-b8f6-4285-df94-2382ef533e4a"
   },
   "outputs": [],
   "source": [
    "import datasets\n",
    "import numpy as np\n",
    "from transformers import BertTokenizerFast\n",
    "from transformers import DataCollatorForTokenClassification\n",
    "from transformers import AutoModelForTokenClassification\n",
    "\n",
    "from transformers import TrainingArguments, Trainer\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "TCgMrBFQ3O-m",
    "outputId": "5503f975-519e-4422-9e67-44cd1a791536"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['id', 'tokens', 'pos_tags', 'chunk_tags', 'ner_tags'],\n",
       "        num_rows: 14041\n",
       "    })\n",
       "    validation: Dataset({\n",
       "        features: ['id', 'tokens', 'pos_tags', 'chunk_tags', 'ner_tags'],\n",
       "        num_rows: 3250\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['id', 'tokens', 'pos_tags', 'chunk_tags', 'ner_tags'],\n",
       "        num_rows: 3453\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "conll2003 = datasets.load_dataset(\"conll2003\")\n",
    "conll2003"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "3vSxfkT73R7R",
    "outputId": "14993c36-77ed-4169-dc17-dd1856e2662d"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'train': (14041, 5), 'validation': (3250, 5), 'test': (3453, 5)}"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conll2003.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "CZ7gkGsR3VnE",
    "outputId": "5b8ba210-7312-4d2a-ca2b-f7bf7d3b439b"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'id': '0',\n",
       " 'tokens': ['EU',\n",
       "  'rejects',\n",
       "  'German',\n",
       "  'call',\n",
       "  'to',\n",
       "  'boycott',\n",
       "  'British',\n",
       "  'lamb',\n",
       "  '.'],\n",
       " 'pos_tags': [22, 42, 16, 21, 35, 37, 16, 21, 7],\n",
       " 'chunk_tags': [11, 21, 11, 12, 21, 22, 11, 12, 0],\n",
       " 'ner_tags': [3, 0, 7, 0, 0, 0, 7, 0, 0]}"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conll2003[\"train\"][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "88_au5l23YKJ",
    "outputId": "e82e6b7e-8c36-4410-aa65-5f62b690f0d7"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Sequence(feature=ClassLabel(names=['O', 'B-PER', 'I-PER', 'B-ORG', 'I-ORG', 'B-LOC', 'I-LOC', 'B-MISC', 'I-MISC'], id=None), length=-1, id=None)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conll2003[\"train\"].features[\"ner_tags\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 36
    },
    "id": "Fn8asff53dQC",
    "outputId": "51afadf9-264d-4a0e-8223-9f67043c4651"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'The shared task of CoNLL-2003 concerns language-independent named entity recognition. We will concentrate on\\nfour types of named entities: persons, locations, organizations and names of miscellaneous entities that do\\nnot belong to the previous three groups.\\n\\nThe CoNLL-2003 shared task data files contain four columns separated by a single space. Each word has been put on\\na separate line and there is an empty line after each sentence. The first item on each line is a word, the second\\na part-of-speech (POS) tag, the third a syntactic chunk tag and the fourth the named entity tag. The chunk tags\\nand the named entity tags have the format I-TYPE which means that the word is inside a phrase of type TYPE. Only\\nif two phrases of the same type immediately follow each other, the first word of the second phrase will have tag\\nB-TYPE to show that it starts a new phrase. A word with tag O is not part of a phrase. Note the dataset uses IOB2\\ntagging scheme, whereas the original dataset uses IOB1.\\n\\nFor more details see https://www.clips.uantwerpen.be/conll2003/ner/ and https://www.aclweb.org/anthology/W03-0419\\n'"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conll2003['train'].description"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "id": "s-szOW9w3w6A"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Users\\Admin\\anaconda3\\envs\\pytorch_env\\Lib\\site-packages\\transformers\\tokenization_utils_base.py:1601: FutureWarning: `clean_up_tokenization_spaces` was not set. It will be set to `True` by default. This behavior will be depracted in transformers v4.45, and will be then set to `False` by default. For more details check this issue: https://github.com/huggingface/transformers/issues/31884\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "tokenizer = BertTokenizerFast.from_pretrained(\"nlpaueb/legal-bert-base-uncased\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "8Jk3cBaI4BjY",
    "outputId": "f7e39fda-a1ba-4ac7-ec7b-340d366e5cec"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'id': '0',\n",
       " 'tokens': ['EU',\n",
       "  'rejects',\n",
       "  'German',\n",
       "  'call',\n",
       "  'to',\n",
       "  'boycott',\n",
       "  'British',\n",
       "  'lamb',\n",
       "  '.'],\n",
       " 'pos_tags': [22, 42, 16, 21, 35, 37, 16, 21, 7],\n",
       " 'chunk_tags': [11, 21, 11, 12, 21, 22, 11, 12, 0],\n",
       " 'ner_tags': [3, 0, 7, 0, 0, 0, 7, 0, 0]}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conll2003['train'][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 53
    },
    "id": "dE2Kq36i4Gyr",
    "outputId": "b204b01e-2b29-464c-cd3b-2c4f1dc54185"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[None, 0, 1, 2, 3, 4, 5, 6, 6, 7, 8, None]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "' As we can see, it returns a list with the same number of elements as our processed input ids, mapping special tokens to None and all other tokens to their respective word. This way, we can align the labels with the processed input ids. '"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "example_text = conll2003['train'][0]\n",
    "\n",
    "tokenized_input = tokenizer(example_text[\"tokens\"], is_split_into_words=True)\n",
    "\n",
    "tokens = tokenizer.convert_ids_to_tokens(tokenized_input[\"input_ids\"])\n",
    "\n",
    "word_ids = tokenized_input.word_ids()\n",
    "\n",
    "print(word_ids)\n",
    "\n",
    "''' As we can see, it returns a list with the same number of elements as our processed input ids, mapping special tokens to None and all other tokens to their respective word. This way, we can align the labels with the processed input ids. '''\n",
    "\n",
    "# tokenized_input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "rIgXRNog4RPa",
    "outputId": "f01ac480-2381-4e6e-b3cb-8fd7753f11e2"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['[CLS]',\n",
       " 'eu',\n",
       " 'rejects',\n",
       " 'german',\n",
       " 'call',\n",
       " 'to',\n",
       " 'boycott',\n",
       " 'britis',\n",
       " '##h',\n",
       " 'lamb',\n",
       " '.',\n",
       " '[SEP]']"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokens = tokenizer.convert_ids_to_tokens(tokenized_input[\"input_ids\"])\n",
    "tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "J3o876hP4TuS",
    "outputId": "d15fd7e3-5843-4a15-d5a2-756a8a0ea1a8"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(9, 12)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(example_text['ner_tags']), len(tokenized_input[\"input_ids\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "id": "cAWtqrFW4Xnz"
   },
   "outputs": [],
   "source": [
    "def tokenize_and_align_labels(examples, label_all_tokens=True):\n",
    "    tokenized_inputs = tokenizer(examples[\"tokens\"], truncation=True, is_split_into_words=True)\n",
    "    labels = []\n",
    "    for i, label in enumerate(examples[\"ner_tags\"]):\n",
    "        word_ids = tokenized_inputs.word_ids(batch_index=i)\n",
    "        # word_ids() => Return a list mapping the tokens\n",
    "        # to their actual word in the initial sentence.\n",
    "        # It Returns a list indicating the word corresponding to each token.\n",
    "        previous_word_idx = None\n",
    "        label_ids = []\n",
    "        # Special tokens like `` and `<\\s>` are originally mapped to None\n",
    "        # We need to set the label to -100 so they are automatically ignored in the loss function.\n",
    "        for word_idx in word_ids:\n",
    "            if word_idx is None:\n",
    "                # set –100 as the label for these special tokens\n",
    "                label_ids.append(-100)\n",
    "            # For the other tokens in a word, we set the label to either the current label or -100, depending on\n",
    "            # the label_all_tokens flag.\n",
    "            elif word_idx != previous_word_idx:\n",
    "                # if current word_idx is != prev then its the most regular case\n",
    "                # and add the corresponding token\n",
    "                label_ids.append(label[word_idx])\n",
    "            else:\n",
    "                # to take care of sub-words which have the same word_idx\n",
    "                # set -100 as well for them, but only if label_all_tokens == False\n",
    "                label_ids.append(label[word_idx] if label_all_tokens else -100)\n",
    "                # mask the subword representations after the first subword\n",
    "\n",
    "            previous_word_idx = word_idx\n",
    "        labels.append(label_ids)\n",
    "    tokenized_inputs[\"labels\"] = labels\n",
    "    return tokenized_inputs\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Kx533oFG4dez",
    "outputId": "771cf0f2-62b6-4e0b-99b9-9352a4ce65b8"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'id': ['4'],\n",
       " 'tokens': [['Germany',\n",
       "   \"'s\",\n",
       "   'representative',\n",
       "   'to',\n",
       "   'the',\n",
       "   'European',\n",
       "   'Union',\n",
       "   \"'s\",\n",
       "   'veterinary',\n",
       "   'committee',\n",
       "   'Werner',\n",
       "   'Zwingmann',\n",
       "   'said',\n",
       "   'on',\n",
       "   'Wednesday',\n",
       "   'consumers',\n",
       "   'should',\n",
       "   'buy',\n",
       "   'sheepmeat',\n",
       "   'from',\n",
       "   'countries',\n",
       "   'other',\n",
       "   'than',\n",
       "   'Britain',\n",
       "   'until',\n",
       "   'the',\n",
       "   'scientific',\n",
       "   'advice',\n",
       "   'was',\n",
       "   'clearer',\n",
       "   '.']],\n",
       " 'pos_tags': [[22,\n",
       "   27,\n",
       "   21,\n",
       "   35,\n",
       "   12,\n",
       "   22,\n",
       "   22,\n",
       "   27,\n",
       "   16,\n",
       "   21,\n",
       "   22,\n",
       "   22,\n",
       "   38,\n",
       "   15,\n",
       "   22,\n",
       "   24,\n",
       "   20,\n",
       "   37,\n",
       "   21,\n",
       "   15,\n",
       "   24,\n",
       "   16,\n",
       "   15,\n",
       "   22,\n",
       "   15,\n",
       "   12,\n",
       "   16,\n",
       "   21,\n",
       "   38,\n",
       "   17,\n",
       "   7]],\n",
       " 'chunk_tags': [[11,\n",
       "   11,\n",
       "   12,\n",
       "   13,\n",
       "   11,\n",
       "   12,\n",
       "   12,\n",
       "   11,\n",
       "   12,\n",
       "   12,\n",
       "   12,\n",
       "   12,\n",
       "   21,\n",
       "   13,\n",
       "   11,\n",
       "   12,\n",
       "   21,\n",
       "   22,\n",
       "   11,\n",
       "   13,\n",
       "   11,\n",
       "   1,\n",
       "   13,\n",
       "   11,\n",
       "   17,\n",
       "   11,\n",
       "   12,\n",
       "   12,\n",
       "   21,\n",
       "   1,\n",
       "   0]],\n",
       " 'ner_tags': [[5,\n",
       "   0,\n",
       "   0,\n",
       "   0,\n",
       "   0,\n",
       "   3,\n",
       "   4,\n",
       "   0,\n",
       "   0,\n",
       "   0,\n",
       "   1,\n",
       "   2,\n",
       "   0,\n",
       "   0,\n",
       "   0,\n",
       "   0,\n",
       "   0,\n",
       "   0,\n",
       "   0,\n",
       "   0,\n",
       "   0,\n",
       "   0,\n",
       "   0,\n",
       "   5,\n",
       "   0,\n",
       "   0,\n",
       "   0,\n",
       "   0,\n",
       "   0,\n",
       "   0,\n",
       "   0]]}"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "conll2003['train'][4:5]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "0f6BfSwL4879",
    "outputId": "d369517d-1b3b-4f0f-a036-e96916f134d5"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'input_ids': [[101, 1221, 110, 163, 900, 211, 207, 274, 403, 110, 163, 2824, 195, 526, 532, 188, 3298, 13898, 235, 4149, 786, 222, 15034, 2765, 305, 2778, 4899, 3681, 238, 779, 231, 268, 6852, 890, 598, 207, 2137, 2580, 246, 969, 577, 117, 102]], 'token_type_ids': [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], 'attention_mask': [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]], 'labels': [[-100, 5, 0, 0, 0, 0, 0, 3, 4, 0, 0, 0, 0, 0, 1, 1, 1, 2, 2, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 5, 5, 0, 0, 0, 0, 0, 0, 0, 0, -100]]}\n"
     ]
    }
   ],
   "source": [
    "q = tokenize_and_align_labels(conll2003['train'][4:5])\n",
    "print(q)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "W2GYlktZ5BIO",
    "outputId": "02d58dd5-499d-41c0-a2fa-84a8d467bd7d"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[CLS]___________________________________ -100\n",
      "germany_________________________________ 5\n",
      "'_______________________________________ 0\n",
      "s_______________________________________ 0\n",
      "representative__________________________ 0\n",
      "to______________________________________ 0\n",
      "the_____________________________________ 0\n",
      "european________________________________ 3\n",
      "union___________________________________ 4\n",
      "'_______________________________________ 0\n",
      "s_______________________________________ 0\n",
      "veterinar_______________________________ 0\n",
      "##y_____________________________________ 0\n",
      "committee_______________________________ 0\n",
      "we______________________________________ 1\n",
      "##r_____________________________________ 1\n",
      "##ner___________________________________ 1\n",
      "zw______________________________________ 2\n",
      "##ing___________________________________ 2\n",
      "##mann__________________________________ 2\n",
      "said____________________________________ 0\n",
      "on______________________________________ 0\n",
      "wednesday_______________________________ 0\n",
      "consumers_______________________________ 0\n",
      "should__________________________________ 0\n",
      "buy_____________________________________ 0\n",
      "sheep___________________________________ 0\n",
      "##meat__________________________________ 0\n",
      "from____________________________________ 0\n",
      "countries_______________________________ 0\n",
      "other___________________________________ 0\n",
      "than____________________________________ 0\n",
      "brita___________________________________ 5\n",
      "##in____________________________________ 5\n",
      "until___________________________________ 0\n",
      "the_____________________________________ 0\n",
      "scientific______________________________ 0\n",
      "advice__________________________________ 0\n",
      "was_____________________________________ 0\n",
      "clear___________________________________ 0\n",
      "##er____________________________________ 0\n",
      "._______________________________________ 0\n",
      "[SEP]___________________________________ -100\n"
     ]
    }
   ],
   "source": [
    "for token, label in zip(tokenizer.convert_ids_to_tokens(q[\"input_ids\"][0]),q[\"labels\"][0]):\n",
    "    print(f\"{token:_<40} {label}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "id": "n0-6HLGk5E2c"
   },
   "outputs": [],
   "source": [
    "## Applying on entire data\n",
    "tokenized_datasets = conll2003.map(tokenize_and_align_labels, batched=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "DtkCf13G5LTa",
    "outputId": "9a24853a-445c-417e-c1c4-aa1d087a3f1a"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'id': '0',\n",
       " 'tokens': ['EU',\n",
       "  'rejects',\n",
       "  'German',\n",
       "  'call',\n",
       "  'to',\n",
       "  'boycott',\n",
       "  'British',\n",
       "  'lamb',\n",
       "  '.'],\n",
       " 'pos_tags': [22, 42, 16, 21, 35, 37, 16, 21, 7],\n",
       " 'chunk_tags': [11, 21, 11, 12, 21, 22, 11, 12, 0],\n",
       " 'ner_tags': [3, 0, 7, 0, 0, 0, 7, 0, 0],\n",
       " 'input_ids': [101,\n",
       "  501,\n",
       "  5714,\n",
       "  1600,\n",
       "  1842,\n",
       "  211,\n",
       "  21215,\n",
       "  3585,\n",
       "  178,\n",
       "  7846,\n",
       "  117,\n",
       "  102],\n",
       " 'token_type_ids': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],\n",
       " 'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1],\n",
       " 'labels': [-100, 3, 0, 7, 0, 0, 0, 7, 7, 0, 0, -100]}"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenized_datasets['train'][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "JXXdGYng5QGQ",
    "outputId": "a5838416-3b20-4222-9eeb-07df173b7291"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of BertForTokenClassification were not initialized from the model checkpoint at nlpaueb/legal-bert-base-uncased and are newly initialized: ['classifier.bias', 'classifier.weight']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    }
   ],
   "source": [
    "model = AutoModelForTokenClassification.from_pretrained(\"nlpaueb/legal-bert-base-uncased\", num_labels=9)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "id": "ThGE02OU5bx6"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Users\\Admin\\anaconda3\\envs\\pytorch_env\\Lib\\site-packages\\transformers\\training_args.py:1525: FutureWarning: `evaluation_strategy` is deprecated and will be removed in version 4.46 of 🤗 Transformers. Use `eval_strategy` instead\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "args = TrainingArguments(\n",
    "\"test-ner\",\n",
    "evaluation_strategy = \"epoch\",\n",
    "learning_rate=2e-5,\n",
    "per_device_train_batch_size=16,\n",
    "per_device_eval_batch_size=16,\n",
    "num_train_epochs=3,\n",
    "weight_decay=0.01,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "id": "tatPx4f3716w"
   },
   "outputs": [],
   "source": [
    "data_collator = DataCollatorForTokenClassification(tokenizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "r_waZ5D677VW",
    "outputId": "2989f2b2-a5b3-44cd-8597-cd97b33f68ce"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\Admin\\AppData\\Local\\Temp\\ipykernel_11916\\4212495798.py:1: FutureWarning: load_metric is deprecated and will be removed in the next major version of datasets. Use 'evaluate.load' instead, from the new library 🤗 Evaluate: https://huggingface.co/docs/evaluate\n",
      "  metric = datasets.load_metric(\"seqeval\",trust_remote_code=True)\n"
     ]
    }
   ],
   "source": [
    "metric = datasets.load_metric(\"seqeval\",trust_remote_code=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "id": "YUBlwhCW8ANN"
   },
   "outputs": [],
   "source": [
    "example = conll2003['train'][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "3zpUOQz_8Fex",
    "outputId": "aa7b900c-6452-44c6-b6cf-c428d86fe77b"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['O', 'B-PER', 'I-PER', 'B-ORG', 'I-ORG', 'B-LOC', 'I-LOC', 'B-MISC', 'I-MISC']"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "label_list = conll2003[\"train\"].features[\"ner_tags\"].feature.names\n",
    "\n",
    "label_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "1-Uo0k7v8J83",
    "outputId": "64951f59-c5b6-4036-e1fd-2211f716f9f4"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3\n",
      "0\n",
      "7\n",
      "0\n",
      "0\n",
      "0\n",
      "7\n",
      "0\n",
      "0\n"
     ]
    }
   ],
   "source": [
    "for i in example[\"ner_tags\"]:\n",
    "  print(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "WXM4zusz8PlZ",
    "outputId": "1460456a-3e36-4fe3-c76f-9134bcb424b4"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['B-ORG', 'O', 'B-MISC', 'O', 'O', 'O', 'B-MISC', 'O', 'O']"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "labels = [label_list[i] for i in example[\"ner_tags\"]]\n",
    "labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "OUsKPEVw8UDw",
    "outputId": "39ffefaa-4e4c-41a0-b117-5ab0f52987fc"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'MISC': {'precision': 1.0, 'recall': 1.0, 'f1': 1.0, 'number': 2},\n",
       " 'ORG': {'precision': 1.0, 'recall': 1.0, 'f1': 1.0, 'number': 1},\n",
       " 'overall_precision': 1.0,\n",
       " 'overall_recall': 1.0,\n",
       " 'overall_f1': 1.0,\n",
       " 'overall_accuracy': 1.0}"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "metric.compute(predictions=[labels], references=[labels])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "id": "3mk-nAO58Yc8"
   },
   "outputs": [],
   "source": [
    "def compute_metrics(eval_preds):\n",
    "    pred_logits, labels = eval_preds\n",
    "\n",
    "    pred_logits = np.argmax(pred_logits, axis=2)\n",
    "    # the logits and the probabilities are in the same order,\n",
    "    # so we don’t need to apply the softmax\n",
    "\n",
    "    # We remove all the values where the label is -100\n",
    "    predictions = [\n",
    "        [label_list[eval_preds] for (eval_preds, l) in zip(prediction, label) if l != -100]\n",
    "        for prediction, label in zip(pred_logits, labels)\n",
    "    ]\n",
    "\n",
    "    true_labels = [\n",
    "      [label_list[l] for (eval_preds, l) in zip(prediction, label) if l != -100]\n",
    "       for prediction, label in zip(pred_logits, labels)\n",
    "   ]\n",
    "    results = metric.compute(predictions=predictions, references=true_labels)\n",
    "\n",
    "    return {\n",
    "          \"precision\": results[\"overall_precision\"],\n",
    "          \"recall\": results[\"overall_recall\"],\n",
    "          \"f1\": results[\"overall_f1\"],\n",
    "          \"accuracy\": results[\"overall_accuracy\"],\n",
    "  }\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8Z9Q7LIr8eBz"
   },
   "source": [
    "# Training:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "id": "HHF7tCSP8dOr"
   },
   "outputs": [],
   "source": [
    "trainer = Trainer(\n",
    "   model,\n",
    "   args,\n",
    "   train_dataset=tokenized_datasets[\"train\"],\n",
    "   eval_dataset=tokenized_datasets[\"validation\"],\n",
    "   data_collator=data_collator,\n",
    "   tokenizer=tokenizer,\n",
    "   compute_metrics=compute_metrics\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 75
    },
    "id": "gWP9GVil8m6D",
    "outputId": "bade1a45-2b34-4efa-b1c4-1cd99b7aac25"
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "efa8b49e546b430994e7a839d8b1c9f6",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/2634 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Users\\Admin\\anaconda3\\envs\\pytorch_env\\Lib\\site-packages\\transformers\\models\\bert\\modeling_bert.py:439: UserWarning: 1Torch was not compiled with flash attention. (Triggered internally at C:\\cb\\pytorch_1000000000000\\work\\aten\\src\\ATen\\native\\transformers\\cuda\\sdp_utils.cpp:555.)\n",
      "  attn_output = torch.nn.functional.scaled_dot_product_attention(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 0.3091, 'grad_norm': 1.8334522247314453, 'learning_rate': 1.6203492786636296e-05, 'epoch': 0.57}\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "dc0caf6ccccd4443b973796bc8e45042",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/204 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.10415639728307724, 'eval_precision': 0.8785548865215378, 'eval_recall': 0.8790434702011308, 'eval_f1': 0.878799110452187, 'eval_accuracy': 0.969333097886868, 'eval_runtime': 2.8434, 'eval_samples_per_second': 1143.012, 'eval_steps_per_second': 71.746, 'epoch': 1.0}\n",
      "{'loss': 0.127, 'grad_norm': 7.049642086029053, 'learning_rate': 1.240698557327259e-05, 'epoch': 1.14}\n",
      "{'loss': 0.0889, 'grad_norm': 2.821964979171753, 'learning_rate': 8.610478359908885e-06, 'epoch': 1.71}\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "341809c0b54b4a5d9aae9f894997cfc4",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/204 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.0890548974275589, 'eval_precision': 0.9023015066087439, 'eval_recall': 0.9048104550931504, 'eval_f1': 0.9035542391706776, 'eval_accuracy': 0.9752339749249515, 'eval_runtime': 2.922, 'eval_samples_per_second': 1112.269, 'eval_steps_per_second': 69.816, 'epoch': 2.0}\n",
      "{'loss': 0.0684, 'grad_norm': 6.190309524536133, 'learning_rate': 4.8139711465451785e-06, 'epoch': 2.28}\n",
      "{'loss': 0.0549, 'grad_norm': 2.7784011363983154, 'learning_rate': 1.0174639331814731e-06, 'epoch': 2.85}\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "11f910b5de49434d99faccd1cae46cb5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/204 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 0.08921829611063004, 'eval_precision': 0.9010163904404359, 'eval_recall': 0.9120400407822783, 'eval_f1': 0.9064947029018886, 'eval_accuracy': 0.9754547059862264, 'eval_runtime': 3.1305, 'eval_samples_per_second': 1038.18, 'eval_steps_per_second': 65.166, 'epoch': 3.0}\n",
      "{'train_runtime': 163.9536, 'train_samples_per_second': 256.92, 'train_steps_per_second': 16.066, 'train_loss': 0.12552894981140972, 'epoch': 3.0}\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=2634, training_loss=0.12552894981140972, metrics={'train_runtime': 163.9536, 'train_samples_per_second': 256.92, 'train_steps_per_second': 16.066, 'total_flos': 1101309468061158.0, 'train_loss': 0.12552894981140972, 'epoch': 3.0})"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "id": "4TvBaSGNR2Bf"
   },
   "outputs": [],
   "source": [
    "model.save_pretrained(\"ner_model\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('tokenizer\\\\tokenizer_config.json',\n",
       " 'tokenizer\\\\special_tokens_map.json',\n",
       " 'tokenizer\\\\vocab.txt',\n",
       " 'tokenizer\\\\added_tokens.json',\n",
       " 'tokenizer\\\\tokenizer.json')"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tokenizer.save_pretrained(\"tokenizer\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "id2label = {\n",
    "    str(i): label for i,label in enumerate(label_list)\n",
    "}\n",
    "label2id = {\n",
    "    label: str(i) for i,label in enumerate(label_list)\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'0': 'O',\n",
       " '1': 'B-PER',\n",
       " '2': 'I-PER',\n",
       " '3': 'B-ORG',\n",
       " '4': 'I-ORG',\n",
       " '5': 'B-LOC',\n",
       " '6': 'I-LOC',\n",
       " '7': 'B-MISC',\n",
       " '8': 'I-MISC'}"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "id2label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'O': '0',\n",
       " 'B-PER': '1',\n",
       " 'I-PER': '2',\n",
       " 'B-ORG': '3',\n",
       " 'I-ORG': '4',\n",
       " 'B-LOC': '5',\n",
       " 'I-LOC': '6',\n",
       " 'B-MISC': '7',\n",
       " 'I-MISC': '8'}"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "label2id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "config = json.load(open(\"ner_model/config.json\"))\n",
    "config[\"id2label\"] = id2label\n",
    "config[\"label2id\"] = label2id\n",
    "json.dump(config, open(\"ner_model/config.json\",\"w\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_fine_tuned = AutoModelForTokenClassification.from_pretrained(\"ner_model\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Hardware accelerator e.g. GPU is available in the environment, but no `device` argument is passed to the `Pipeline` object. Model will be on CPU.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{'entity': 'B-PER', 'score': 0.9974859, 'index': 1, 'word': 'pris', 'start': 0, 'end': 4}, {'entity': 'B-PER', 'score': 0.9980908, 'index': 2, 'word': '##cilla', 'start': 4, 'end': 9}, {'entity': 'I-PER', 'score': 0.9983791, 'index': 3, 'word': 'li', 'start': 10, 'end': 12}, {'entity': 'I-PER', 'score': 0.9983045, 'index': 4, 'word': '##cu', 'start': 12, 'end': 14}, {'entity': 'I-PER', 'score': 0.99832314, 'index': 5, 'word': '##p', 'start': 14, 'end': 15}, {'entity': 'B-ORG', 'score': 0.9954332, 'index': 10, 'word': 'dun', 'start': 36, 'end': 39}, {'entity': 'B-ORG', 'score': 0.9952342, 'index': 11, 'word': '##kin', 'start': 39, 'end': 42}, {'entity': 'I-ORG', 'score': 0.9958896, 'index': 12, 'word': 'don', 'start': 43, 'end': 46}, {'entity': 'I-ORG', 'score': 0.9959781, 'index': 13, 'word': '##ut', 'start': 46, 'end': 48}, {'entity': 'I-ORG', 'score': 0.9942637, 'index': 14, 'word': '##s', 'start': 48, 'end': 49}]\n"
     ]
    }
   ],
   "source": [
    "nlp = pipeline(\"ner\", model=model_fine_tuned, tokenizer=tokenizer)\n",
    "\n",
    "\n",
    "example = \"Priscilla Licup is the President of Dunkin Donuts\"\n",
    "\n",
    "ner_results = nlp(example)\n",
    "\n",
    "print(ner_results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "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.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}