finetune.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "fb92b09b-8114-4e13-9b92-3d3875cadb6e",
   "metadata": {},
   "source": [
    "# Fine-tune\n",
    "\n",
    "As an extra exercise, we're going to finetune the model to perform better on the type of conversations we're interested in. To do so, we're using [LoRA](https://huggingface.co/docs/peft/conceptual_guides/lora) technique. The goal of LoRA is to fine-tune just the most important parts of the model by injecting new trainable weights to specific layers of the transformer architecture, while keeping the original model checkpoint frozen.\n",
    "\n",
    "> 📝 **Note:** this notebook is based on [this post by Geronimo](https://medium.com/@geronimo7/phinetuning-2-0-28a2be6de110), the code was taking and adapted from their explanations."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a56e294-b3a8-40af-bddf-d41fd28897bf",
   "metadata": {},
   "source": [
    "## Create an artificial dataset\n",
    "\n",
    "In order to fine-tune our model, we first need a dataset to fine-tune it on. We don't have real customer-agent conversations, se we're going to create artificial data instead. Ours, is a tiny (only 4 samples) and toy dataset created using [Google Bard](https://bard.google.com/chat) asking it to give use some sample conversations between a `customer` and an `agent` on imaginary issues (samples were generated only for the `support` class).\n",
    "\n",
    "> 📝 **Note:** this dataset was created for the task of text/response generation. Similarly, we could create another dataset to train the model for the classification task (`support`, `sales` and `joke`).\n",
    "\n",
    "The result is a tiny dataset of the form (here `content` was truncated to just the first sentence):\n",
    "\n",
    "```json\n",
    "{ \"data\": [\n",
    "    { \"conversation\" : [\n",
    "      { \"role\": \"customer\", \"content\": \"Having trouble tracking package\" },\n",
    "      { \"role\": \"agent\", \"content\": \"Apologize for inconvenience, look into it\" },\n",
    "      { \"role\": \"customer\", \"content\": \"Need package for important event\" },\n",
    "      { \"role\": \"agent\", \"content\": \"Understand concern, package shipped on [date], in transit\" },\n",
    "      { \"role\": \"customer\", \"content\": \"Thank you for help\" },\n",
    "      { \"role\": \"agent\", \"content\": \"You're welcome, let me know if you have other questions\" }\n",
    "    ]},\n",
    "    { \"conversation\" : [\n",
    "      { \"role\": \"customer\", \"content\": \"Worry package, shipped on [date], haven't received it\" },\n",
    "      { \"role\": \"agent\", \"content\": \"Apologize for delay, investigate further\" },\n",
    "      { \"role\": \"customer\", \"content\": \"Thank you for looking into it\" },\n",
    "      { \"role\": \"agent\", \"content\": \"Understand concern, file missing package report, update you\" },\n",
    "      { \"role\": \"customer\", \"content\": \"Thank you, relieved you're taking it seriously\" },\n",
    "      { \"role\": \"agent\", \"content\": \"You're welcome, let me know if you have other questions\" }\n",
    "    ]}\n",
    "  ]\n",
    "}\n",
    "```\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "1082a6ba-4ab6-4f40-8c58-cf619e3f459c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/grego/code/jobsearch/ipglobal-2/env/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "dataset = load_dataset(\"json\", data_files=\"finetune_dataset.json\", field=\"data\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a400a02c-903f-4cd4-88ef-b7eb798e9112",
   "metadata": {},
   "source": [
    "Let's print a sample sentence."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4c040dba-7e83-4580-ae5f-4bf21d738216",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'messages': [{'content': \"I'm having trouble tracking my package. It was supposed to arrive yesterday, but I don't have any tracking information.\",\n",
       "   'role': 'customer'},\n",
       "  {'content': 'I apologize for the inconvenience, Mr. Smith. Let me look into this for you. Please provide me with your order number, AB0002345.',\n",
       "   'role': 'agent'},\n",
       "  {'content': \"Thank you. I'm concerned because I need this package for an important event.\",\n",
       "   'role': 'customer'},\n",
       "  {'content': \"I understand your concern. According to our records, your package was shipped on Jan 4th and is currently in transit. The expected delivery date is Jan 1st. I'll keep an eye on the tracking information for you and let you know if there are any updates.\",\n",
       "   'role': 'agent'},\n",
       "  {'content': 'Thank you for your help. I appreciate it.', 'role': 'customer'},\n",
       "  {'content': \"You're welcome, Mr. Smith. Please let me know if you have any other questions.\",\n",
       "   'role': 'agent'}]}"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset['train'][0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18333582-6afb-46f4-9aaf-a2e432313a8f",
   "metadata": {},
   "source": [
    "Now create a test split. As it only has 4 samples, we'll randomly select one of them as test."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "96c9ff79-a9f0-4c3c-ae4a-e99460897389",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = dataset[\"train\"].train_test_split(test_size=0.25)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f75833c5-dca7-4a42-a766-4c594e9d2814",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['messages'],\n",
       "        num_rows: 3\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['messages'],\n",
       "        num_rows: 1\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44a19998-9618-48c9-89f8-44203c9df347",
   "metadata": {},
   "source": [
    "## Finetuning the model\n",
    "\n",
    "Now we can start finetuning the model. The steps are:\n",
    "\n",
    "+ Load the model\n",
    "+ Load and adapt the tokenizer\n",
    "+ Create the LoRA configuration and addapt the model\n",
    "+ Tokenize the dataset\n",
    "+ Define the collate function\n",
    "+ Finally, trainig the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "1d6e7c81-8917-4ddf-868b-1efbf0ac7df6",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading checkpoint shards: 100%|██████████████████| 2/2 [00:02<00:00,  1.03s/it]\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from transformers import AutoModelForCausalLM, BitsAndBytesConfig\n",
    "\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    \"microsoft/phi-2\",\n",
    "    device_map=\"auto\",\n",
    "    quantization_config=BitsAndBytesConfig(\n",
    "        load_in_4bit=True,\n",
    "        bnb_4bit_compute_dtype=torch.bfloat16,\n",
    "        bnb_4bit_quant_type=\"nf4\",\n",
    "    ),\n",
    "    torch_dtype=torch.bfloat16,\n",
    "    trust_remote_code=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38caf62a-45b5-42dd-b59c-91a4a806052c",
   "metadata": {},
   "source": [
    "Load tokenizer and add `<PAD>` token. In the original code, extra tokens were added (`<|im_start|>` and `<|im_end|>`), this is because they used [ChatML](https://cobusgreyling.medium.com/the-introduction-of-chat-markup-language-chatml-is-important-for-a-number-of-reasons-5061f6fe2a85) format. We're going to stick the simpler [chat format](https://huggingface.co/microsoft/phi-2#chat-format) used by `microsoft/phi-2` so we only need the `<PAD>` token.\n",
    "\n",
    "```text\n",
    "Alice: I don't know why, I'm struggling to maintain focus while studying. Any suggestions?\n",
    "Bob: Well, have you tried creating a study schedule and sticking to it?\n",
    "Alice: Yes, I have, but it doesn't seem to help much.\n",
    "Bob: Hmm, maybe you should try studying in a quiet environment, like the library.\n",
    "Alice: ...\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "9aa7c593-0f59-443d-bad2-4a6086bdd1d0",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Embedding(50304, 2560)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers import AutoTokenizer\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"microsoft/phi-2\", use_fast=False)   \n",
    "tokenizer.add_tokens([\"<PAD>\"])\n",
    "tokenizer.pad_token = \"<PAD>\"\n",
    "\n",
    "model.resize_token_embeddings(\n",
    "    new_num_tokens=len(tokenizer),\n",
    "    pad_to_multiple_of=64)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c947e0b4-7b0b-4871-980d-89388b03c70c",
   "metadata": {},
   "source": [
    "We create the LoRA adapters and add them to the model. Hopefully, we only need to define the `LoraConfig`, and most of the job is done by the [PEFT](https://huggingface.co/docs/peft/v0.7.1/en/conceptual_guides/lora#lora) library."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "9b18a27f-10bf-4939-9134-840aac382038",
   "metadata": {},
   "outputs": [],
   "source": [
    "from peft import prepare_model_for_kbit_training, LoraConfig, get_peft_model\n",
    "\n",
    "lora_config = LoraConfig(\n",
    "    r=32,\n",
    "    lora_alpha=32,\n",
    "    target_modules = ['Wqkv','out_proj'],\n",
    "    lora_dropout=0.1,\n",
    "    bias=\"none\",\n",
    "    modules_to_save = [\"lm_head\", \"embed_tokens\"],\n",
    "    task_type=\"CAUSAL_LM\"\n",
    ")\n",
    "\n",
    "model = prepare_model_for_kbit_training(model, use_gradient_checkpointing = False)\n",
    "model = get_peft_model(model, lora_config)\n",
    "model.config.use_cache = False"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ec50ab9-5819-48a2-a44d-5e190879281f",
   "metadata": {},
   "source": [
    "Now, we can tokenize our dataset with `dataset.map()`. In addition to `input_ids` and `attention_mask`, we use `labels` as well to tell the model what's the expected output. We only want the model to learn the `agent` messages, ignoring the `customer` ones; so we label them with the `IGNORE_INDEX` id."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "c7025b16-ae72-4896-9863-6d63462d48f7",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "num_proc must be <= 3. Reducing num_proc to 3 for dataset of size 3.\n",
      "Map (num_proc=3): 100%|████████████████████| 3/3 [00:00<00:00, 20.67 examples/s]\n",
      "num_proc must be <= 1. Reducing num_proc to 1 for dataset of size 1.\n",
      "Map: 100%|████████████████████████████████| 1/1 [00:00<00:00, 162.97 examples/s]\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "from functools import partial\n",
    "\n",
    "\n",
    "IGNORE_INDEX = -100\n",
    "template = \"{role}: {content}\\n\"\n",
    "\n",
    "def tokenize(input, max_length):\n",
    "    input_ids, attention_mask, labels = [], [], []\n",
    "\n",
    "    for i, msg in enumerate(input[\"messages\"]):\n",
    "        chat_msg = template.format(**msg)\n",
    "        msg_tokenized = tokenizer(chat_msg, truncation=False, add_special_tokens=False)\n",
    "\n",
    "        input_ids += msg_tokenized[\"input_ids\"]\n",
    "        attention_mask += msg_tokenized[\"attention_mask\"]\n",
    "        labels += [IGNORE_INDEX]*len(msg_tokenized[\"input_ids\"]) \\\n",
    "                    if msg[\"role\"] == \"customer\" \\\n",
    "                    else msg_tokenized[\"input_ids\"]\n",
    "    return {\n",
    "        \"input_ids\": input_ids[:max_length], \n",
    "        \"attention_mask\": attention_mask[:max_length],\n",
    "        \"labels\": labels[:max_length],\n",
    "    }\n",
    "\n",
    "dataset_tokenized = dataset.map(\n",
    "    partial(tokenize, max_length=1024), \n",
    "    batched = False,\n",
    "    num_proc = os.cpu_count(),\n",
    "    remove_columns = dataset[\"train\"].column_names\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2924f03f-2632-4495-bc4f-27fd3cb6411a",
   "metadata": {},
   "source": [
    "We define the collate function. It is in charge of putting samples together into batches. It also pads the inputs so they all have the same length."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "72f7aa1f-5a59-4c30-9c66-46058ac37e34",
   "metadata": {},
   "outputs": [],
   "source": [
    "def collate(elements):\n",
    "    tokens = [e[\"input_ids\"] for e in elements]\n",
    "    tokens_maxlen = max([len(t) for t in tokens])\n",
    "\n",
    "    for i, sample in enumerate(elements):\n",
    "        input_ids = sample[\"input_ids\"]\n",
    "        labels = sample[\"labels\"]\n",
    "        attention_mask = sample[\"attention_mask\"]\n",
    "\n",
    "        pad_len = tokens_maxlen-len(input_ids)\n",
    "        input_ids.extend( pad_len * [tokenizer.pad_token_id] )\n",
    "        labels.extend( pad_len * [IGNORE_INDEX] )\n",
    "        attention_mask.extend( pad_len * [0] )\n",
    "    batch={\n",
    "        \"input_ids\": torch.tensor( [e[\"input_ids\"] for e in elements] ),\n",
    "        \"labels\": torch.tensor( [e[\"labels\"] for e in elements] ),\n",
    "        \"attention_mask\": torch.tensor( [e[\"attention_mask\"] for e in elements] ),\n",
    "    }\n",
    "    return batch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ab2d731-9c89-4f39-bab5-5e723f2520e7",
   "metadata": {},
   "source": [
    "## Train\n",
    "\n",
    "Finally we can train the model. As this is such a simple dataset, the training process ends very quickly. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "194e09e3-702e-42fb-9e3b-a35d37ddae16",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='5' max='5' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [5/5 00:07, Epoch 5/5]\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",
       "      <th>Validation Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>0.273800</td>\n",
       "      <td>1.441532</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>0.271300</td>\n",
       "      <td>1.420660</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3</td>\n",
       "      <td>0.269900</td>\n",
       "      <td>1.416498</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4</td>\n",
       "      <td>0.267500</td>\n",
       "      <td>1.397182</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5</td>\n",
       "      <td>0.251100</td>\n",
       "      <td>1.390604</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=5, training_loss=0.26674823760986327, metrics={'train_runtime': 8.7028, 'train_samples_per_second': 1.724, 'train_steps_per_second': 0.575, 'total_flos': 45798332398080.0, 'train_loss': 0.26674823760986327, 'epoch': 5.0})"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers import TrainingArguments, Trainer\n",
    "\n",
    "bs=1\n",
    "ga_steps=16  # gradient acc. steps\n",
    "epochs=5\n",
    "lr=0.00002\n",
    "\n",
    "steps_per_epoch=len(dataset_tokenized[\"train\"])//(bs*ga_steps)\n",
    "\n",
    "args = TrainingArguments(\n",
    "    output_dir=\"out\",\n",
    "    per_device_train_batch_size=bs,\n",
    "    per_device_eval_batch_size=16,\n",
    "    evaluation_strategy=\"steps\",\n",
    "    logging_steps=1,\n",
    "    eval_steps=steps_per_epoch,\n",
    "    save_steps=steps_per_epoch,\n",
    "    gradient_accumulation_steps=ga_steps,\n",
    "    num_train_epochs=epochs,\n",
    "    lr_scheduler_type=\"constant\",\n",
    "    optim=\"paged_adamw_32bit\",\n",
    "    learning_rate=lr,\n",
    "    group_by_length=False,\n",
    "    bf16=True,\n",
    "    ddp_find_unused_parameters=False,\n",
    ")\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    tokenizer=tokenizer,\n",
    "    args=args,\n",
    "    data_collator=collate,\n",
    "    train_dataset=dataset_tokenized[\"train\"],\n",
    "    eval_dataset=dataset_tokenized[\"test\"],\n",
    ")\n",
    "\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2be75e6-add7-4aea-9877-0b8bc1bc816c",
   "metadata": {},
   "source": [
    "## Discussion\n",
    "\n",
    "However simple, this toy example contains all the code parts needed to fine-tune the `microsoft/phi-2` model with a custom dataset using the [LoRA](https://huggingface.co/docs/peft/conceptual_guides/lora) technique. For a useful use case, a real dataset should be used with actual chat data. But this script can be easily adapted to train on any chat data, provided it's in the right format (see [section above](#create-an-artificial-dataset))."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python3 Ipglobal 2",
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}