Codes/01 - Fine-tuning.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9cfae24c-01e5-480b-ad5e-2aec772e0983",
   "metadata": {},
   "source": [
    "### Checking for availability of GPUs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8229d250-fd3c-4ee6-ba55-dcfbeba9b06e",
   "metadata": {},
   "outputs": [],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b9081a3-e5bb-45c5-a88b-4ef4e0490574",
   "metadata": {},
   "source": [
    "### Installing the necessary libraries\n",
    "Please note that this installation step might fail if no GPU is available."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9c4d533-ec98-4b0f-a263-c3c8f649547e",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install --upgrade pip\n",
    "!pip install -q -U bitsandbytes\n",
    "!pip install -q -U git+https://github.com/huggingface/transformers.git\n",
    "!pip install -q -U git+https://github.com/huggingface/peft.git\n",
    "!pip install -q -U git+https://github.com/huggingface/accelerate.git\n",
    "!pip install -q trl xformers wandb datasets einops gradio sentencepiece\n",
    "!pip install flash-attn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3425af9b-12cd-46e5-80f5-30fd2f7fed67",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Libraries used for analyzing the dataset\n",
    "!pip install matplotlib\n",
    "!pip install statistics"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0ddd32e-2c21-4abd-8e5c-db85efb194ab",
   "metadata": {},
   "source": [
    "### Importing the necessary libraries and preparation\n",
    "Here, we set the repos of the base model (to be fine-tuned), the dataset (on which the base model will be fine-tuned) and the output model (the base model fine-tuned on our custom dataset)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1eb19e95-6c02-45a5-99a3-5bb966ed2d0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import (\n",
    "    AutoModelForCausalLM, \n",
    "    AutoTokenizer, \n",
    "    BitsAndBytesConfig,\n",
    "    HfArgumentParser,\n",
    "    TrainingArguments,\n",
    "    pipeline, \n",
    "    logging, \n",
    "    TextStreamer,\n",
    "    GenerationConfig\n",
    ")\n",
    "from peft import (\n",
    "    LoraConfig, \n",
    "    PeftModel, \n",
    "    prepare_model_for_kbit_training, \n",
    "    get_peft_model\n",
    ")\n",
    "import os\n",
    "import torch\n",
    "import wandb\n",
    "import platform\n",
    "import warnings\n",
    "from datasets import load_dataset\n",
    "from trl import SFTTrainer\n",
    "from huggingface_hub import notebook_login\n",
    "\n",
    "# Base model\n",
    "base_model = \"mistralai/Mistral-Nemo-Instruct-2407\"\n",
    "\n",
    "# Custom dataset and new model to be created\n",
    "dataset_name = \"dataset/repo\"\n",
    "\n",
    "# New model to be created\n",
    "new_model = \"model/repo\"\n",
    "\n",
    "# Ignore all warnings\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d4c91d86-47fc-4ff4-bac4-9a2006138d0a",
   "metadata": {},
   "source": [
    "### Logging into Hugging Face\n",
    "Since both our dataset and the model we want to create are (meant to be) private, it is better to already log into our HF account."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7593199-8ec5-433b-b195-0393a8ebe2eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We already log into our own hugging face hub with a 'WRITE' token\n",
    "!huggingface-cli login --token MY_HF_WRITE_TOKEN"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6effec67-70fa-4362-b0e2-6e9adb80ae18",
   "metadata": {},
   "source": [
    "*Remark:* an alternative is to simply write\n",
    "```\n",
    "notebook_login()\n",
    "```\n",
    "and then type the HF 'WRITE' token.\n",
    "\n",
    "### Loading the dataset\n",
    "\n",
    "Our dataset here is of the form:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e96f2849-3059-42d4-8753-abe2c2fffdf1",
   "metadata": {},
   "source": [
    "```json\n",
    "   [\n",
    "       {\n",
    "           \"source\": \"Chat no 1-1\",\n",
    "           \"text\": [\n",
    "               {\"role\": \"user\", \"content\": \"My system prompt\\n\\nQuestion 11\"},\n",
    "               {\"role\": \"assistant\", \"content\": \"Answer 11\"}\n",
    "           ]\n",
    "       },\n",
    "       {\n",
    "           \"source\": \"Chat no 1-2\",\n",
    "           \"text\": [\n",
    "               {\"role\": \"user\", \"content\": \"Question 11\"},\n",
    "               {\"role\": \"assistant\", \"content\": \"Answer 11\"},\n",
    "               {\"role\": \"user\", \"content\": \"My system prompt\\n\\nQuestion 12\"},\n",
    "               {\"role\": \"assistant\", \"content\": \"Answer 12\"}\n",
    "           ]\n",
    "       },\n",
    "       {\n",
    "           \"source\": \"Chat no 2\",\n",
    "           \"text\": [\n",
    "               {\"role\": \"user\", \"content\": \"My system prompt\\n\\nQuestion 21\"},\n",
    "               {\"role\": \"assistant\", \"content\": \"Answer 21\"}\n",
    "   ]\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "647c3363-8791-4e21-b3f5-a3bc08af4d74",
   "metadata": {},
   "source": [
    "We will convert it later on using the chat template that comes with the tokenizer.\n",
    "\n",
    "*Remark:* Note that the system prompt is included in the very last user message of each \"chats\" that make up the dataset. This is because, even though Nemo supports system prompts (as per the Jinja code given after executing ```tokenizer.chat_template```), there seem to be an issue with the tokenizer in this respect. It will produce the same effect anyway (the model is supposed to be aware with this way of doing)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "374bf038-86d6-4b82-b223-4f0a2e4cb9c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Loading the dataset\n",
    "dataset = load_dataset(dataset_name, split=\"train\", encoding='latin-1')\n",
    "\n",
    "# We \"shuffle\" the dataset\n",
    "dataset = dataset.shuffle()\n",
    "\n",
    "dataset[\"text\"][0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "070da9d3-fb2d-466e-a423-c39beedc019a",
   "metadata": {},
   "source": [
    "### Loading the base model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9db212ec-902e-4d30-bfb5-a04148c729da",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The quantization to apply to the base model\n",
    "bnb_config = BitsAndBytesConfig(\n",
    "    load_in_4bit= True,\n",
    "    bnb_4bit_quant_type= \"nf4\",\n",
    "    bnb_4bit_compute_dtype= torch.bfloat16,\n",
    "    bnb_4bit_use_double_quant= False,\n",
    ")\n",
    "\n",
    "# We load quantized base model\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    base_model,\n",
    "    quantization_config=bnb_config,\n",
    "    device_map=\"auto\",\n",
    "    attn_implementation=\"flash_attention_2\"\n",
    ")\n",
    "\n",
    "model.config.use_cache = False\n",
    "model.config.pretraining_tp = 1\n",
    "model.gradient_checkpointing_enable()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c9c2f0e-40ae-4579-953c-cd021ca2291d",
   "metadata": {},
   "source": [
    "### Loading the tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f33deef-9b44-4717-9f85-ee0956a81443",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load tokenizer\n",
    "tokenizer = AutoTokenizer.from_pretrained(base_model, trust_remote_code=True)\n",
    "tokenizer.padding_side = 'right'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e967f58-e84a-4672-9754-a700f5900a87",
   "metadata": {},
   "source": [
    "To avoid breaking French punctuation rules, we have to make sure that the following defaults as ```False```."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1daa6cd-c681-4973-b786-471bdfe75517",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer.clean_up_tokenization_spaces"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "112c4c11-6472-4703-a94f-349194155466",
   "metadata": {},
   "source": [
    "We now have to set the PAD token. But, first of all, we need to check whether the tokenizer already has one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa1d2581-c15e-445f-acba-bd3a40634ca3",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(tokenizer.pad_token)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8773667-e68c-4769-8898-23e76b624b71",
   "metadata": {},
   "source": [
    "If not, then we have to set it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f8c5e22-719e-48c2-91a8-1028650f5df3",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer.pad_token = tokenizer.unk_token"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da41f925-dc4b-40f6-80b4-f29b2fb87891",
   "metadata": {},
   "source": [
    "*Remark:* The value set for the PAD token might differ from a model to another. Sometimes ```<pad>``` or ```[pad]``` might be a better choice. The best way for choosing the right value is to check the tokenizer's vocabulary.\n",
    "\n",
    "Let's check that it worked:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5cd53fa-3fa1-4cdd-85a9-2109fa346d29",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(tokenizer.pad_token)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4012f9c9-86b3-42b8-8049-1c920602552d",
   "metadata": {},
   "source": [
    "Now, let's add the PAD and EOS tokens to the tokenizer:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1a91233-10dd-409a-adae-2ccb084c5965",
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer.add_pad_token = True\n",
    "tokenizer.add_eos_token = True"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99376e19-5215-4a4b-b266-4c6a97f3764d",
   "metadata": {},
   "source": [
    "### Formatting the dataset\n",
    "Having loaded and configured the tokenizer, we can now apply the appropriate chat template to our dataset, namely:\n",
    "```\n",
    "<s>[INST] My first question. [/INST]The model's answer to my question.</s> [INST] A follow-up question. [/INST]The model's answer to my follow-up question.</s>\n",
    "```\n",
    "Note that, to avoid having extra/undesired \"Assistant:\" or \"User:\" tags placed at the end of each instances of the dataset, it is necessary to specify ```add_generation_prompt=False``` when applying the chat template."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2fbbaa34-5e03-4432-a20d-0b0032cb27f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import Dataset\n",
    "import json\n",
    "\n",
    "# Creating a list with elements of the dataset properly formatted\n",
    "formatted_dataset = [tokenizer.apply_chat_template(conv, tokenize=False, add_generation_prompt=False)[3:] for conv in dataset[\"text\"]]\n",
    "\n",
    "# Converted dataset with appropriate chat template\n",
    "ds = Dataset.from_dict({\"text\": formatted_dataset})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9e85540-727e-43ec-a30f-26ecb8ca79aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's have a look\n",
    "print(ds[\"text\"][100])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "032888d6-44c2-434f-8ec6-3b3fe63b88a4",
   "metadata": {},
   "source": [
    "*Remark:* We added ```[3:]``` to elements of ```formatted_dataset``` in order to remove the very first ```<s>``` (BOS) token because the tokenizer adds it automatically. Otherwise we would have two BOS at the beginning, which might bring some trouble. To double check it suffices to decode the encoded data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "391a83bd-1da0-4021-a2dc-35cf1788196a",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(tokenizer.decode(tokenizer.encode(ds[\"text\"][35])))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "956ff7d6-40c7-4d58-bb38-0304631cd73f",
   "metadata": {},
   "source": [
    "The very first BOS and the very last EOS have indeed been added by the tokenizer, matching the chat template!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33a56f40-bf0c-41d3-bb13-44484ca5924b",
   "metadata": {},
   "source": [
    "### Rapid analysis of the tokenized dataset\n",
    "Unfortunately, LLMs have limited (yet modulable) context windows, so it is worth analyzing our dataset in order to know which context length better fits our needs. Indeed, later on we will need to specify the length of the inputs used for fine-tuning (i.e. to set ```max_seq_length``` in the ```SFTTrainer()```). \n",
    "\n",
    "Concretely, we will have a look at the distribution of the length of each tokenized example. We will first (i) compute the list of all such tokenized examples and then (ii) plot it to analyze it more conveniently."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "785b1c8d-bfcb-4944-a7bf-3e91cc25940d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We compute the list of all the lengths of the tokenized examples\n",
    "length_examples = []\n",
    "for i in range(len(ds)):\n",
    "    size = len(tokenizer.encode(ds[\"text\"][i]))\n",
    "    length_examples.append(size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f059aff-f2f5-4bd5-a0f3-3bb3fc7a252f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from collections import Counter\n",
    "\n",
    "# We count the occurrences of each integer\n",
    "counter = Counter(length_examples)\n",
    "\n",
    "# We get the integers and their corresponding counts\n",
    "x = list(counter.keys())\n",
    "y = list(counter.values())\n",
    "\n",
    "# We sort the integers (x-axis)\n",
    "x_sorted, y_sorted = zip(*sorted(zip(x, y)))\n",
    "\n",
    "# We plot the values\n",
    "plt.bar(x_sorted, y_sorted)\n",
    "plt.xlabel('Number of tokens')\n",
    "plt.ylabel('Number of Occurrences')\n",
    "plt.title('Distribution of lengths of tokenized dataset examples')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9f6efb5-1970-495a-a2c3-088cd1beb62e",
   "metadata": {},
   "source": [
    "To understand more precisely the dataset, we can calculate, say, the $90$th percentile value of the lengths of our tokenized examples. Roughly speaking, this means determining, for any $p\\in(0,1)$, the smallest integer $N=N(p)\\in\\mathbb{N}$ such that\n",
    "<p style=\"text-align: center;\">$F_X(N)=\\mathbb{P}(X\\leq N)=p,$</p>\n",
    "where $F_X$ is the cumulative distribution function of the tokenized dataset and $X$ the length of a random example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d93579ca-fe54-4cd3-a35d-ae8a609c21da",
   "metadata": {},
   "outputs": [],
   "source": [
    "from statistics import mean\n",
    "import numpy as np\n",
    "\n",
    "# Target percentile\n",
    "p = 0.9\n",
    "\n",
    "def find_percentile_value(arr,percentile):\n",
    "    if not arr:\n",
    "        raise ValueError(\"The list should not be empty\")\n",
    "    arr.sort()  # Sort the list\n",
    "    # Calculate the index for the pth percentile\n",
    "    index = int(percentile * len(arr)) - 1\n",
    "    # Handle edge cases where index might be out of bounds\n",
    "    index = max(0, min(index, len(arr) - 1))\n",
    "    return arr[index]\n",
    "\n",
    "percentile_value = find_percentile_value(length_examples,p)\n",
    "print(\"The minimum value is:\", min(length_examples))\n",
    "print(\"The arithmetic mean is:\", int(np.floor(mean(length_examples))))\n",
    "print(\"The \"+str(int(p*100))+\"th percentile value is:\", percentile_value)\n",
    "print(\"The maximum value is:\", max(length_examples))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5d994a68-d9cb-4dcb-8a85-891c81db9d37",
   "metadata": {},
   "source": [
    "### Data collation\n",
    "\n",
    "We now need to pass the dataset through a data collator. The point here is to \"mask\" the \"user questions\" (prompts) within the dataset, so that the model learns the \"answers\" only (there is no point in learning the questions, they should only serve as context). \n",
    "\n",
    "What the collator does is that it will use the instruction/response templates to locate the starts of user input and starts of assistant response. It then fill those sections with $-100$ tokens, which will be ignored during training (that’s how the user inputs are masked to prevent them from contributing to the loss)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16edd4fc-4fd9-4cdf-95a7-73a8ec18ae66",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import DataCollatorForLanguageModeling\n",
    "import torch\n",
    "\n",
    "class CustomDataCollator(DataCollatorForLanguageModeling):\n",
    "    def __init__(self, tokenizer, mlm=False):\n",
    "        super().__init__(tokenizer=tokenizer, mlm=mlm)\n",
    "        self.inst_token = tokenizer.encode('[INST]', add_special_tokens=False)[0]\n",
    "        self.inst_end_token = tokenizer.encode('[/INST]', add_special_tokens=False)[0]\n",
    "        self.eos_token = tokenizer.eos_token_id\n",
    "\n",
    "    def torch_call(self, examples):\n",
    "        batch = super().torch_call(examples)\n",
    "\n",
    "        for i, input_ids in enumerate(batch['input_ids']):\n",
    "            # Find positions of [INST], [/INST], and EOS tokens\n",
    "            inst_positions = torch.where(input_ids == self.inst_token)[0]\n",
    "            inst_end_positions = torch.where(input_ids == self.inst_end_token)[0]\n",
    "            eos_positions = torch.where(input_ids == self.eos_token)[0]\n",
    "\n",
    "            if len(inst_positions) > 0 and len(eos_positions) > 0:\n",
    "                # Find the start of the last assistant response\n",
    "                last_inst_end_pos = inst_end_positions[-1]\n",
    "                \n",
    "                # Mask everything before the last assistant response\n",
    "                batch['labels'][i, :last_inst_end_pos+1] = -100\n",
    "\n",
    "                # If there's an EOS token after the last [/INST], unmask until that EOS\n",
    "                if eos_positions[-1] > last_inst_end_pos:\n",
    "                    batch['labels'][i, last_inst_end_pos+1:eos_positions[-1]+1] = input_ids[last_inst_end_pos+1:eos_positions[-1]+1]\n",
    "                else:\n",
    "                    # If no EOS after last [/INST], unmask until the end\n",
    "                    batch['labels'][i, last_inst_end_pos+1:] = input_ids[last_inst_end_pos+1:]\n",
    "\n",
    "        return batch\n",
    "\n",
    "# We now define the collator\n",
    "collator = CustomDataCollator(tokenizer=tokenizer)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2c2bc01-e500-4271-86fd-fed5912e3627",
   "metadata": {},
   "source": [
    "Let's check that the data collator has appropriately handled the question/answer turns:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "497183f2-cd85-49d0-a25b-386c421b45dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We apply the collator to a sample of the dataset\n",
    "collated_data = collator([tokenizer(ds[\"text\"][1])])['labels'][0]\n",
    "\n",
    "# We convert the above to a list (so as to be able to print everything)\n",
    "collated_data_list = collated_data.tolist()\n",
    "print(collated_data_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0af6110c-0bf5-4031-9eb8-3886522df583",
   "metadata": {},
   "source": [
    "*Remark:* The only thing that should not be masked as \"$-100$\" tokens in the very last \"assistant\" message."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38ae6f7c-6041-4ccd-9e83-21432702afea",
   "metadata": {},
   "source": [
    "### Setting LoRA parameters and modules\n",
    "\n",
    "Here, we will do two things: setting the LoRA parameters $r$ and $\\alpha$, and setting the target modules to be trained.\n",
    "\n",
    "##### LoRA parameters\n",
    "The values $\\alpha=r$ or $\\alpha=2r$ (considered a \"sweet spot\"), with $r$ a power of $2$, are conventional (see LoRA paper in the link below). The latter is often used to have more emphasis and the new data. \n",
    "(See: https://arxiv.org/pdf/2106.09685)\n",
    "\n",
    "##### LoRA modules\n",
    "The LoRA modules are, essentially, the matrices that the model allows us to train. Their name may depend on the model, which is why it is recommended to look at what it looks like via the command:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26baed38-f3bf-4b4b-b259-0b36347edec1",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b258567-6632-4bdf-9e96-9e5abb09192f",
   "metadata": {},
   "source": [
    "(Look for linear modules.) \n",
    "In most cases (Mistral, Llama, etc), however, the modules are something like:\n",
    "```\n",
    "[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\", \"gate_proj\", \"up_proj\", \"down_proj\", \"lm_head\"]\n",
    "```\n",
    "but they sometimes bear different names (e.g. with Falcon models)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9072728a-2d22-4bec-a1ad-fb29907df44b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Adding the adapters in the layers\n",
    "model = prepare_model_for_kbit_training(model)\n",
    "\n",
    "# LoRA config\n",
    "peft_config = LoraConfig(\n",
    "        r=256,\n",
    "        lora_alpha=256,\n",
    "        lora_dropout=0.05, # Conventional\n",
    "        bias=\"none\",\n",
    "        task_type=\"CAUSAL_LM\",\n",
    "        target_modules=\"all-linear\"\n",
    "    )\n",
    "\n",
    "# We add the LoRA \"adapters\" to the model\n",
    "model = get_peft_model(model, peft_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f685d579-1451-442b-b725-b9c0fc53091d",
   "metadata": {},
   "source": [
    "*Remark:* In most cases, one can simply set:\n",
    "```\n",
    "target_modules = 'all-linear'\n",
    "```\n",
    "but this might raise an error for some models (e.g. with Google's Gemma-2B).\n",
    "\n",
    "(See: https://stackoverflow.com/questions/76768226/target-modules-for-applying-peft-lora-on-different-models)\n",
    "\n",
    "### Setting Weights & Biases for monitoring purposes\n",
    "\n",
    "This will allow us to monitor the training directly from W&B."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac77ca8c-88e6-41ea-9e20-677cdd443804",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Monitering the LLM from WandB\n",
    "wandb.login(key = \"MY_WANDB_KEY\")\n",
    "run = wandb.init(project='MyProjectName', job_type=\"training\", anonymous=\"allow\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "beb29aae-0239-4b61-b3d4-0bc6310495ae",
   "metadata": {},
   "source": [
    "### Setting training parameters\n",
    "Here, we set the hyperparameters (epochs, batch, etc.) as well as the SFT (i.e. Supervised Fine-Tuning) parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b86ef8be-9fd4-43a8-8086-6397e7a4e10c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Hyperparamter\n",
    "training_arguments = TrainingArguments(\n",
    "    output_dir= \"./results\",\n",
    "    num_train_epochs= 3,              # total number of training epochs\n",
    "    per_device_train_batch_size= 4,   # batch size per device during training (the higher the better)\n",
    "    gradient_accumulation_steps= 8,\n",
    "    optim = \"paged_adamw_8bit\",\n",
    "    save_steps= 1000,                 # Save checkpoints #every 50 steps\n",
    "    logging_steps= 15,                # When to start reporting loss\n",
    "    learning_rate= 2e-5,              # We take a not too small value given that new tokens have been added\n",
    "    weight_decay= 0.001,              # strength of weight decay\n",
    "    fp16= False,\n",
    "    bf16= False,\n",
    "    max_grad_norm= 0.3,\n",
    "    max_steps= -1,\n",
    "    warmup_ratio= 0.3,                # number of warmup steps for learning rate scheduler\n",
    "    group_by_length= True,            # Group sequences into batches with same length (saves memory and speeds up training considerably)\n",
    "    lr_scheduler_type=\"constant\",\n",
    "    report_to=\"wandb\",\n",
    ")\n",
    "\n",
    "# Setting SFT parameters\n",
    "trainer = SFTTrainer(\n",
    "    model=model,                         # the instantiated HF Transformers model to be trained\n",
    "    train_dataset=ds,                    # Training dataset\n",
    "    data_collator=collator,              # Making sure that data collator correction is taken into account\n",
    "    peft_config=peft_config,\n",
    "    max_seq_length=4096,\n",
    "    dataset_text_field=\"text\",           # Column of the dataset to be used for training\n",
    "    tokenizer=tokenizer,\n",
    "    args=training_arguments,             # training arguments, defined above\n",
    "    packing= False,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c65bf6b7-fd2b-497a-bc6e-00b6a7187d62",
   "metadata": {},
   "source": [
    "### Training and saving the model\n",
    "A progression bar will soon appear displaying a number of steps which corresponds to\n",
    "<p style=\"text-align: center;\">$ \\#\\verb\"steps\" = \\left\\lceil\\displaystyle\\frac{\\verb\"total_dataset_size\"\\times\\#\\verb\"epochs\"}{\\verb\"batch_size\"}\\right\\rceil, $</p>\n",
    "where\n",
    "<p style=\"text-align: center;\">$\\verb\"batch_size\" = \\verb\"per_device_train_batch_size\" \\times \\verb\"gradient_accumulation_steps\",$</p>\n",
    "\n",
    "and $\\left\\lceil\\cdot\\right\\rceil$ is the ceiling function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8182f4fa-da51-467a-96bd-93c6047822be",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training the model\n",
    "trainer.train()\n",
    "\n",
    "# Save the fine-tuned model\n",
    "trainer.model.save_pretrained(new_model)\n",
    "wandb.finish()\n",
    "model.config.use_cache = True\n",
    "model.eval()\n",
    "\n",
    "#######################################################\n",
    "# We save the model in the same cell, to avoid weird  #\n",
    "#    behaviors induced by GPU providers (Colab...)    #\n",
    "#######################################################\n",
    "\n",
    "# Clear the memory footprint (NB: pretrained model has been saved to \"new_model\")\n",
    "del model, trainer\n",
    "torch.cuda.empty_cache()\n",
    "print(\"Memory footprint has been cleared.\")\n",
    "\n",
    "# Reload the base model\n",
    "base_model_reload = AutoModelForCausalLM.from_pretrained(\n",
    "    base_model, low_cpu_mem_usage=True,\n",
    "    return_dict=True,torch_dtype=torch.bfloat16,\n",
    "    device_map= \"auto\")\n",
    "print(\"Base model has been reloaded.\")\n",
    "\n",
    "# Merging the adapter with model\n",
    "model = PeftModel.from_pretrained(base_model_reload, new_model)\n",
    "model = model.merge_and_unload()\n",
    "print(\"Adapter has been merged to the base model.\")\n",
    "\n",
    "# Reload tokenizer\n",
    "tokenizer = AutoTokenizer.from_pretrained(base_model, trust_remote_code=True)\n",
    "tokenizer.padding_side = \"right\"\n",
    "tokenizer.pad_token = tokenizer.unk_token\n",
    "tokenizer.add_pad_token = True\n",
    "print(\"Tokenizer has been reloaded.\")\n",
    "\n",
    "# Pushing the merged model to hugging face hub\n",
    "model.push_to_hub(repo_id='New_model_name', private=True)\n",
    "tokenizer.push_to_hub(repo_id='New_model_name', private=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eae976cf-aa6c-4b62-bbbc-331505c99719",
   "metadata": {},
   "source": [
    "Et voilà !…"
   ]
  }
 ],
 "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.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}