Upload 3 files

lora_clm_accelerate_big_model_inference.ipynb

@@ -0,0 +1,481 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "71fbfca2",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "===================================BUG REPORT===================================\n",
+      "Welcome to bitsandbytes. For bug reports, please submit your error trace to: https://github.com/TimDettmers/bitsandbytes/issues\n",
+      "For effortless bug reporting copy-paste your error into this form: https://docs.google.com/forms/d/e/1FAIpQLScPB8emS3Thkp66nvqwmjTEgxp8Y9ufuWTzFyr9kJ5AoI47dQ/viewform?usp=sf_link\n",
+      "================================================================================\n",
+      "CUDA SETUP: CUDA runtime path found: /home/sourab/miniconda3/envs/ml/lib/libcudart.so\n",
+      "CUDA SETUP: Highest compute capability among GPUs detected: 7.5\n",
+      "CUDA SETUP: Detected CUDA version 117\n",
+      "CUDA SETUP: Loading binary /home/sourab/miniconda3/envs/ml/lib/python3.10/site-packages/bitsandbytes/libbitsandbytes_cuda117.so...\n"
+     ]
+    }
+   ],
+   "source": [
+    "from transformers import AutoModelForCausalLM\n",
+    "from peft import PeftModel, PeftConfig\n",
+    "import torch\n",
+    "from datasets import load_dataset\n",
+    "import os\n",
+    "from transformers import AutoTokenizer\n",
+    "from torch.utils.data import DataLoader\n",
+    "from transformers import default_data_collator, get_linear_schedule_with_warmup\n",
+    "from tqdm import tqdm\n",
+    "from datasets import load_dataset\n",
+    "\n",
+    "device = \"cuda\"\n",
+    "model_name_or_path = \"bigscience/bloomz-7b1\"\n",
+    "tokenizer_name_or_path = \"bigscience/bloomz-7b1\"\n",
+    "dataset_name = \"twitter_complaints\"\n",
+    "text_column = \"Tweet text\"\n",
+    "label_column = \"text_label\"\n",
+    "max_length = 64\n",
+    "lr = 1e-3\n",
+    "num_epochs = 50\n",
+    "batch_size = 8"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e1a3648b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from datasets import load_dataset\n",
+    "\n",
+    "dataset = load_dataset(\"ought/raft\", dataset_name)\n",
+    "\n",
+    "classes = [k.replace(\"_\", \" \") for k in dataset[\"train\"].features[\"Label\"].names]\n",
+    "print(classes)\n",
+    "dataset = dataset.map(\n",
+    "    lambda x: {\"text_label\": [classes[label] for label in x[\"Label\"]]},\n",
+    "    batched=True,\n",
+    "    num_proc=1,\n",
+    ")\n",
+    "print(dataset)\n",
+    "dataset[\"train\"][0]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "fe12d4d3",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "3\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "10cabeec92ab428f9a660ebaecbaf865",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Running tokenizer on dataset:   0%|          | 0/1 [00:00<?, ?ba/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "8a344e989ab34c71b230acee68b477e8",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Running tokenizer on dataset:   0%|          | 0/4 [00:00<?, ?ba/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "# data preprocessing\n",
+    "tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)\n",
+    "if tokenizer.pad_token_id is None:\n",
+    "    tokenizer.pad_token_id = tokenizer.eos_token_id\n",
+    "target_max_length = max([len(tokenizer(class_label)[\"input_ids\"]) for class_label in classes])\n",
+    "print(target_max_length)\n",
+    "\n",
+    "\n",
+    "def preprocess_function(examples):\n",
+    "    batch_size = len(examples[text_column])\n",
+    "    inputs = [f\"{text_column} : {x} Label : \" for x in examples[text_column]]\n",
+    "    targets = [str(x) for x in examples[label_column]]\n",
+    "    model_inputs = tokenizer(inputs)\n",
+    "    labels = tokenizer(targets, add_special_tokens=False)  # don't add bos token because we concatenate with inputs\n",
+    "    for i in range(batch_size):\n",
+    "        sample_input_ids = model_inputs[\"input_ids\"][i]\n",
+    "        label_input_ids = labels[\"input_ids\"][i] + [tokenizer.eos_token_id]\n",
+    "        # print(i, sample_input_ids, label_input_ids)\n",
+    "        model_inputs[\"input_ids\"][i] = sample_input_ids + label_input_ids\n",
+    "        labels[\"input_ids\"][i] = [-100] * len(sample_input_ids) + label_input_ids\n",
+    "        model_inputs[\"attention_mask\"][i] = [1] * len(model_inputs[\"input_ids\"][i])\n",
+    "    # print(model_inputs)\n",
+    "    for i in range(batch_size):\n",
+    "        sample_input_ids = model_inputs[\"input_ids\"][i]\n",
+    "        label_input_ids = labels[\"input_ids\"][i]\n",
+    "        model_inputs[\"input_ids\"][i] = [tokenizer.pad_token_id] * (\n",
+    "            max_length - len(sample_input_ids)\n",
+    "        ) + sample_input_ids\n",
+    "        model_inputs[\"attention_mask\"][i] = [0] * (max_length - len(sample_input_ids)) + model_inputs[\n",
+    "            \"attention_mask\"\n",
+    "        ][i]\n",
+    "        labels[\"input_ids\"][i] = [-100] * (max_length - len(sample_input_ids)) + label_input_ids\n",
+    "        model_inputs[\"input_ids\"][i] = torch.tensor(model_inputs[\"input_ids\"][i][:max_length])\n",
+    "        model_inputs[\"attention_mask\"][i] = torch.tensor(model_inputs[\"attention_mask\"][i][:max_length])\n",
+    "        labels[\"input_ids\"][i] = torch.tensor(labels[\"input_ids\"][i][:max_length])\n",
+    "    model_inputs[\"labels\"] = labels[\"input_ids\"]\n",
+    "    return model_inputs\n",
+    "\n",
+    "\n",
+    "processed_datasets = dataset.map(\n",
+    "    preprocess_function,\n",
+    "    batched=True,\n",
+    "    num_proc=1,\n",
+    "    remove_columns=dataset[\"train\"].column_names,\n",
+    "    load_from_cache_file=False,\n",
+    "    desc=\"Running tokenizer on dataset\",\n",
+    ")\n",
+    "\n",
+    "train_dataset = processed_datasets[\"train\"]\n",
+    "\n",
+    "\n",
+    "train_dataloader = DataLoader(\n",
+    "    train_dataset, shuffle=True, collate_fn=default_data_collator, batch_size=batch_size, pin_memory=True\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2795b9d0",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def test_preprocess_function(examples):\n",
+    "    batch_size = len(examples[text_column])\n",
+    "    inputs = [f\"{text_column} : {x} Label : \" for x in examples[text_column]]\n",
+    "    model_inputs = tokenizer(inputs)\n",
+    "    # print(model_inputs)\n",
+    "    for i in range(batch_size):\n",
+    "        sample_input_ids = model_inputs[\"input_ids\"][i]\n",
+    "        model_inputs[\"input_ids\"][i] = [tokenizer.pad_token_id] * (\n",
+    "            max_length - len(sample_input_ids)\n",
+    "        ) + sample_input_ids\n",
+    "        model_inputs[\"attention_mask\"][i] = [0] * (max_length - len(sample_input_ids)) + model_inputs[\n",
+    "            \"attention_mask\"\n",
+    "        ][i]\n",
+    "        model_inputs[\"input_ids\"][i] = torch.tensor(model_inputs[\"input_ids\"][i][:max_length])\n",
+    "        model_inputs[\"attention_mask\"][i] = torch.tensor(model_inputs[\"attention_mask\"][i][:max_length])\n",
+    "    return model_inputs\n",
+    "\n",
+    "\n",
+    "processed_datasets = dataset.map(\n",
+    "    test_preprocess_function,\n",
+    "    batched=True,\n",
+    "    num_proc=1,\n",
+    "    remove_columns=dataset[\"train\"].column_names,\n",
+    "    load_from_cache_file=False,\n",
+    "    desc=\"Running tokenizer on dataset\",\n",
+    ")\n",
+    "\n",
+    "eval_dataset = processed_datasets[\"train\"]\n",
+    "test_dataset = processed_datasets[\"test\"]\n",
+    "\n",
+    "eval_dataloader = DataLoader(eval_dataset, collate_fn=default_data_collator, batch_size=batch_size, pin_memory=True)\n",
+    "test_dataloader = DataLoader(test_dataset, collate_fn=default_data_collator, batch_size=batch_size, pin_memory=True)\n",
+    "print(next(iter(eval_dataloader)))\n",
+    "print(next(iter(test_dataloader)))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "42b14a11",
+   "metadata": {},
+   "source": [
+    "You can load model from hub or local\n",
+    "\n",
+    "- Load model from Hugging Face Hub, you can change to your own model id\n",
+    "```python\n",
+    "peft_model_id = \"username/twitter_complaints_bigscience_bloomz-7b1_LORA_CAUSAL_LM\"\n",
+    "```\n",
+    "- Or load model form local\n",
+    "```python\n",
+    "peft_model_id = \"twitter_complaints_bigscience_bloomz-7b1_LORA_CAUSAL_LM\"\n",
+    "```"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "9caac014",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/home/sourab/pet/src/peft/tuners/lora.py:143: UserWarning: fan_in_fan_out is set to True but the target module is not a Conv1D. Setting fan_in_fan_out to False.\n",
+      "  warnings.warn(\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "bc38030106a14173a1363eb1ee388eda",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Downloading:   0%|          | 0.00/15.8M [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "from peft import PeftModel, PeftConfig\n",
+    "\n",
+    "max_memory = {0: \"1GIB\", 1: \"1GIB\", 2: \"2GIB\", 3: \"10GIB\", \"cpu\": \"30GB\"}\n",
+    "peft_model_id = \"smangrul/twitter_complaints_bigscience_bloomz-7b1_LORA_CAUSAL_LM\"\n",
+    "config = PeftConfig.from_pretrained(peft_model_id)\n",
+    "model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, device_map=\"auto\", max_memory=max_memory)\n",
+    "model = PeftModel.from_pretrained(model, peft_model_id, device_map=\"auto\", max_memory=max_memory)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 35,
+   "id": "6fac10b5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "2a08ee6d",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'base_model.model.transformer.word_embeddings': 3,\n",
+       " 'base_model.model.lm_head': 3,\n",
+       " 'base_model.model.transformer.word_embeddings_layernorm': 3,\n",
+       " 'base_model.model.transformer.h.0': 3,\n",
+       " 'base_model.model.transformer.h.1': 3,\n",
+       " 'base_model.model.transformer.h.2': 3,\n",
+       " 'base_model.model.transformer.h.3': 3,\n",
+       " 'base_model.model.transformer.h.4': 3,\n",
+       " 'base_model.model.transformer.h.5': 3,\n",
+       " 'base_model.model.transformer.h.6': 3,\n",
+       " 'base_model.model.transformer.h.7': 3,\n",
+       " 'base_model.model.transformer.h.8': 'cpu',\n",
+       " 'base_model.model.transformer.h.9': 'cpu',\n",
+       " 'base_model.model.transformer.h.10': 'cpu',\n",
+       " 'base_model.model.transformer.h.11': 'cpu',\n",
+       " 'base_model.model.transformer.h.12': 'cpu',\n",
+       " 'base_model.model.transformer.h.13': 'cpu',\n",
+       " 'base_model.model.transformer.h.14': 'cpu',\n",
+       " 'base_model.model.transformer.h.15': 'cpu',\n",
+       " 'base_model.model.transformer.h.16': 'cpu',\n",
+       " 'base_model.model.transformer.h.17': 'cpu',\n",
+       " 'base_model.model.transformer.h.18': 'cpu',\n",
+       " 'base_model.model.transformer.h.19': 'cpu',\n",
+       " 'base_model.model.transformer.h.20': 'cpu',\n",
+       " 'base_model.model.transformer.h.21': 'cpu',\n",
+       " 'base_model.model.transformer.h.22': 'cpu',\n",
+       " 'base_model.model.transformer.h.23': 'cpu',\n",
+       " 'base_model.model.transformer.h.24': 'cpu',\n",
+       " 'base_model.model.transformer.h.25': 'cpu',\n",
+       " 'base_model.model.transformer.h.26': 'cpu',\n",
+       " 'base_model.model.transformer.h.27': 'cpu',\n",
+       " 'base_model.model.transformer.h.28': 'cpu',\n",
+       " 'base_model.model.transformer.h.29': 'cpu',\n",
+       " 'base_model.model.transformer.ln_f': 'cpu'}"
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "model.hf_device_map"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 34,
+   "id": "b33be5e6",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "@HondaCustSvc Your customer service has been horrible during the recall process. I will never purchase a Honda again.\n",
+      "{'input_ids': tensor([[227985,   5484,    915,   2566, 216744,     38,   1316,     54,  42705,\n",
+      "          32465,  52166,   9440,   1809,   3784,  88483,   9411,    368,  84342,\n",
+      "           4451,     17,    473,   2152,  11705,  82406,    267,  51591,   5734,\n",
+      "             17,  77658,    915,    210]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
+      "         1, 1, 1, 1, 1, 1, 1]])}\n",
+      "tensor([[227985,   5484,    915,   2566, 216744,     38,   1316,     54,  42705,\n",
+      "          32465,  52166,   9440,   1809,   3784,  88483,   9411,    368,  84342,\n",
+      "           4451,     17,    473,   2152,  11705,  82406,    267,  51591,   5734,\n",
+      "             17,  77658,    915,    210,  16449,   5952,      3,      3,      3,\n",
+      "              3,      3,      3,      3,      3]])\n",
+      "['Tweet text : @HondaCustSvc Your customer service has been horrible during the recall process. I will never purchase a Honda again. Label : complaint']\n"
+     ]
+    }
+   ],
+   "source": [
+    "model.eval()\n",
+    "i = 89\n",
+    "inputs = tokenizer(f'{text_column} : {dataset[\"test\"][i][\"Tweet text\"]} Label : ', return_tensors=\"pt\")\n",
+    "print(dataset[\"test\"][i][\"Tweet text\"])\n",
+    "print(inputs)\n",
+    "\n",
+    "with torch.no_grad():\n",
+    "    outputs = model.generate(input_ids=inputs[\"input_ids\"], max_new_tokens=10)\n",
+    "    print(outputs)\n",
+    "    print(tokenizer.batch_decode(outputs.detach().cpu().numpy(), skip_special_tokens=True))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "b6d6cd5b",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|███████████████████████████████████████████████████████████████████████████���████████████████| 7/7 [01:42<00:00, 14.70s/it]\n"
+     ]
+    }
+   ],
+   "source": [
+    "model.eval()\n",
+    "eval_preds = []\n",
+    "for _, batch in enumerate(tqdm(eval_dataloader)):\n",
+    "    batch = {k: v for k, v in batch.items() if k != \"labels\"}\n",
+    "    with torch.no_grad():\n",
+    "        outputs = model.generate(**batch, max_new_tokens=10)\n",
+    "    preds = outputs[:, max_length:].detach().cpu().numpy()\n",
+    "    eval_preds.extend(tokenizer.batch_decode(preds, skip_special_tokens=True))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "61264abe",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "accuracy=100.0\n",
+      "eval_preds[:10]=['no complaint', 'no complaint', 'complaint', 'complaint', 'no complaint', 'no complaint', 'no complaint', 'complaint', 'complaint', 'no complaint']\n",
+      "dataset['train'][label_column][:10]=['no complaint', 'no complaint', 'complaint', 'complaint', 'no complaint', 'no complaint', 'no complaint', 'complaint', 'complaint', 'no complaint']\n"
+     ]
+    }
+   ],
+   "source": [
+    "correct = 0\n",
+    "total = 0\n",
+    "for pred, true in zip(eval_preds, dataset[\"train\"][label_column]):\n",
+    "    if pred.strip() == true.strip():\n",
+    "        correct += 1\n",
+    "    total += 1\n",
+    "accuracy = correct / total * 100\n",
+    "print(f\"{accuracy=}\")\n",
+    "print(f\"{eval_preds[:10]=}\")\n",
+    "print(f\"{dataset['train'][label_column][:10]=}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a70802a3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model.eval()\n",
+    "test_preds = []\n",
+    "\n",
+    "for _, batch in enumerate(tqdm(test_dataloader)):\n",
+    "    batch = {k: v for k, v in batch.items() if k != \"labels\"}\n",
+    "    with torch.no_grad():\n",
+    "        outputs = model.generate(**batch, max_new_tokens=10)\n",
+    "    preds = outputs[:, max_length:].detach().cpu().numpy()\n",
+    "    test_preds.extend(tokenizer.batch_decode(preds, skip_special_tokens=True))\n",
+    "    if len(test_preds) > 100:\n",
+    "        break\n",
+    "test_preds"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e1c4ad9c",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "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.10.4"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "aee8b7b246df8f9039afb4144a1f6fd8d2ca17a180786b69acc140d282b71a49"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

lora_clm_with_additional_tokens.ipynb

@@ -0,0 +1,1012 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "5f239612-620e-4430-8685-9fdc6b179b41",
+   "metadata": {},
+   "source": [
+    "# Training PEFT models with new tokens being added to the embedding layers and tokenizer\n",
+    "\n",
+    "In this example, we will learn how to train a LoRA model when adding new tokens to the tokenizer and model. \n",
+    "This is a common usecase when doing the following:\n",
+    "1. Instruction finetuning with new tokens beind added such as `<|user|>`, `<|assistant|>`, `<|system|>`, `</s>`, `<s>` to properly format the conversations\n",
+    "2. Finetuning on a specific language wherein language spoecific tokens are added, e.g., korean tokens being added to vocabulary for finetuning LLM on Korean datasets.\n",
+    "3. Instruction finetuning to return outputs in certain format to enable agent behaviour new tokens such as `<|FUNCTIONS|>`, `<|BROWSE|>`, `<|TEXT2IMAGE|>`, `<|ASR|>`, `<|TTS|>`, `<|GENERATECODE|>`, `<|RAG|>`.\n",
+    "\n",
+    "In such cases, you add the Embedding modules to the LORA `target_modules`. PEFT will take care of saving the embedding layers with the new added tokens along with the adapter weights that were trained on the specific initialization of the embeddings weights of the added tokens."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b27c55e8-edaa-4059-90bc-d6096d596902",
+   "metadata": {},
+   "source": [
+    "Let's import the necessary libraries"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "6f864c90",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "\n",
+    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"3\"\n",
+    "os.environ[\"WANDB_PROJECT\"] = \"PeftExamples\"\n",
+    "import transformers\n",
+    "from peft import (\n",
+    "    LoraConfig,\n",
+    "    PeftConfig,\n",
+    "    PeftModel,\n",
+    "    get_peft_model,\n",
+    "    prepare_model_for_int8_training,\n",
+    ")\n",
+    "from transformers import (\n",
+    "    AutoModelForCausalLM,\n",
+    "    AutoTokenizer,\n",
+    "    HfArgumentParser,\n",
+    "    TrainingArguments,\n",
+    "    Trainer,\n",
+    "    default_data_collator,\n",
+    ")\n",
+    "import torch\n",
+    "from dataclasses import dataclass, field\n",
+    "from typing import Optional\n",
+    "from dataclass_csv import DataclassReader\n",
+    "from torch.utils.data import Dataset, DataLoader\n",
+    "\n",
+    "from enum import Enum"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "74950a3f-bb63-4ce5-9e2b-1b83f92b13a2",
+   "metadata": {},
+   "source": [
+    "## Prepare Model and Tokenizer"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "76763f5e-64b2-409b-8845-ae5589f8a4e0",
+   "metadata": {},
+   "source": [
+    "Now, we will be adding 27 new tokens as well as replace the existing pad, bos and eos tokens of the model."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "fd0498ea-547e-418d-bf13-c9abafdd5476",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class SpecialTokens(str, Enum):\n",
+    "    begin_target = \"<|begintarget|>\"\n",
+    "    end_target = \"<|endtarget|>\"\n",
+    "    begin_context = \"<|begincontext|>\"\n",
+    "    end_context = \"<|endcontext|>\"\n",
+    "    system = \"<|system|>\"\n",
+    "    user = \"<|user|>\"\n",
+    "    begin_last_user_utterance = \"<|beginlastuserutterance|>\"\n",
+    "    end_last_user_utterance = \"<|endlastuserutterance|>\"\n",
+    "    begin_dsts = \"<|begindsts|>\"\n",
+    "    end_dsts = \"<|enddsts|>\"\n",
+    "    begin_dst = \"<|begindst|>\"\n",
+    "    end_dst = \"<|enddst|>\"\n",
+    "    begin_belief = \"<|beginbelief|>\"\n",
+    "    end_belief = \"<|endbelief|>\"\n",
+    "    begin_response = \"<|beginresponse|>\"\n",
+    "    end_response = \"<|endresponse|>\"\n",
+    "    begin_action = \"<|beginaction|>\"\n",
+    "    end_action = \"<|endaction|>\"\n",
+    "    begin_user_action = \"<|beginuseraction|>\"\n",
+    "    end_user_action = \"<|enduseraction|>\"\n",
+    "    sys_actions = \"<|sysactions|>\"\n",
+    "    begin_intent = \"<|beginintent|>\"\n",
+    "    end_intent = \"<|endintent|>\"\n",
+    "    begin_requested_slots = \"<|beginrequestedslots|>\"\n",
+    "    end_requested_slots = \"<|endrequestedslots|>\"\n",
+    "    pad_token = \"<|pad|>\"\n",
+    "    bos_token = \"<|startoftext|>\"\n",
+    "\n",
+    "    @classmethod\n",
+    "    def list(cls):\n",
+    "        return [c.value for c in cls]"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ae4a4255-5f13-4eef-a024-4f1de0f2173b",
+   "metadata": {},
+   "source": [
+    "We will be finetuning Mistral-7B model. Let's load the tokenizer and add the special tokens followed by loading the base model and resizzing the embedding layers to accomodate the newly added tokens."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "f0eedef9",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "91c67b6377fc4dd7977bf544de784d51",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/plain": [
+       "Embedding(32027, 4096)"
+      ]
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "model_name = \"mistralai/Mistral-7B-v0.1\"\n",
+    "tokenizer = AutoTokenizer.from_pretrained(\n",
+    "    model_name,\n",
+    "    pad_token=SpecialTokens.pad_token.value,\n",
+    "    bos_token=SpecialTokens.bos_token.value,\n",
+    "    eos_token=SpecialTokens.end_target.value,\n",
+    "    additional_special_tokens=SpecialTokens.list(),\n",
+    ")\n",
+    "model = AutoModelForCausalLM.from_pretrained(\n",
+    "    model_name,\n",
+    "    low_cpu_mem_usage=True\n",
+    "    # use_flash_attention_2=True, # leading to an error\n",
+    ")\n",
+    "model.resize_token_embeddings(len(tokenizer))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "88439ed6-9974-4918-80df-ec78b05b4185",
+   "metadata": {},
+   "source": [
+    "## Apply LoRA"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "80967087",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "trainable params: 31,886,720 || all params: 7,273,840,000 || trainable%: 0.43837532857472805\n",
+      "None\n",
+      "PeftModel(\n",
+      "  (base_model): LoraModel(\n",
+      "    (model): MistralForCausalLM(\n",
+      "      (model): MistralModel(\n",
+      "        (embed_tokens): lora.Embedding(\n",
+      "          (base_layer): Embedding(32027, 4096)\n",
+      "          (lora_dropout): ModuleDict(\n",
+      "            (default): Identity()\n",
+      "          )\n",
+      "          (lora_A): ModuleDict()\n",
+      "          (lora_B): ModuleDict()\n",
+      "          (lora_embedding_A): ParameterDict(  (default): Parameter containing: [torch.FloatTensor of size 64x32027])\n",
+      "          (lora_embedding_B): ParameterDict(  (default): Parameter containing: [torch.FloatTensor of size 4096x64])\n",
+      "        )\n",
+      "        (layers): ModuleList(\n",
+      "          (0-31): 32 x MistralDecoderLayer(\n",
+      "            (self_attn): MistralAttention(\n",
+      "              (q_proj): lora.Linear(\n",
+      "                (base_layer): Linear(in_features=4096, out_features=4096, bias=False)\n",
+      "                (lora_dropout): ModuleDict(\n",
+      "                  (default): Identity()\n",
+      "                )\n",
+      "                (lora_A): ModuleDict(\n",
+      "                  (default): Linear(in_features=4096, out_features=64, bias=False)\n",
+      "                )\n",
+      "                (lora_B): ModuleDict(\n",
+      "                  (default): Linear(in_features=64, out_features=4096, bias=False)\n",
+      "                )\n",
+      "                (lora_embedding_A): ParameterDict()\n",
+      "                (lora_embedding_B): ParameterDict()\n",
+      "              )\n",
+      "              (k_proj): Linear(in_features=4096, out_features=1024, bias=False)\n",
+      "              (v_proj): lora.Linear(\n",
+      "                (base_layer): Linear(in_features=4096, out_features=1024, bias=False)\n",
+      "                (lora_dropout): ModuleDict(\n",
+      "                  (default): Identity()\n",
+      "                )\n",
+      "                (lora_A): ModuleDict(\n",
+      "                  (default): Linear(in_features=4096, out_features=64, bias=False)\n",
+      "                )\n",
+      "                (lora_B): ModuleDict(\n",
+      "                  (default): Linear(in_features=64, out_features=1024, bias=False)\n",
+      "                )\n",
+      "                (lora_embedding_A): ParameterDict()\n",
+      "                (lora_embedding_B): ParameterDict()\n",
+      "              )\n",
+      "              (o_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
+      "              (rotary_emb): MistralRotaryEmbedding()\n",
+      "            )\n",
+      "            (mlp): MistralMLP(\n",
+      "              (gate_proj): Linear(in_features=4096, out_features=14336, bias=False)\n",
+      "              (up_proj): Linear(in_features=4096, out_features=14336, bias=False)\n",
+      "              (down_proj): Linear(in_features=14336, out_features=4096, bias=False)\n",
+      "              (act_fn): SiLU()\n",
+      "            )\n",
+      "            (input_layernorm): MistralRMSNorm()\n",
+      "            (post_attention_layernorm): MistralRMSNorm()\n",
+      "          )\n",
+      "        )\n",
+      "        (norm): MistralRMSNorm()\n",
+      "      )\n",
+      "      (lm_head): lora.Linear(\n",
+      "        (base_layer): Linear(in_features=4096, out_features=32027, bias=False)\n",
+      "        (lora_dropout): ModuleDict(\n",
+      "          (default): Identity()\n",
+      "        )\n",
+      "        (lora_A): ModuleDict(\n",
+      "          (default): Linear(in_features=4096, out_features=64, bias=False)\n",
+      "        )\n",
+      "        (lora_B): ModuleDict(\n",
+      "          (default): Linear(in_features=64, out_features=32027, bias=False)\n",
+      "        )\n",
+      "        (lora_embedding_A): ParameterDict()\n",
+      "        (lora_embedding_B): ParameterDict()\n",
+      "      )\n",
+      "    )\n",
+      "  )\n",
+      ")\n"
+     ]
+    }
+   ],
+   "source": [
+    "config = LoraConfig(\n",
+    "    r=64, lora_alpha=128, lora_dropout=0.0, target_modules=[\"embed_tokens\", \"lm_head\", \"q_proj\", \"v_proj\"]\n",
+    ")\n",
+    "model = get_peft_model(model, config)\n",
+    "print(model.print_trainable_parameters())\n",
+    "print(model)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "15ac9945-4fcb-45f4-9478-d99a25a519cc",
+   "metadata": {},
+   "source": [
+    "## Preapre Dataset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "c6980d59-42d4-4a27-84cc-a9719302088b",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "33d9539232da48f3ae922216b98ae462",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Running tokenizer on dataset:   0%|          | 0/986 [00:00<?, ? examples/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "b7a33811d93742099140240cad91b679",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Running tokenizer on dataset:   0%|          | 0/247 [00:00<?, ? examples/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "from datasets import load_dataset\n",
+    "\n",
+    "dataset = load_dataset(\"smangrul/assistant_chatbot_dataset\")\n",
+    "dataset = dataset[\"train\"].train_test_split(0.2)\n",
+    "\n",
+    "text_column = \"context\"\n",
+    "label_column = \"target\"\n",
+    "max_length = 512\n",
+    "\n",
+    "\n",
+    "def preprocess_function(examples):\n",
+    "    batch_size = len(examples[text_column])\n",
+    "    targets = [str(x) for x in examples[label_column]]\n",
+    "    model_inputs = tokenizer(examples[text_column])\n",
+    "    labels = tokenizer(targets, add_special_tokens=False)  # don't add bos token because we concatenate with inputs\n",
+    "    for i in range(batch_size):\n",
+    "        sample_input_ids = model_inputs[\"input_ids\"][i]\n",
+    "        label_input_ids = labels[\"input_ids\"][i] + [tokenizer.eos_token_id]\n",
+    "        # print(i, sample_input_ids, label_input_ids)\n",
+    "        model_inputs[\"input_ids\"][i] = sample_input_ids + label_input_ids\n",
+    "        labels[\"input_ids\"][i] = [-100] * len(sample_input_ids) + label_input_ids\n",
+    "        model_inputs[\"attention_mask\"][i] = [1] * len(model_inputs[\"input_ids\"][i])\n",
+    "    # print(model_inputs)\n",
+    "    for i in range(batch_size):\n",
+    "        sample_input_ids = model_inputs[\"input_ids\"][i]\n",
+    "        label_input_ids = labels[\"input_ids\"][i]\n",
+    "        model_inputs[\"input_ids\"][i] = [tokenizer.pad_token_id] * (\n",
+    "            max_length - len(sample_input_ids)\n",
+    "        ) + sample_input_ids\n",
+    "        model_inputs[\"attention_mask\"][i] = [0] * (max_length - len(sample_input_ids)) + model_inputs[\n",
+    "            \"attention_mask\"\n",
+    "        ][i]\n",
+    "        labels[\"input_ids\"][i] = [-100] * (max_length - len(sample_input_ids)) + label_input_ids\n",
+    "        model_inputs[\"input_ids\"][i] = model_inputs[\"input_ids\"][i][:max_length]\n",
+    "        model_inputs[\"attention_mask\"][i] = model_inputs[\"attention_mask\"][i][:max_length]\n",
+    "        labels[\"input_ids\"][i] = labels[\"input_ids\"][i][:max_length]\n",
+    "    model_inputs[\"labels\"] = labels[\"input_ids\"]\n",
+    "    return model_inputs\n",
+    "\n",
+    "\n",
+    "processed_datasets = dataset.map(\n",
+    "    preprocess_function,\n",
+    "    batched=True,\n",
+    "    num_proc=1,\n",
+    "    remove_columns=dataset[\"train\"].column_names,\n",
+    "    load_from_cache_file=False,\n",
+    "    desc=\"Running tokenizer on dataset\",\n",
+    ")\n",
+    "\n",
+    "train_dataset = processed_datasets[\"train\"]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "5671b1ee-dca4-4705-8399-5c2967b9fb5c",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "Dataset({\n",
+       "    features: ['input_ids', 'attention_mask', 'labels'],\n",
+       "    num_rows: 986\n",
+       "})"
+      ]
+     },
+     "execution_count": 6,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "train_dataset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "3f38888e-4382-415b-869d-7202a816606a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_dataloader = DataLoader(\n",
+    "    train_dataset, shuffle=True, collate_fn=default_data_collator, batch_size=8, pin_memory=True\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "53b9e552-4c5d-43e8-a9cd-8073af8d4280",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'input_ids': tensor([[32002, 32002, 32002,  ..., 32017, 32001, 32001],\n",
+       "         [32002, 32002, 32002,  ..., 32017, 32001, 32001],\n",
+       "         [32002, 32002, 32002,  ..., 32017, 32001, 32001],\n",
+       "         ...,\n",
+       "         [32002, 32002, 32002,  ..., 32017, 32001, 32001],\n",
+       "         [32002, 32002, 32002,  ..., 32017, 32001, 32001],\n",
+       "         [32002, 32002, 32002,  ..., 32017, 32001, 32001]]),\n",
+       " 'attention_mask': tensor([[0, 0, 0,  ..., 1, 1, 1],\n",
+       "         [0, 0, 0,  ..., 1, 1, 1],\n",
+       "         [0, 0, 0,  ..., 1, 1, 1],\n",
+       "         ...,\n",
+       "         [0, 0, 0,  ..., 1, 1, 1],\n",
+       "         [0, 0, 0,  ..., 1, 1, 1],\n",
+       "         [0, 0, 0,  ..., 1, 1, 1]]),\n",
+       " 'labels': tensor([[ -100,  -100,  -100,  ..., 32017, 32001, 32001],\n",
+       "         [ -100,  -100,  -100,  ..., 32017, 32001, 32001],\n",
+       "         [ -100,  -100,  -100,  ..., 32017, 32001, 32001],\n",
+       "         ...,\n",
+       "         [ -100,  -100,  -100,  ..., 32017, 32001, 32001],\n",
+       "         [ -100,  -100,  -100,  ..., 32017, 32001, 32001],\n",
+       "         [ -100,  -100,  -100,  ..., 32017, 32001, 32001]])}"
+      ]
+     },
+     "execution_count": 8,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "next(iter(train_dataloader))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "7de31ee2-185e-4658-9ad1-ae5f6bc3a611",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "\"<|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|pad|><|startoftext|><|begincontext|><|user|> Can you find me place to eat?<|system|> What kind of food would you like to have and where would you like me to search in?<|user|> Food kind of California will be perfect in SF.<|system|> There are 10 restaurants, Al's Place is one of the good restaurant in San Francisco.<|user|> Can you look for any other restaurant?<|system|> Alta Msp is one of the good restaurant in San Francisco.<|beginlastuserutterance|> Can you find me the address?<|endlastuserutterance|><|endcontext|><|begintarget|><|begindsts|><|begindst|><|beginintent|> FindRestaurants<|endintent|><|beginrequestedslots|> Restaurants^street_address<|endrequestedslots|><|beginbelief|> Restaurants^city->SF~San Francisco|Restaurants^cuisine->California<|endbelief|><|enddst|><|enddsts|><|beginuseraction|> REQUEST->Restaurants^street_address~<|enduseraction|><|beginaction|> INFORM->Restaurants^street_address~1275 Minnesota Street<|endaction|><|beginresponse|> The street address of the restaurant is 1275 Minnesota Street.<|endresponse|><|endtarget|><|endtarget|>\""
+      ]
+     },
+     "execution_count": 9,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "tokenizer.decode(train_dataset[0][\"input_ids\"])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "239d1c83-196d-471e-9bf7-5f36dafa9894",
+   "metadata": {},
+   "source": [
+    "# Train the model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "ec80d6ee",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Detected kernel version 5.4.0, which is below the recommended minimum of 5.5.0; this can cause the process to hang. It is recommended to upgrade the kernel to the minimum version or higher.\n",
+      "Failed to detect the name of this notebook, you can set it manually with the WANDB_NOTEBOOK_NAME environment variable to enable code saving.\n",
+      "\u001b[34m\u001b[1mwandb\u001b[0m: Currently logged in as: \u001b[33msmangrul\u001b[0m. Use \u001b[1m`wandb login --relogin`\u001b[0m to force relogin\n"
+     ]
+    },
+    {
+     "data": {
+      "text/html": [
+       "Tracking run with wandb version 0.16.0"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "Run data is saved locally in <code>/raid/sourab/temp/wandb/run-20231128_230934-edod21gq</code>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       "Syncing run <strong><a href='https://wandb.ai/smangrul/PeftExamples/runs/edod21gq' target=\"_blank\">ethereal-eon-1</a></strong> to <a href='https://wandb.ai/smangrul/PeftExamples' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       " View project at <a href='https://wandb.ai/smangrul/PeftExamples' target=\"_blank\">https://wandb.ai/smangrul/PeftExamples</a>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/html": [
+       " View run at <a href='https://wandb.ai/smangrul/PeftExamples/runs/edod21gq' target=\"_blank\">https://wandb.ai/smangrul/PeftExamples/runs/edod21gq</a>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...\n"
+     ]
+    },
+    {
+     "data": {
+      "text/html": [
+       "\n",
+       "    <div>\n",
+       "      \n",
+       "      <progress value='246' max='246' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
+       "      [246/246 05:51, Epoch 2/2]\n",
+       "    </div>\n",
+       "    <table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       " <tr style=\"text-align: left;\">\n",
+       "      <th>Step</th>\n",
+       "      <th>Training Loss</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <td>10</td>\n",
+       "      <td>5.189800</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>20</td>\n",
+       "      <td>3.745500</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>30</td>\n",
+       "      <td>2.371500</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>40</td>\n",
+       "      <td>1.630200</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>50</td>\n",
+       "      <td>1.302600</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>60</td>\n",
+       "      <td>0.999400</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>70</td>\n",
+       "      <td>0.704100</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>80</td>\n",
+       "      <td>0.527800</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>90</td>\n",
+       "      <td>0.509700</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>100</td>\n",
+       "      <td>0.382300</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>110</td>\n",
+       "      <td>0.318200</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>120</td>\n",
+       "      <td>0.323500</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>130</td>\n",
+       "      <td>0.263400</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>140</td>\n",
+       "      <td>0.290900</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>150</td>\n",
+       "      <td>0.277400</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>160</td>\n",
+       "      <td>0.232800</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>170</td>\n",
+       "      <td>0.223600</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>180</td>\n",
+       "      <td>0.229600</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>190</td>\n",
+       "      <td>0.233100</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>200</td>\n",
+       "      <td>0.210200</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>210</td>\n",
+       "      <td>0.245800</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>220</td>\n",
+       "      <td>0.197300</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>230</td>\n",
+       "      <td>0.210100</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <td>240</td>\n",
+       "      <td>0.209800</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table><p>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/plain": [
+       "TrainOutput(global_step=246, training_loss=0.8516577879587809, metrics={'train_runtime': 354.9013, 'train_samples_per_second': 5.556, 'train_steps_per_second': 0.693, 'total_flos': 4.318233532091597e+16, 'train_loss': 0.8516577879587809, 'epoch': 2.0})"
+      ]
+     },
+     "execution_count": 10,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "training_args = TrainingArguments(\n",
+    "    output_dir=\"mistral_lora_clm_with_added_tokens\",\n",
+    "    num_train_epochs=2,\n",
+    "    save_total_limit=5,\n",
+    "    per_device_train_batch_size=8,\n",
+    "    warmup_steps=10,\n",
+    "    weight_decay=0.0001,\n",
+    "    dataloader_drop_last=True,\n",
+    "    bf16=True,\n",
+    "    logging_steps=10,\n",
+    "    learning_rate=1e-5,\n",
+    "    gradient_checkpointing=True,\n",
+    "    gradient_checkpointing_kwargs={\"use_reentrant\": False},\n",
+    "    remove_unused_columns=False,\n",
+    "    hub_model_id=\"smangrul/mistral_lora_clm_with_added_tokens\",\n",
+    "    push_to_hub=True,\n",
+    "    hub_private_repo=True,\n",
+    ")\n",
+    "trainer = Trainer(\n",
+    "    model=model,\n",
+    "    args=training_args,\n",
+    "    train_dataset=train_dataset,\n",
+    "    data_collator=default_data_collator,\n",
+    ")\n",
+    "# model.config.use_cache = False\n",
+    "trainer.train()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7bc1cbed-4eb9-4aaa-ab5f-5b91bf432307",
+   "metadata": {},
+   "source": [
+    "# Check the model output on a sample from evaluation dataset"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "71851793",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "context=\"<|begincontext|><|user|>Can you find me a place to eat please?<|system|>Where at? And what kind of cuisine are you craving?<|user|>Somewhere in SF, and I am really craving Thai food at the moment!<|system|>I found a bunch of restaurants, there's actually 10 that you might like in San Francisco, one of them being Baan Thai House & Wine Bar<|user|>How can I reach them? And what's their address?<|system|>You can reach them by phone at 415-379-4505 and visit them at 534 Irving Street<|beginlastuserutterance|>Great, that restaurant sounds good<|endlastuserutterance|><|endcontext|>\" \n",
+      "\n",
+      " target_predicted='<|begintarget|><|begindsts|><|begindst|><|beginintent|> FindRestaurants<|endintent|><|beginbelief|> Restaurants^city->SF~San Francisco|Restaurants^cuisine->Thai|Restaurants^restaurant_name->Baan Thai House & Wine Bar<|endbelief|><|enddst|><|enddsts|><|beginuseraction|> REQUEST->Restaurants^phone_number~|REQUEST->Restaurants^street_address~<|enduseraction|><|beginaction|> INFORM->Restaurants^phone_number~415-379-4505|INFORM->Restaurants^street_address~534 Irving Street<|endaction|><|beginresponse|> Great, the phone number is 415-379-4505 and the address is 534 Irving Street<|endresponse|><|endtarget|>' \n",
+      "\n",
+      " target='<|begintarget|><|begindsts|><|begindst|><|beginintent|>FindRestaurants<|endintent|><|beginbelief|>Restaurants^city->SF~San Francisco|Restaurants^cuisine->Thai|Restaurants^restaurant_name->Baan Thai House & Wine Bar<|endbelief|><|enddst|><|enddsts|><|beginuseraction|>SELECT->Restaurants^~<|enduseraction|><|beginaction|>OFFER_INTENT->Restaurants^intent~ReserveRestaurant<|endaction|><|beginresponse|>Want me to book a table?<|endresponse|><|endtarget|>'\n"
+     ]
+    }
+   ],
+   "source": [
+    "import random\n",
+    "\n",
+    "i = random.randint(0, len(dataset[\"test\"]))\n",
+    "context = dataset[\"test\"][i][\"context\"]\n",
+    "\n",
+    "batch = tokenizer(context, return_tensors=\"pt\")\n",
+    "batch = {k: v.to(\"cuda\") for k, v in batch.items()}\n",
+    "model.eval()\n",
+    "output_tokens = model.generate(\n",
+    "    **batch,\n",
+    "    max_new_tokens=256,\n",
+    "    do_sample=True,\n",
+    "    temperature=0.2,\n",
+    "    top_p=0.95,\n",
+    "    top_k=50,\n",
+    "    eos_token_id=tokenizer.eos_token_id,\n",
+    "    pad_token_id=tokenizer.pad_token_id,\n",
+    ")\n",
+    "target_predicted = tokenizer.decode(output_tokens[0], skip_special_tokens=False).split(\"<|endcontext|>\")[1]\n",
+    "target = dataset[\"test\"][i][\"target\"]\n",
+    "print(f\"{context=} \\n\\n {target_predicted=} \\n\\n {target=}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "f940a660-2f7c-4a3a-b412-3f037aedb890",
+   "metadata": {},
+   "source": [
+    "# Save the Adapter model "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7ebe05e9-9b93-42f6-bba8-46b8cc3d100f",
+   "metadata": {},
+   "source": [
+    "When the lora layers are applied to embedding layers, the corresponding base model embedding layers are also saved. "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "3d7459ba-caa8-4f10-aa70-89be4541cbdf",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/raid/sourab/peft/src/peft/utils/save_and_load.py:128: UserWarning: Setting `is_embedding_layer_resized` to `True` as embedding layers found in `target_modules`\n",
+      "  warnings.warn(\"Setting `is_embedding_layer_resized` to `True` as embedding layers found in `target_modules`\")\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "8d23186832014f209939ab83e79da011",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Upload 3 LFS files:   0%|          | 0/3 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "a3d831bc7d8843038364e821aacff5f1",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "adapter_model.safetensors:   0%|          | 0.00/1.18G [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "84cc7a2a3a474bb791d61e2357dd229e",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "events.out.tfevents.1701209373.hf-dgx-01.667111.0:   0%|          | 0.00/8.52k [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "7ce2025dd01647599c00578044512c8c",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "training_args.bin:   0%|          | 0.00/4.79k [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "text/plain": [
+       "CommitInfo(commit_url='https://huggingface.co/smangrul/mistral_lora_clm_with_added_tokens/commit/60ed7ea8bef10ce46d7a64229481dd1ad0e3d1c5', commit_message='Upload model', commit_description='', oid='60ed7ea8bef10ce46d7a64229481dd1ad0e3d1c5', pr_url=None, pr_revision=None, pr_num=None)"
+      ]
+     },
+     "execution_count": 12,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "trainer.push_to_hub()\n",
+    "trainer.model.push_to_hub(training_args.output_dir)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "66812cc4-f9a3-46c4-bcee-0cba03950685",
+   "metadata": {},
+   "source": [
+    "# Check the model loading is working as expected and generating plausible outputs."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "589c46d7-d567-40b4-ab7d-e0a9e1cab40e",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "f98524da95b64a29a9016c6067313b2b",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "aaae3bc0f52f45bbaab60687b71fc4cf",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "adapter_config.json:   0%|          | 0.00/637 [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "1fc5754f41784d1aba00b93551894579",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "adapter_model.safetensors:   0%|          | 0.00/1.18G [00:00<?, ?B/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "context=\"<|begincontext|><|user|>Can you find me a place to eat please?<|system|>Where at? And what kind of cuisine are you craving?<|user|>Somewhere in SF, and I am really craving Thai food at the moment!<|system|>I found a bunch of restaurants, there's actually 10 that you might like in San Francisco, one of them being Baan Thai House & Wine Bar<|user|>How can I reach them? And what's their address?<|system|>You can reach them by phone at 415-379-4505 and visit them at 534 Irving Street<|beginlastuserutterance|>Great, that restaurant sounds good<|endlastuserutterance|><|endcontext|>\" \n",
+      "\n",
+      " target_predicted='<|begintarget|><|begindsts|><|begindst|><|beginintent|> FindRestaurant<|endintent|><|beginbelief|> Restaurants^city->SF~San Francisco|Restaurants^cuisine->Thai|Restaurants^restaurant_name->Baan Thai House & Wine Bar<|endbelief|><|enddst|><|enddsts|><|beginuseraction|> REQUEST->Restaurants^phone_number~|REQUEST->Restaurants^street_address~<|enduseraction|><|beginaction|> INFORM->Restaurants^phone_number~415-379-4505|INFORM->Restaurants^street_address~534 Irving Street<|endaction|><|beginresponse|> The phone number is 415-379-4505 and the address is 534 Irving Street<|endresponse|><|endtarget|>' \n",
+      "\n",
+      " target='<|begintarget|><|begindsts|><|begindst|><|beginintent|>FindRestaurants<|endintent|><|beginbelief|>Restaurants^city->SF~San Francisco|Restaurants^cuisine->Thai|Restaurants^restaurant_name->Baan Thai House & Wine Bar<|endbelief|><|enddst|><|enddsts|><|beginuseraction|>SELECT->Restaurants^~<|enduseraction|><|beginaction|>OFFER_INTENT->Restaurants^intent~ReserveRestaurant<|endaction|><|beginresponse|>Want me to book a table?<|endresponse|><|endtarget|>'\n"
+     ]
+    }
+   ],
+   "source": [
+    "from peft import PeftModel\n",
+    "\n",
+    "inference_model = AutoModelForCausalLM.from_pretrained(\n",
+    "    model_name,\n",
+    "    low_cpu_mem_usage=True,\n",
+    "    # use_flash_attention_2=True,\n",
+    ")\n",
+    "inference_model.resize_token_embeddings(len(tokenizer))\n",
+    "\n",
+    "inference_model = PeftModel.from_pretrained(inference_model, \"smangrul/mistral_lora_clm_with_added_tokens\")\n",
+    "inference_model.to(\"cuda\")\n",
+    "inference_model.eval()\n",
+    "\n",
+    "output_tokens = inference_model.generate(\n",
+    "    **batch,\n",
+    "    max_new_tokens=256,\n",
+    "    do_sample=True,\n",
+    "    temperature=0.2,\n",
+    "    top_p=0.95,\n",
+    "    top_k=50,\n",
+    "    eos_token_id=tokenizer.eos_token_id,\n",
+    "    pad_token_id=tokenizer.pad_token_id,\n",
+    ")\n",
+    "\n",
+    "target_predicted = tokenizer.decode(output_tokens[0], skip_special_tokens=False).split(\"<|endcontext|>\")[1]\n",
+    "print(f\"{context=} \\n\\n {target_predicted=} \\n\\n {target=}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fd57f6e8-761f-4e0b-941c-f6973e13b186",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "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.10.13"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

prompt_tuning_clm.ipynb

@@ -0,0 +1,1229 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "71fbfca2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from transformers import AutoModelForCausalLM\n",
+    "from peft import get_peft_config, get_peft_model, PromptTuningInit, PromptTuningConfig, TaskType, PeftType\n",
+    "import torch\n",
+    "from datasets import load_dataset\n",
+    "import os\n",
+    "from transformers import AutoTokenizer\n",
+    "from torch.utils.data import DataLoader\n",
+    "from transformers import default_data_collator, get_linear_schedule_with_warmup\n",
+    "from tqdm import tqdm\n",
+    "from datasets import load_dataset\n",
+    "\n",
+    "device = \"cuda\"\n",
+    "model_name_or_path = \"bigscience/bloomz-560m\"\n",
+    "tokenizer_name_or_path = \"bigscience/bloomz-560m\"\n",
+    "peft_config = PromptTuningConfig(\n",
+    "    task_type=TaskType.CAUSAL_LM,\n",
+    "    prompt_tuning_init=PromptTuningInit.TEXT,\n",
+    "    num_virtual_tokens=8,\n",
+    "    prompt_tuning_init_text=\"Classify if the tweet is a complaint or not:\",\n",
+    "    tokenizer_name_or_path=model_name_or_path,\n",
+    ")\n",
+    "\n",
+    "dataset_name = \"twitter_complaints\"\n",
+    "checkpoint_name = f\"{dataset_name}_{model_name_or_path}_{peft_config.peft_type}_{peft_config.task_type}_v1.pt\".replace(\n",
+    "    \"/\", \"_\"\n",
+    ")\n",
+    "text_column = \"Tweet text\"\n",
+    "label_column = \"text_label\"\n",
+    "max_length = 64\n",
+    "lr = 3e-2\n",
+    "num_epochs = 50\n",
+    "batch_size = 8"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e1a3648b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from datasets import load_dataset\n",
+    "\n",
+    "dataset = load_dataset(\"ought/raft\", dataset_name)\n",
+    "\n",
+    "classes = [k.replace(\"_\", \" \") for k in dataset[\"train\"].features[\"Label\"].names]\n",
+    "print(classes)\n",
+    "dataset = dataset.map(\n",
+    "    lambda x: {\"text_label\": [classes[label] for label in x[\"Label\"]]},\n",
+    "    batched=True,\n",
+    "    num_proc=1,\n",
+    ")\n",
+    "print(dataset)\n",
+    "dataset[\"train\"][0]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fe12d4d3",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# data preprocessing\n",
+    "tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)\n",
+    "if tokenizer.pad_token_id is None:\n",
+    "    tokenizer.pad_token_id = tokenizer.eos_token_id\n",
+    "target_max_length = max([len(tokenizer(class_label)[\"input_ids\"]) for class_label in classes])\n",
+    "print(target_max_length)\n",
+    "\n",
+    "\n",
+    "def preprocess_function(examples):\n",
+    "    batch_size = len(examples[text_column])\n",
+    "    inputs = [f\"{text_column} : {x} Label : \" for x in examples[text_column]]\n",
+    "    targets = [str(x) for x in examples[label_column]]\n",
+    "    model_inputs = tokenizer(inputs)\n",
+    "    labels = tokenizer(targets, add_special_tokens=False)  # don't add bos token because we concatenate with inputs\n",
+    "    for i in range(batch_size):\n",
+    "        sample_input_ids = model_inputs[\"input_ids\"][i]\n",
+    "        label_input_ids = labels[\"input_ids\"][i] + [tokenizer.eos_token_id]\n",
+    "        # print(i, sample_input_ids, label_input_ids)\n",
+    "        model_inputs[\"input_ids\"][i] = sample_input_ids + label_input_ids\n",
+    "        labels[\"input_ids\"][i] = [-100] * len(sample_input_ids) + label_input_ids\n",
+    "        model_inputs[\"attention_mask\"][i] = [1] * len(model_inputs[\"input_ids\"][i])\n",
+    "    # print(model_inputs)\n",
+    "    for i in range(batch_size):\n",
+    "        sample_input_ids = model_inputs[\"input_ids\"][i]\n",
+    "        label_input_ids = labels[\"input_ids\"][i]\n",
+    "        model_inputs[\"input_ids\"][i] = [tokenizer.pad_token_id] * (\n",
+    "            max_length - len(sample_input_ids)\n",
+    "        ) + sample_input_ids\n",
+    "        model_inputs[\"attention_mask\"][i] = [0] * (max_length - len(sample_input_ids)) + model_inputs[\n",
+    "            \"attention_mask\"\n",
+    "        ][i]\n",
+    "        labels[\"input_ids\"][i] = [-100] * (max_length - len(sample_input_ids)) + label_input_ids\n",
+    "        model_inputs[\"input_ids\"][i] = torch.tensor(model_inputs[\"input_ids\"][i][:max_length])\n",
+    "        model_inputs[\"attention_mask\"][i] = torch.tensor(model_inputs[\"attention_mask\"][i][:max_length])\n",
+    "        labels[\"input_ids\"][i] = torch.tensor(labels[\"input_ids\"][i][:max_length])\n",
+    "    model_inputs[\"labels\"] = labels[\"input_ids\"]\n",
+    "    return model_inputs\n",
+    "\n",
+    "\n",
+    "processed_datasets = dataset.map(\n",
+    "    preprocess_function,\n",
+    "    batched=True,\n",
+    "    num_proc=1,\n",
+    "    remove_columns=dataset[\"train\"].column_names,\n",
+    "    load_from_cache_file=False,\n",
+    "    desc=\"Running tokenizer on dataset\",\n",
+    ")\n",
+    "\n",
+    "train_dataset = processed_datasets[\"train\"]\n",
+    "eval_dataset = processed_datasets[\"train\"]\n",
+    "\n",
+    "\n",
+    "train_dataloader = DataLoader(\n",
+    "    train_dataset, shuffle=True, collate_fn=default_data_collator, batch_size=batch_size, pin_memory=True\n",
+    ")\n",
+    "eval_dataloader = DataLoader(eval_dataset, collate_fn=default_data_collator, batch_size=batch_size, pin_memory=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "641b21fe",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def test_preprocess_function(examples):\n",
+    "    batch_size = len(examples[text_column])\n",
+    "    inputs = [f\"{text_column} : {x} Label : \" for x in examples[text_column]]\n",
+    "    model_inputs = tokenizer(inputs)\n",
+    "    # print(model_inputs)\n",
+    "    for i in range(batch_size):\n",
+    "        sample_input_ids = model_inputs[\"input_ids\"][i]\n",
+    "        model_inputs[\"input_ids\"][i] = [tokenizer.pad_token_id] * (\n",
+    "            max_length - len(sample_input_ids)\n",
+    "        ) + sample_input_ids\n",
+    "        model_inputs[\"attention_mask\"][i] = [0] * (max_length - len(sample_input_ids)) + model_inputs[\n",
+    "            \"attention_mask\"\n",
+    "        ][i]\n",
+    "        model_inputs[\"input_ids\"][i] = torch.tensor(model_inputs[\"input_ids\"][i][:max_length])\n",
+    "        model_inputs[\"attention_mask\"][i] = torch.tensor(model_inputs[\"attention_mask\"][i][:max_length])\n",
+    "    return model_inputs\n",
+    "\n",
+    "\n",
+    "test_dataset = dataset[\"test\"].map(\n",
+    "    test_preprocess_function,\n",
+    "    batched=True,\n",
+    "    num_proc=1,\n",
+    "    remove_columns=dataset[\"train\"].column_names,\n",
+    "    load_from_cache_file=False,\n",
+    "    desc=\"Running tokenizer on dataset\",\n",
+    ")\n",
+    "\n",
+    "test_dataloader = DataLoader(test_dataset, collate_fn=default_data_collator, batch_size=batch_size, pin_memory=True)\n",
+    "next(iter(test_dataloader))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "accc5012",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "next(iter(train_dataloader))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "218df807",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "len(test_dataloader)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "47d1fedf",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "next(iter(test_dataloader))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a773e092",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# creating model\n",
+    "model = AutoModelForCausalLM.from_pretrained(model_name_or_path)\n",
+    "model = get_peft_model(model, peft_config)\n",
+    "model.print_trainable_parameters()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "b2f91568",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# model\n",
+    "# optimizer and lr scheduler\n",
+    "optimizer = torch.optim.AdamW(model.parameters(), lr=lr)\n",
+    "lr_scheduler = get_linear_schedule_with_warmup(\n",
+    "    optimizer=optimizer,\n",
+    "    num_warmup_steps=0,\n",
+    "    num_training_steps=(len(train_dataloader) * num_epochs),\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "id": "e4fb69fc",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:01<00:00,  5.68it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 21.48it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=0: train_ppl=tensor(2.2720e+13, device='cuda:0') train_epoch_loss=tensor(30.7543, device='cuda:0') eval_ppl=tensor(483597.5625, device='cuda:0') eval_epoch_loss=tensor(13.0890, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 10.91it/s]\n",
+      "100%|████████████████��███████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 20.96it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=1: train_ppl=tensor(452658.3750, device='cuda:0') train_epoch_loss=tensor(13.0229, device='cuda:0') eval_ppl=tensor(275088.1875, device='cuda:0') eval_epoch_loss=tensor(12.5248, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 10.90it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 21.41it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=2: train_ppl=tensor(199203.3906, device='cuda:0') train_epoch_loss=tensor(12.2021, device='cuda:0') eval_ppl=tensor(143637.0312, device='cuda:0') eval_epoch_loss=tensor(11.8750, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 10.92it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 21.31it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=3: train_ppl=tensor(114743.9531, device='cuda:0') train_epoch_loss=tensor(11.6505, device='cuda:0') eval_ppl=tensor(54962., device='cuda:0') eval_epoch_loss=tensor(10.9144, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 10.81it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 21.34it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=4: train_ppl=tensor(40786.5977, device='cuda:0') train_epoch_loss=tensor(10.6161, device='cuda:0') eval_ppl=tensor(18342.5430, device='cuda:0') eval_epoch_loss=tensor(9.8170, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 10.89it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 21.34it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=5: train_ppl=tensor(14023.0830, device='cuda:0') train_epoch_loss=tensor(9.5485, device='cuda:0') eval_ppl=tensor(6316.8540, device='cuda:0') eval_epoch_loss=tensor(8.7510, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████████████████████████████████████████████████████████████████████████████████████��█| 7/7 [00:00<00:00, 10.84it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 21.32it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=6: train_ppl=tensor(5635.3262, device='cuda:0') train_epoch_loss=tensor(8.6368, device='cuda:0') eval_ppl=tensor(2476.5776, device='cuda:0') eval_epoch_loss=tensor(7.8146, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 10.88it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 21.30it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=7: train_ppl=tensor(1818.4940, device='cuda:0') train_epoch_loss=tensor(7.5058, device='cuda:0') eval_ppl=tensor(934.1146, device='cuda:0') eval_epoch_loss=tensor(6.8396, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 10.05it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 18.97it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=8: train_ppl=tensor(645.2143, device='cuda:0') train_epoch_loss=tensor(6.4696, device='cuda:0') eval_ppl=tensor(361.9093, device='cuda:0') eval_epoch_loss=tensor(5.8914, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00,  9.67it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 19.12it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=9: train_ppl=tensor(293.8047, device='cuda:0') train_epoch_loss=tensor(5.6829, device='cuda:0') eval_ppl=tensor(215.8185, device='cuda:0') eval_epoch_loss=tensor(5.3744, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 10.54it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 20.83it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=10: train_ppl=tensor(191.2377, device='cuda:0') train_epoch_loss=tensor(5.2535, device='cuda:0') eval_ppl=tensor(177.1512, device='cuda:0') eval_epoch_loss=tensor(5.1770, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████���███████████████████████████| 7/7 [00:00<00:00, 10.02it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 18.98it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=11: train_ppl=tensor(153.6052, device='cuda:0') train_epoch_loss=tensor(5.0344, device='cuda:0') eval_ppl=tensor(126.6154, device='cuda:0') eval_epoch_loss=tensor(4.8412, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00,  9.54it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 18.78it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=12: train_ppl=tensor(122.8925, device='cuda:0') train_epoch_loss=tensor(4.8113, device='cuda:0') eval_ppl=tensor(97.3331, device='cuda:0') eval_epoch_loss=tensor(4.5781, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00,  9.66it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 19.72it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=13: train_ppl=tensor(84.8845, device='cuda:0') train_epoch_loss=tensor(4.4413, device='cuda:0') eval_ppl=tensor(70.3213, device='cuda:0') eval_epoch_loss=tensor(4.2531, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:01<00:00,  6.73it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 16.07it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=14: train_ppl=tensor(64.6705, device='cuda:0') train_epoch_loss=tensor(4.1693, device='cuda:0') eval_ppl=tensor(50.4688, device='cuda:0') eval_epoch_loss=tensor(3.9214, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00,  8.41it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 15.63it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=15: train_ppl=tensor(44.2937, device='cuda:0') train_epoch_loss=tensor(3.7908, device='cuda:0') eval_ppl=tensor(34.8210, device='cuda:0') eval_epoch_loss=tensor(3.5502, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00,  8.31it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 15.67it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=16: train_ppl=tensor(30.0995, device='cuda:0') train_epoch_loss=tensor(3.4045, device='cuda:0') eval_ppl=tensor(24.7703, device='cuda:0') eval_epoch_loss=tensor(3.2096, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00,  8.31it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 15.59it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=17: train_ppl=tensor(23.3086, device='cuda:0') train_epoch_loss=tensor(3.1488, device='cuda:0') eval_ppl=tensor(20.8131, device='cuda:0') eval_epoch_loss=tensor(3.0356, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00,  8.29it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 16.04it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=18: train_ppl=tensor(16.4479, device='cuda:0') train_epoch_loss=tensor(2.8002, device='cuda:0') eval_ppl=tensor(12.0876, device='cuda:0') eval_epoch_loss=tensor(2.4922, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00,  8.37it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 15.37it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=19: train_ppl=tensor(11.1977, device='cuda:0') train_epoch_loss=tensor(2.4157, device='cuda:0') eval_ppl=tensor(9.0399, device='cuda:0') eval_epoch_loss=tensor(2.2016, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00,  8.23it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 17.29it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=20: train_ppl=tensor(8.1847, device='cuda:0') train_epoch_loss=tensor(2.1023, device='cuda:0') eval_ppl=tensor(6.7486, device='cuda:0') eval_epoch_loss=tensor(1.9093, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|█████████████████��██████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00,  8.30it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 15.58it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=21: train_ppl=tensor(6.1145, device='cuda:0') train_epoch_loss=tensor(1.8107, device='cuda:0') eval_ppl=tensor(5.5931, device='cuda:0') eval_epoch_loss=tensor(1.7215, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00,  8.34it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 21.36it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=22: train_ppl=tensor(5.2963, device='cuda:0') train_epoch_loss=tensor(1.6670, device='cuda:0') eval_ppl=tensor(5.0573, device='cuda:0') eval_epoch_loss=tensor(1.6208, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 10.84it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 21.26it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=23: train_ppl=tensor(4.7485, device='cuda:0') train_epoch_loss=tensor(1.5578, device='cuda:0') eval_ppl=tensor(3.6277, device='cuda:0') eval_epoch_loss=tensor(1.2886, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 10.84it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 21.31it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=24: train_ppl=tensor(3.4080, device='cuda:0') train_epoch_loss=tensor(1.2261, device='cuda:0') eval_ppl=tensor(3.0467, device='cuda:0') eval_epoch_loss=tensor(1.1141, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 10.88it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 21.25it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=25: train_ppl=tensor(3.3052, device='cuda:0') train_epoch_loss=tensor(1.1955, device='cuda:0') eval_ppl=tensor(2.7784, device='cuda:0') eval_epoch_loss=tensor(1.0219, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 10.86it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 21.22it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=26: train_ppl=tensor(2.9487, device='cuda:0') train_epoch_loss=tensor(1.0814, device='cuda:0') eval_ppl=tensor(2.9471, device='cuda:0') eval_epoch_loss=tensor(1.0808, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 10.85it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 21.25it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=27: train_ppl=tensor(2.8738, device='cuda:0') train_epoch_loss=tensor(1.0556, device='cuda:0') eval_ppl=tensor(2.5801, device='cuda:0') eval_epoch_loss=tensor(0.9478, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 10.84it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 21.28it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=28: train_ppl=tensor(2.3241, device='cuda:0') train_epoch_loss=tensor(0.8433, device='cuda:0') eval_ppl=tensor(2.2198, device='cuda:0') eval_epoch_loss=tensor(0.7974, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 10.84it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 20.89it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=29: train_ppl=tensor(2.0376, device='cuda:0') train_epoch_loss=tensor(0.7118, device='cuda:0') eval_ppl=tensor(1.8572, device='cuda:0') eval_epoch_loss=tensor(0.6191, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00,  9.76it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 18.83it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=30: train_ppl=tensor(1.8301, device='cuda:0') train_epoch_loss=tensor(0.6044, device='cuda:0') eval_ppl=tensor(1.8864, device='cuda:0') eval_epoch_loss=tensor(0.6347, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00,  9.80it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 19.81it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=31: train_ppl=tensor(1.7301, device='cuda:0') train_epoch_loss=tensor(0.5482, device='cuda:0') eval_ppl=tensor(1.6340, device='cuda:0') eval_epoch_loss=tensor(0.4910, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 10.60it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 19.11it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=32: train_ppl=tensor(1.5842, device='cuda:0') train_epoch_loss=tensor(0.4601, device='cuda:0') eval_ppl=tensor(1.6179, device='cuda:0') eval_epoch_loss=tensor(0.4811, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 10.11it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 18.35it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=33: train_ppl=tensor(1.5193, device='cuda:0') train_epoch_loss=tensor(0.4183, device='cuda:0') eval_ppl=tensor(1.5543, device='cuda:0') eval_epoch_loss=tensor(0.4410, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00,  9.59it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 18.60it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=34: train_ppl=tensor(1.5402, device='cuda:0') train_epoch_loss=tensor(0.4319, device='cuda:0') eval_ppl=tensor(1.4924, device='cuda:0') eval_epoch_loss=tensor(0.4004, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00,  9.80it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 19.63it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=35: train_ppl=tensor(1.4410, device='cuda:0') train_epoch_loss=tensor(0.3654, device='cuda:0') eval_ppl=tensor(1.3888, device='cuda:0') eval_epoch_loss=tensor(0.3284, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:01<00:00,  6.60it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 15.36it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=36: train_ppl=tensor(1.3675, device='cuda:0') train_epoch_loss=tensor(0.3130, device='cuda:0') eval_ppl=tensor(1.4001, device='cuda:0') eval_epoch_loss=tensor(0.3366, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00,  8.40it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 15.58it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=37: train_ppl=tensor(1.4197, device='cuda:0') train_epoch_loss=tensor(0.3505, device='cuda:0') eval_ppl=tensor(1.3214, device='cuda:0') eval_epoch_loss=tensor(0.2787, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00,  8.27it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 15.56it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=38: train_ppl=tensor(1.3855, device='cuda:0') train_epoch_loss=tensor(0.3261, device='cuda:0') eval_ppl=tensor(1.3501, device='cuda:0') eval_epoch_loss=tensor(0.3001, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00,  8.25it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 15.57it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=39: train_ppl=tensor(1.3643, device='cuda:0') train_epoch_loss=tensor(0.3107, device='cuda:0') eval_ppl=tensor(1.3549, device='cuda:0') eval_epoch_loss=tensor(0.3037, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00,  8.28it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 15.41it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=40: train_ppl=tensor(1.3093, device='cuda:0') train_epoch_loss=tensor(0.2695, device='cuda:0') eval_ppl=tensor(1.3233, device='cuda:0') eval_epoch_loss=tensor(0.2801, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00,  8.24it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 15.51it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=41: train_ppl=tensor(1.3108, device='cuda:0') train_epoch_loss=tensor(0.2706, device='cuda:0') eval_ppl=tensor(1.3440, device='cuda:0') eval_epoch_loss=tensor(0.2957, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00,  8.78it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 15.61it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=42: train_ppl=tensor(1.2944, device='cuda:0') train_epoch_loss=tensor(0.2581, device='cuda:0') eval_ppl=tensor(1.2711, device='cuda:0') eval_epoch_loss=tensor(0.2399, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00,  8.29it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 15.56it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=43: train_ppl=tensor(1.2616, device='cuda:0') train_epoch_loss=tensor(0.2323, device='cuda:0') eval_ppl=tensor(1.2449, device='cuda:0') eval_epoch_loss=tensor(0.2190, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 10.85it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 21.27it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=44: train_ppl=tensor(1.2478, device='cuda:0') train_epoch_loss=tensor(0.2214, device='cuda:0') eval_ppl=tensor(1.2202, device='cuda:0') eval_epoch_loss=tensor(0.1990, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 10.85it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 21.31it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=45: train_ppl=tensor(1.2350, device='cuda:0') train_epoch_loss=tensor(0.2111, device='cuda:0') eval_ppl=tensor(1.2180, device='cuda:0') eval_epoch_loss=tensor(0.1972, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 10.86it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 21.33it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=46: train_ppl=tensor(1.2277, device='cuda:0') train_epoch_loss=tensor(0.2052, device='cuda:0') eval_ppl=tensor(1.2077, device='cuda:0') eval_epoch_loss=tensor(0.1887, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 10.87it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 21.35it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=47: train_ppl=tensor(1.2037, device='cuda:0') train_epoch_loss=tensor(0.1854, device='cuda:0') eval_ppl=tensor(1.2041, device='cuda:0') eval_epoch_loss=tensor(0.1857, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 10.83it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 21.29it/s]\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=48: train_ppl=tensor(1.2026, device='cuda:0') train_epoch_loss=tensor(0.1845, device='cuda:0') eval_ppl=tensor(1.1982, device='cuda:0') eval_epoch_loss=tensor(0.1808, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 10.86it/s]\n",
+      "100%|████████████████████████████████████████████████████████████████████████████████████████████| 7/7 [00:00<00:00, 21.35it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "epoch=49: train_ppl=tensor(1.2005, device='cuda:0') train_epoch_loss=tensor(0.1827, device='cuda:0') eval_ppl=tensor(1.1968, device='cuda:0') eval_epoch_loss=tensor(0.1796, device='cuda:0')\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "# training and evaluation\n",
+    "model = model.to(device)\n",
+    "\n",
+    "for epoch in range(num_epochs):\n",
+    "    model.train()\n",
+    "    total_loss = 0\n",
+    "    for step, batch in enumerate(tqdm(train_dataloader)):\n",
+    "        batch = {k: v.to(device) for k, v in batch.items()}\n",
+    "        #         print(batch)\n",
+    "        #         print(batch[\"input_ids\"].shape)\n",
+    "        outputs = model(**batch)\n",
+    "        loss = outputs.loss\n",
+    "        total_loss += loss.detach().float()\n",
+    "        loss.backward()\n",
+    "        optimizer.step()\n",
+    "        lr_scheduler.step()\n",
+    "        optimizer.zero_grad()\n",
+    "\n",
+    "    model.eval()\n",
+    "    eval_loss = 0\n",
+    "    eval_preds = []\n",
+    "    for step, batch in enumerate(tqdm(eval_dataloader)):\n",
+    "        batch = {k: v.to(device) for k, v in batch.items()}\n",
+    "        with torch.no_grad():\n",
+    "            outputs = model(**batch)\n",
+    "        loss = outputs.loss\n",
+    "        eval_loss += loss.detach().float()\n",
+    "        eval_preds.extend(\n",
+    "            tokenizer.batch_decode(torch.argmax(outputs.logits, -1).detach().cpu().numpy(), skip_special_tokens=True)\n",
+    "        )\n",
+    "\n",
+    "    eval_epoch_loss = eval_loss / len(eval_dataloader)\n",
+    "    eval_ppl = torch.exp(eval_epoch_loss)\n",
+    "    train_epoch_loss = total_loss / len(train_dataloader)\n",
+    "    train_ppl = torch.exp(train_epoch_loss)\n",
+    "    print(f\"{epoch=}: {train_ppl=} {train_epoch_loss=} {eval_ppl=} {eval_epoch_loss=}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 29,
+   "id": "53752a7b",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "@TommyHilfiger Dramatic shopping exp. ordered 6 jeans same size (30/32) 2 fits / 2 too large / 2 too slim : same brand &gt; different sizing\n",
+      "{'input_ids': tensor([[227985,   5484,    915,   2566, 226154, 126015,   5385,    259, 239364,\n",
+      "           3396,  70823,   5853,     17,  57247,   1231, 191040,   5025,   7869,\n",
+      "            375,   2324, 149349,     12,    415, 122321,    897,    415,  10136,\n",
+      "          10021,    897,    415,  10136,   6497,    381,    915,   5025,  51950,\n",
+      "          66869,   5955,    272,  20311,  77658,    915,    210]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,\n",
+      "         1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])}\n",
+      "tensor([[227985,   5484,    915,   2566, 226154, 126015,   5385,    259, 239364,\n",
+      "           3396,  70823,   5853,     17,  57247,   1231, 191040,   5025,   7869,\n",
+      "            375,   2324, 149349,     12,    415, 122321,    897,    415,  10136,\n",
+      "          10021,    897,    415,  10136,   6497,    381,    915,   5025,  51950,\n",
+      "          66869,   5955,    272,  20311,  77658,    915,    210,  16449,   5952,\n",
+      "              3]], device='cuda:0')\n",
+      "['Tweet text : @TommyHilfiger Dramatic shopping exp. ordered 6 jeans same size (30/32) 2 fits / 2 too large / 2 too slim : same brand &gt; different sizing Label : complaint']\n"
+     ]
+    }
+   ],
+   "source": [
+    "model.eval()\n",
+    "i = 33\n",
+    "inputs = tokenizer(f'{text_column} : {dataset[\"test\"][i][\"Tweet text\"]} Label : ', return_tensors=\"pt\")\n",
+    "print(dataset[\"test\"][i][\"Tweet text\"])\n",
+    "print(inputs)\n",
+    "\n",
+    "with torch.no_grad():\n",
+    "    inputs = {k: v.to(device) for k, v in inputs.items()}\n",
+    "    outputs = model.generate(\n",
+    "        input_ids=inputs[\"input_ids\"], attention_mask=inputs[\"attention_mask\"], max_new_tokens=10, eos_token_id=3\n",
+    "    )\n",
+    "    print(outputs)\n",
+    "    print(tokenizer.batch_decode(outputs.detach().cpu().numpy(), skip_special_tokens=True))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c8f35152",
+   "metadata": {},
+   "source": [
+    "You can push model to hub or save model locally. \n",
+    "\n",
+    "- Option1: Pushing the model to Hugging Face Hub\n",
+    "```python\n",
+    "model.push_to_hub(\n",
+    "    f\"{dataset_name}_{model_name_or_path}_{peft_config.peft_type}_{peft_config.task_type}\".replace(\"/\", \"_\"),\n",
+    "    token = \"hf_...\"\n",
+    ")\n",
+    "```\n",
+    "token (`bool` or `str`, *optional*):\n",
+    "    `token` is to be used for HTTP Bearer authorization when accessing remote files. If `True`, will use the token generated\n",
+    "    when running `huggingface-cli login` (stored in `~/.huggingface`). Will default to `True` if `repo_url`\n",
+    "    is not specified.\n",
+    "    Or you can get your token from https://huggingface.co/settings/token\n",
+    "```\n",
+    "- Or save model locally\n",
+    "```python\n",
+    "peft_model_id = f\"{dataset_name}_{model_name_or_path}_{peft_config.peft_type}_{peft_config.task_type}\".replace(\"/\", \"_\")\n",
+    "model.save_pretrained(peft_model_id)\n",
+    "```"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "id": "d8ba1f8c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# saving model\n",
+    "peft_model_id = f\"{dataset_name}_{model_name_or_path}_{peft_config.peft_type}_{peft_config.task_type}\".replace(\n",
+    "    \"/\", \"_\"\n",
+    ")\n",
+    "model.save_pretrained(peft_model_id)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "id": "4928c7f1",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
+      "To disable this warning, you can either:\n",
+      "\t- Avoid using `tokenizers` before the fork if possible\n",
+      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n",
+      "36K\tbigscience/bloomz-560m_PROMPT_TUNING_CAUSAL_LM/adapter_model.bin\n"
+     ]
+    }
+   ],
+   "source": [
+    "ckpt = f\"{peft_model_id}/adapter_model.bin\"\n",
+    "!du -h $ckpt"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "id": "4d9476e1",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from peft import PeftModel, PeftConfig\n",
+    "\n",
+    "peft_model_id = f\"{dataset_name}_{model_name_or_path}_{peft_config.peft_type}_{peft_config.task_type}\".replace(\n",
+    "    \"/\", \"_\"\n",
+    ")\n",
+    "\n",
+    "config = PeftConfig.from_pretrained(peft_model_id)\n",
+    "model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path)\n",
+    "model = PeftModel.from_pretrained(model, peft_model_id)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 33,
+   "id": "ebe174a6",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "@greateranglia Ok thanks...\n",
+      "{'input_ids': tensor([[227985,   5484,    915,   2566,  14173,   2960,  29906,    387,  20706,\n",
+      "          49337,   1369,  77658,    915,    210]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])}\n",
+      "tensor([[227985,   5484,    915,   2566,  14173,   2960,  29906,    387,  20706,\n",
+      "          49337,   1369,  77658,    915,    210,   1936, 106863,      3]],\n",
+      "       device='cuda:0')\n",
+      "['Tweet text : @greateranglia Ok thanks... Label : no complaint']\n"
+     ]
+    }
+   ],
+   "source": [
+    "model.to(device)\n",
+    "model.eval()\n",
+    "i = 4\n",
+    "inputs = tokenizer(f'{text_column} : {dataset[\"test\"][i][\"Tweet text\"]} Label : ', return_tensors=\"pt\")\n",
+    "print(dataset[\"test\"][i][\"Tweet text\"])\n",
+    "print(inputs)\n",
+    "\n",
+    "with torch.no_grad():\n",
+    "    inputs = {k: v.to(device) for k, v in inputs.items()}\n",
+    "    outputs = model.generate(\n",
+    "        input_ids=inputs[\"input_ids\"], attention_mask=inputs[\"attention_mask\"], max_new_tokens=10, eos_token_id=3\n",
+    "    )\n",
+    "    print(outputs)\n",
+    "    print(tokenizer.batch_decode(outputs.detach().cpu().numpy(), skip_special_tokens=True))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "24041ee1",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "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.10.5"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "aee8b7b246df8f9039afb4144a1f6fd8d2ca17a180786b69acc140d282b71a49"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}