From 5758a7eb1c156b2c513930af666acb2d4a2cb3c3 Mon Sep 17 00:00:00 2001
From: JINO ROHIT <find.jinorohit@gmail.com>
Date: Thu, 10 Oct 2024 14:34:16 +0530
Subject: [PATCH] ENH LoRA notebook for NER task (#2126)

---
 .../token_classification/peft_lora_ner.ipynb  | 780 ++++++++++++++++++
 1 file changed, 780 insertions(+)
 create mode 100644 examples/token_classification/peft_lora_ner.ipynb

diff --git a/examples/token_classification/peft_lora_ner.ipynb b/examples/token_classification/peft_lora_ner.ipynb
new file mode 100644
index 0000000000..fae9b94b4e
--- /dev/null
+++ b/examples/token_classification/peft_lora_ner.ipynb
@@ -0,0 +1,780 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Named Entity Recognition with Peft Model 🤗\n",
+    "\n",
+    "##### In this notebook, we will learn how to perform Named Entity Recognition(NER) on the CoNLL-2003 dataset using the Trainer class\n",
+    "\n",
+    "##### This notebook has been adapted from the main NLP course here - https://huggingface.co/learn/nlp-course/chapter7/2?fw=pt#fine-tuning-the-model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "#install the required libraries\n",
+    "!pip install -q datasets evaluate transformers seqeval"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Import required libraries\n",
+    "from datasets import load_dataset\n",
+    "from transformers import AutoTokenizer, AutoModelForTokenClassification, DataCollatorForTokenClassification, TrainingArguments, Trainer, pipeline\n",
+    "from peft import get_peft_model, LoraConfig, TaskType\n",
+    "import evaluate\n",
+    "import numpy as np\n",
+    "from huggingface_hub import notebook_login"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "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",
+      "})\n"
+     ]
+    }
+   ],
+   "source": [
+    "raw_datasets = load_dataset(\"conll2003\")\n",
+    "print(raw_datasets)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "['EU', 'rejects', 'German', 'call', 'to', 'boycott', 'British', 'lamb', '.']"
+      ]
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# Look at the tokens of the first training example\n",
+    "raw_datasets[\"train\"][0][\"tokens\"]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[3, 0, 7, 0, 0, 0, 7, 0, 0]"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# Look at the NER tags of the first training example\n",
+    "raw_datasets[\"train\"][0][\"ner_tags\"]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "['O', 'B-PER', 'I-PER', 'B-ORG', 'I-ORG', 'B-LOC', 'I-LOC', 'B-MISC', 'I-MISC']"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# Get the label names for the NER tags\n",
+    "ner_feature = raw_datasets[\"train\"].features[\"ner_tags\"]\n",
+    "label_names = ner_feature.feature.names\n",
+    "label_names"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "EU    rejects German call to boycott British lamb . \n",
+      "B-ORG O       B-MISC O    O  O       B-MISC  O    O \n"
+     ]
+    }
+   ],
+   "source": [
+    "words = raw_datasets[\"train\"][0][\"tokens\"]\n",
+    "labels = raw_datasets[\"train\"][0][\"ner_tags\"]\n",
+    "line1 = \"\"\n",
+    "line2 = \"\"\n",
+    "for word, label in zip(words, labels):\n",
+    "    full_label = label_names[label]\n",
+    "    max_length = max(len(word), len(full_label))\n",
+    "    line1 += word + \" \" * (max_length - len(word) + 1)\n",
+    "    line2 += full_label + \" \" * (max_length - len(full_label) + 1)\n",
+    "\n",
+    "print(line1)\n",
+    "print(line2)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "e:\\open_source\\peft-folder\\ner-examples\\.venv\\Lib\\site-packages\\transformers\\tokenization_utils_base.py:1617: FutureWarning: `clean_up_tokenization_spaces` was not set. It will be set to `True` by default. This behavior will be deprecated 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": [
+    "# Load the tokenizer\n",
+    "model_checkpoint = \"bert-base-cased\"\n",
+    "tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "['[CLS]',\n",
+       " 'EU',\n",
+       " 'rejects',\n",
+       " 'German',\n",
+       " 'call',\n",
+       " 'to',\n",
+       " 'boycott',\n",
+       " 'British',\n",
+       " 'la',\n",
+       " '##mb',\n",
+       " '.',\n",
+       " '[SEP]']"
+      ]
+     },
+     "execution_count": 8,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# Tokenize the first training example\n",
+    "inputs = tokenizer(raw_datasets[\"train\"][0][\"tokens\"], is_split_into_words=True)\n",
+    "inputs.tokens()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def align_labels_with_tokens(labels, word_ids):\n",
+    "    new_labels = []\n",
+    "    current_word = None\n",
+    "    for word_id in word_ids:\n",
+    "        if word_id != current_word:\n",
+    "            # Start of a new word!\n",
+    "            current_word = word_id\n",
+    "            label = -100 if word_id is None else labels[word_id]\n",
+    "            new_labels.append(label)\n",
+    "        elif word_id is None:\n",
+    "            # Special token\n",
+    "            new_labels.append(-100)\n",
+    "        else:\n",
+    "            # Same word as previous token\n",
+    "            label = labels[word_id]\n",
+    "            # If the label is B-XXX we change it to I-XXX\n",
+    "            if label % 2 == 1:\n",
+    "                label += 1\n",
+    "            new_labels.append(label)\n",
+    "\n",
+    "    return new_labels"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[3, 0, 7, 0, 0, 0, 7, 0, 0]\n",
+      "[-100, 3, 0, 7, 0, 0, 0, 7, 0, 0, 0, -100]\n"
+     ]
+    }
+   ],
+   "source": [
+    "labels = raw_datasets[\"train\"][0][\"ner_tags\"]\n",
+    "word_ids = inputs.word_ids()\n",
+    "print(labels)\n",
+    "print(align_labels_with_tokens(labels, word_ids))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def tokenize_and_align_labels(examples):\n",
+    "    tokenized_inputs = tokenizer(\n",
+    "        examples[\"tokens\"], truncation=True, is_split_into_words=True\n",
+    "    )\n",
+    "    all_labels = examples[\"ner_tags\"]\n",
+    "    new_labels = []\n",
+    "    for i, labels in enumerate(all_labels):\n",
+    "        word_ids = tokenized_inputs.word_ids(i)\n",
+    "        new_labels.append(align_labels_with_tokens(labels, word_ids))\n",
+    "\n",
+    "    tokenized_inputs[\"labels\"] = new_labels\n",
+    "    return tokenized_inputs"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "tokenized_datasets = raw_datasets.map(\n",
+    "    tokenize_and_align_labels,\n",
+    "    batched=True,\n",
+    "    remove_columns=raw_datasets[\"train\"].column_names,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "data_collator = DataCollatorForTokenClassification(tokenizer=tokenizer)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[-100, 3, 0, 7, 0, 0, 0, 7, 0, 0, 0, -100]\n",
+      "[-100, 1, 2, -100]\n"
+     ]
+    }
+   ],
+   "source": [
+    "for i in range(2):\n",
+    "    print(tokenized_datasets[\"train\"][i][\"labels\"])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "metric = evaluate.load(\"seqeval\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Create label mappings\n",
+    "id2label = {i: label for i, label in enumerate(label_names)}\n",
+    "label2id = {v: k for k, v in id2label.items()}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Some weights of BertForTokenClassification were not initialized from the model checkpoint at bert-base-cased 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": [
+    "# Load the pre-trained model\n",
+    "model = AutoModelForTokenClassification.from_pretrained(\n",
+    "    model_checkpoint,\n",
+    "    id2label=id2label,\n",
+    "    label2id=label2id,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "9"
+      ]
+     },
+     "execution_count": 18,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "model.config.num_labels"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "BertForTokenClassification(\n",
+       "  (bert): BertModel(\n",
+       "    (embeddings): BertEmbeddings(\n",
+       "      (word_embeddings): Embedding(28996, 768, padding_idx=0)\n",
+       "      (position_embeddings): Embedding(512, 768)\n",
+       "      (token_type_embeddings): Embedding(2, 768)\n",
+       "      (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
+       "      (dropout): Dropout(p=0.1, inplace=False)\n",
+       "    )\n",
+       "    (encoder): BertEncoder(\n",
+       "      (layer): ModuleList(\n",
+       "        (0-11): 12 x BertLayer(\n",
+       "          (attention): BertAttention(\n",
+       "            (self): BertSdpaSelfAttention(\n",
+       "              (query): Linear(in_features=768, out_features=768, bias=True)\n",
+       "              (key): Linear(in_features=768, out_features=768, bias=True)\n",
+       "              (value): Linear(in_features=768, out_features=768, bias=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "            (output): BertSelfOutput(\n",
+       "              (dense): Linear(in_features=768, out_features=768, bias=True)\n",
+       "              (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
+       "              (dropout): Dropout(p=0.1, inplace=False)\n",
+       "            )\n",
+       "          )\n",
+       "          (intermediate): BertIntermediate(\n",
+       "            (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
+       "            (intermediate_act_fn): GELUActivation()\n",
+       "          )\n",
+       "          (output): BertOutput(\n",
+       "            (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
+       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
+       "            (dropout): Dropout(p=0.1, inplace=False)\n",
+       "          )\n",
+       "        )\n",
+       "      )\n",
+       "    )\n",
+       "  )\n",
+       "  (dropout): Dropout(p=0.1, inplace=False)\n",
+       "  (classifier): Linear(in_features=768, out_features=9, bias=True)\n",
+       ")"
+      ]
+     },
+     "execution_count": 19,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "trainable params: 301,833 || all params: 108,028,434 || trainable%: 0.2794\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Configure LoRA (Low-Rank Adaptation) for fine-tuning\n",
+    "peft_config = LoraConfig(target_modules = [\"query\", \"key\"], task_type = TaskType.TOKEN_CLS)\n",
+    "\n",
+    "model = get_peft_model(model, peft_config)\n",
+    "model.print_trainable_parameters()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def compute_metrics(eval_preds):\n",
+    "    logits, labels = eval_preds\n",
+    "    predictions = np.argmax(logits, axis=-1)\n",
+    "\n",
+    "    # Remove ignored index (special tokens) and convert to labels\n",
+    "    true_labels = [[label_names[l] for l in label if l != -100] for label in labels]\n",
+    "    true_predictions = [\n",
+    "        [label_names[p] for (p, l) in zip(prediction, label) if l != -100]\n",
+    "        for prediction, label in zip(predictions, labels)\n",
+    "    ]\n",
+    "    all_metrics = metric.compute(predictions=true_predictions, references=true_labels)\n",
+    "    return {\n",
+    "        \"precision\": all_metrics[\"overall_precision\"],\n",
+    "        \"recall\": all_metrics[\"overall_recall\"],\n",
+    "        \"f1\": all_metrics[\"overall_f1\"],\n",
+    "        \"accuracy\": all_metrics[\"overall_accuracy\"],\n",
+    "    }"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "60bd54dd23de4822891a157430ff47b9",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "VBox(children=(HTML(value='<center> <img\\nsrc=https://huggingface.co/front/assets/huggingface_logo-noborder.sv…"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "notebook_login()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 23,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "e:\\open_source\\peft-folder\\ner-examples\\.venv\\Lib\\site-packages\\transformers\\training_args.py:1545: 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",
+    "    \"bert-finetuned-ner-lora\",\n",
+    "    evaluation_strategy=\"epoch\",\n",
+    "    per_device_train_batch_size=32, # decrease this for OOM error\n",
+    "    per_device_eval_batch_size=64,\n",
+    "    save_strategy=\"epoch\",\n",
+    "    learning_rate=2e-3,\n",
+    "    num_train_epochs=5,\n",
+    "    weight_decay=0.01,\n",
+    "    load_best_model_at_end=True,\n",
+    "    do_eval=True,\n",
+    "    do_predict=True,\n",
+    "    metric_for_best_model=\"accuracy\",\n",
+    "    label_names=[\"labels\"],\n",
+    "    push_to_hub=True,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 24,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "52a6470e20474df6808fcea70c8533ff",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/2195 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "e:\\open_source\\peft-folder\\ner-examples\\.venv\\Lib\\site-packages\\transformers\\models\\bert\\modeling_bert.py:440: UserWarning: 1Torch was not compiled with flash attention. (Triggered internally at C:\\actions-runner\\_work\\pytorch\\pytorch\\builder\\windows\\pytorch\\aten\\src\\ATen\\native\\transformers\\cuda\\sdp_utils.cpp:555.)\n",
+      "  attn_output = torch.nn.functional.scaled_dot_product_attention(\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "7944af882c2543238c82a6e9fd65209e",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/51 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'eval_loss': 0.09396398812532425, 'eval_precision': 0.8329652996845426, 'eval_recall': 0.8887579939414338, 'eval_f1': 0.859957661618629, 'eval_accuracy': 0.9724936716312474, 'eval_runtime': 5.7905, 'eval_samples_per_second': 561.264, 'eval_steps_per_second': 8.808, 'epoch': 1.0}\n",
+      "{'loss': 0.1864, 'grad_norm': 0.36243128776550293, 'learning_rate': 0.0015444191343963554, 'epoch': 1.14}\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "79c0a83640aa44a1b7c1fb61de11f618",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/51 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'eval_loss': 0.07848585397005081, 'eval_precision': 0.8617173809144496, 'eval_recall': 0.9102995624368899, 'eval_f1': 0.8853424993862018, 'eval_accuracy': 0.9782333549184671, 'eval_runtime': 4.8142, 'eval_samples_per_second': 675.093, 'eval_steps_per_second': 10.594, 'epoch': 2.0}\n",
+      "{'loss': 0.0867, 'grad_norm': 0.23459061980247498, 'learning_rate': 0.0010888382687927107, 'epoch': 2.28}\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "7af499f0415a41e59ba46c078f837fa0",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/51 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'eval_loss': 0.07186892628669739, 'eval_precision': 0.8772632590930941, 'eval_recall': 0.9214069336923595, 'eval_f1': 0.8987934006402364, 'eval_accuracy': 0.9796903514452228, 'eval_runtime': 4.8141, 'eval_samples_per_second': 675.103, 'eval_steps_per_second': 10.594, 'epoch': 3.0}\n",
+      "{'loss': 0.066, 'grad_norm': 0.2663992941379547, 'learning_rate': 0.0006332574031890661, 'epoch': 3.42}\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "40d883200bb3473b8dd911177685965f",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/51 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'eval_loss': 0.06592053174972534, 'eval_precision': 0.8809218950064021, 'eval_recall': 0.9262874453046113, 'eval_f1': 0.9030352748154225, 'eval_accuracy': 0.9809854594690057, 'eval_runtime': 4.8969, 'eval_samples_per_second': 663.691, 'eval_steps_per_second': 10.415, 'epoch': 4.0}\n",
+      "{'loss': 0.0505, 'grad_norm': 0.1911543756723404, 'learning_rate': 0.0001776765375854214, 'epoch': 4.56}\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "0eaa7fa8947f46a3af1e21d12bf730f8",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "  0%|          | 0/51 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "{'eval_loss': 0.06525895744562149, 'eval_precision': 0.8831439090763566, 'eval_recall': 0.928475260854931, 'eval_f1': 0.9052424317007138, 'eval_accuracy': 0.981515276387826, 'eval_runtime': 4.8357, 'eval_samples_per_second': 672.083, 'eval_steps_per_second': 10.547, 'epoch': 5.0}\n",
+      "{'train_runtime': 213.3684, 'train_samples_per_second': 329.032, 'train_steps_per_second': 10.287, 'train_loss': 0.09303037256749182, 'epoch': 5.0}\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "No files have been modified since last commit. Skipping to prevent empty commit.\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "TrainOutput(global_step=2195, training_loss=0.09303037256749182, metrics={'train_runtime': 213.3684, 'train_samples_per_second': 329.032, 'train_steps_per_second': 10.287, 'total_flos': 1949278889622816.0, 'train_loss': 0.09303037256749182, 'epoch': 5.0})"
+      ]
+     },
+     "execution_count": 24,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "trainer = Trainer(\n",
+    "    model=model,\n",
+    "    args=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",
+    ")\n",
+    "trainer.train()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 33,
+   "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"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": [
+       "[{'entity_group': 'PER',\n",
+       "  'score': np.float32(0.9957055),\n",
+       "  'word': 'Jino',\n",
+       "  'start': 11,\n",
+       "  'end': 15}]"
+      ]
+     },
+     "execution_count": 33,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# Replace this with your own checkpoint\n",
+    "model_checkpoint = \"bert-finetuned-ner-lora\"\n",
+    "token_classifier = pipeline(\n",
+    "    \"token-classification\", model = model_checkpoint, aggregation_strategy = \"simple\"\n",
+    ")\n",
+    "token_classifier(\"My name is Jino.\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "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.1"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}