code/mistral-finetune.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XIyP_0r6zuVc"
   },
   "source": [
    "# Fine-tuning Mistral 7B using QLoRA\n",
    "\n",
    "Mistral 7B is a recent open-source language model developed by MistralAI that consistently delivers state-of-the-art results across a variety of natural language understanding and generation benchmarks. While this model serves as a strong baseline for multiple downstream tasks, it can lack in domain-specific knowledge or proprietary or otherwise sensitive information. Fine-tuning is often used as a means to update a model for a specific task or tasks to better respond to domain-specific prompts. This notebook walks through downloading the Mistral 7B model from Hugging Face, preparing a custom dataset on coding-related tasks and instructions, and using Quantized Low Rank Adaptation (QLoRA) to fine-tune the base model against the dataset. While we focus on a coding-specific task in this example, this methodology can be applied seamlessly to other tasks as well. \n",
    "\n",
    "This workflow is inspired by the posts and repositories [here](https://adithyask.medium.com/a-beginners-guide-to-fine-tuning-mistral-7b-instruct-model-0f39647b20fe) and [here](https://github.com/brevdev/notebooks/blob/main/mistral-finetune.ipynb)\n",
    "\n",
    "### 0. What is LoRA? QLoRA?\n",
    "With regards to Large Language Models (LLMs), fine-tuning is the customization of pretrained models, like Mistral-7B, towards new or more domain-specific instructions and data. This process updates the model weights through retraining either all the parameters of the model (in full fine-tuning), or a certain subset of them (in parameter-efficient fine-tuning, or PEFT). Full fine-tuning may produce better results, but in many cases PEFT is preferred due to it being lesser time-consuming and resource-intensive. \n",
    "\n",
    "Low-Rank Adaptation, or LoRA, is a method of PEFT that uses smaller weight matrices in the retraining as approximation instead of updating the full weight matrix. This rank decomposition optimization technique enables greater memory efficiency and can reduce the size of GPU required in order to perform the fine-tuning successfully. \n",
    "\n",
    "QLoRA is a further optimization that reduces the precision of the model weights as well in order to provide even greater advances in memory and space efficiency. The most common quantization used for this LoRA finetuning workflow is 4-bit quantization, which provides a decent balance between model performance, and fine-tuning feasibility. In theory, incorporating these optimizations means this workflow can even work on an NVIDIA RTX 3090!\n",
    "\n",
    "Alright, enough chit-chat. Let's dive in!\n",
    "\n",
    "First, let's select the level of quantization we would like to use for this fine-tuning project. Choose from ``None``, ``8bit``, or ``4bit`` quantization levels. Keep in mind that the ``None`` option defaults to full 16 bit precision, which mean this workflow will perform **standard LoRA** fine-tuning, while the other two options apply quantization for QLoRA fine-tuning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# DEFINE QUANTIZATION HERE. Choose from (\"none\" | \"8bit\" | \"4bit\")\n",
    "QUANTIZATION = \"4bit\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, let's set our imports."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import transformers\n",
    "from datasets import load_dataset\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig\n",
    "from peft import prepare_model_for_kbit_training, LoraConfig, get_peft_model, PeftModel"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QcE4NTeFyRgd"
   },
   "source": [
    "### 1. Load in the Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FCc64bfnmd3j"
   },
   "source": [
    "While the pretrained Mistral model has some degree of code understanding and generation in addition to English natural language processing tasks, it still falls short in certain cases, which we will explore later in this notebook. For this workflow, we will aim to fine-tune the Mistral 7B model to generate high quality responses to code generation tasks. \n",
    "\n",
    "To accomplish this, we will be using [this dataset](https://huggingface.co/datasets/TokenBender/code_instructions_122k_alpaca_style) from HuggingFace that consists of 122k code instructions that follow the alpaca style of instructions, as well as the ground truth outputs we expect our model to produce. Let's go ahead and load in the dataset, and split the entries into train, test, and validation sets. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "s6f4z8EYmcJ6"
   },
   "outputs": [],
   "source": [
    "dataset = load_dataset(\"TokenBender/code_instructions_122k_alpaca_style\", split='train')\n",
    "dataset = dataset.train_test_split(test_size=0.2)\n",
    "val_test_dataset = dataset['test'].train_test_split(test_size=0.5)\n",
    "\n",
    "train_dataset = dataset[\"train\"]\n",
    "eval_dataset = val_test_dataset[\"train\"]\n",
    "test_dataset = val_test_dataset[\"test\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check that our data splits are correct."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "EmZbX-ltyRge"
   },
   "outputs": [],
   "source": [
    "print(train_dataset)\n",
    "print(eval_dataset)\n",
    "print(test_dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "shz8Xdv-yRgf"
   },
   "source": [
    "### 2. Load In the Base Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MJ-5idQwzvg-"
   },
   "source": [
    "Now, let's now load in the Mistral Model from Huggingface - `mistralai/Mistral-7B-v0.1`. We will aim to use 4-bit quantization, which is a method that significantly reduces the overall memory footprint of the fine-tuning process by reducing precision of the model parameters while preserving performance. This makes it easier to run this fine-tuning workflow on smaller GPU systems, not just A100s!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pre-define quantization configs\n",
    "\n",
    "################## 4bit ##################\n",
    "bb_config_4b = BitsAndBytesConfig(\n",
    "    load_in_4bit=True,\n",
    "    bnb_4bit_use_double_quant=True,\n",
    "    bnb_4bit_quant_type=\"nf4\",\n",
    ")\n",
    "##########################################\n",
    "\n",
    "################## 8bit ##################\n",
    "bb_config_8b = BitsAndBytesConfig(\n",
    "    load_in_8bit=True,\n",
    ")\n",
    "##########################################\n",
    "\n",
    "def quantization_config(quantization):\n",
    "    if quantization == \"8bit\":\n",
    "        return bb_config_8b\n",
    "    else:\n",
    "        return bb_config_4b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "E0Nl5mWL0k2T"
   },
   "outputs": [],
   "source": [
    "%%capture\n",
    "\n",
    "model_id = \"mistralai/Mistral-7B-v0.1\"\n",
    "hf_api_token = os.environ['HUGGING_FACE_HUB_TOKEN']\n",
    "\n",
    "if QUANTIZATION == \"none\":\n",
    "    model = AutoModelForCausalLM.from_pretrained(model_id, token=hf_api_token).to(\"cuda\")\n",
    "else: \n",
    "    model = AutoModelForCausalLM.from_pretrained(model_id, token=hf_api_token, quantization_config=quantization_config(QUANTIZATION))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7fi9wEZYyRgh"
   },
   "source": [
    "### 3. Evaluate Base Model Performance"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Vxbl4ACsyRgi"
   },
   "source": [
    "Before fine-tuning the model, let's first evaluate how well the model does on sample tasks that we intend to fine-tune on, such as generating functions in code, coding syntax and semantics, and general understanding of multiple coding languages. Here, we'll ask it a fairly standard coding question: Write a function to output the prime factorization of 2023 in python, C, and C++. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gOxnx-cAyRgi"
   },
   "outputs": [],
   "source": [
    "base_prompt = \"\"\"Write a function to output the prime factorization of 2023 in python, C, and C++\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's call the model and see what it outputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NidIuFXMyRgi"
   },
   "outputs": [],
   "source": [
    "base_tokenizer = AutoTokenizer.from_pretrained(\n",
    "    model_id,\n",
    "    add_bos_token=True,\n",
    ")\n",
    "\n",
    "model_input = base_tokenizer(base_prompt, return_tensors=\"pt\").to(\"cuda\")\n",
    "\n",
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    print(base_tokenizer.decode(model.generate(**model_input, max_new_tokens=256)[0], skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "dCAWeCzZyRgi"
   },
   "source": [
    "We can see it doesn't do very well out of the box...\n",
    "\n",
    "1. The out-of-the-box model seems to think the prime factorization of 2023 is 13 x 157. This amounts to 2041! The actual answer is 7 x 17 x 17. \n",
    "\n",
    "2. At first glance the python function it outputs is incorrect as well; if we actually run the code, it gives the answer as ``[7, 17, 119, 289, 2023]``. 119, 289, and of course 2023 are not prime factors! \n",
    "\n",
    "While the syntax is generally comprehensible, we can see that there are still issues in the output that could be improved on. Let's attempt to improve the quality of the model's outputs using fine-tuning. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UjNdXolqyRgf"
   },
   "source": [
    "### 4. Format the Data for Fine-Tuning\n",
    "\n",
    "Let's first set up the tokenizer before formatting the dataset. Left-padding is recommended here as it can [reduce memory costs](https://ai.stackexchange.com/questions/41485/while-fine-tuning-a-decoder-only-llm-like-llama-on-chat-dataset-what-kind-of-pa).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "haSUDD9HyRgf"
   },
   "outputs": [],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(\n",
    "    model_id,\n",
    "    model_max_length=512,\n",
    "    padding_side=\"left\",\n",
    "    add_eos_token=True)\n",
    "\n",
    "tokenizer.pad_token = tokenizer.eos_token\n",
    "\n",
    "def tokenize(prompt):\n",
    "    tokenized = tokenizer(\n",
    "        prompt,\n",
    "        truncation=True,\n",
    "        max_length=512,\n",
    "        padding=\"max_length\",\n",
    "    )\n",
    "    tokenized[\"labels\"] = tokenized[\"input_ids\"].copy()\n",
    "    return tokenized"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tJtsbrr6yRgg"
   },
   "source": [
    "We can then reformat the dataset to fit the instruction prompt for fine-tuning. We will enclose the instruction and any inputs given to the model in a ``[INST]`` tag, and attach the correct output afterwards. \n",
    "\n",
    "Then, we tokenize each entry of our dataset using the tokenizer we set up above. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6z9rvnoDyRgg"
   },
   "outputs": [],
   "source": [
    "def process_prompt(data):\n",
    "    new_prompt = f\"\"\"<s>[INST] {data[\"instruction\"]} here are the inputs {data[\"input\"]} [/INST] \\\\n {data[\"output\"]} </s>\"\"\"\n",
    "    return tokenize(new_prompt)\n",
    "\n",
    "tokenized_train_ds = train_dataset.map(process_prompt)\n",
    "tokenized_val_ds = eval_dataset.map(process_prompt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AapDoyfAyRgi"
   },
   "source": [
    "### 5. Set up for QLoRA Fine-Tuning"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Mp2gMi1ZzGET"
   },
   "source": [
    "Now, we are ready to set up our fine-tuning workflow. Let's prepare the model for parameter efficient fine-tuning. We'll also implement a neat function to let us know exactly how many of the model weights will be retrained and how many will be frozen, just to get a good idea for how PEFT is working under the hood. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gkIcwsSU01EB"
   },
   "outputs": [],
   "source": [
    "model.gradient_checkpointing_enable()\n",
    "model = prepare_model_for_kbit_training(model)\n",
    "\n",
    "def print_param_info(model):\n",
    "    \"\"\"\n",
    "    Outputs trainable parameter information.\n",
    "    \"\"\"\n",
    "    trainable_params = 0\n",
    "    all_param = 0\n",
    "    for _, param in model.named_parameters():\n",
    "        all_param += param.numel()\n",
    "        if param.requires_grad:\n",
    "            trainable_params += param.numel()\n",
    "    print(\n",
    "        f\"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "cUYEpEK-yRgj"
   },
   "source": [
    "Next, we can print out the architecture of the model. QLoRA will be applied to all the linear layers of this model. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XshGNsbxyRgj",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "I6mTLuQJyRgj"
   },
   "source": [
    "We can see those layers are: \n",
    "* `q_proj`\n",
    "* `k_proj`\n",
    "* `v_proj`\n",
    "* `o_proj`\n",
    "* `gate_proj`\n",
    "* `up_proj`\n",
    "* `down_proj`\n",
    "* `lm_head`\n",
    "\n",
    "Let's make a note of these and pass them into the LoRA config. You are also able to specify a few other fields in the config. For example: \n",
    "\n",
    "* `r`: This field refers to the rank of the lower-rank matrices you want to use in the adaptation layers of the model, which controls the number of parameters set to be retrained. The higher this number, the more expressiveness you will capture; however, there is added computational cost. \n",
    "\n",
    "* `alpha`: This field refers to the scaling factor for the weights. The weights are scaled by a factor of `alpha/r`, and so the higher this number means more weights are assigned to the LoRA activations.\n",
    "\n",
    "The authors of the original QLoRA paper used the following values: `r=64` and `lora_alpha=16`. While these may be able to generalize well, let's set the defaults here to `r=8` and `lora_alpha=16`. This way, we allocate a greater amount of weights as retrainable on the new fine-tuned data while also minimizing computational complexity. You are free to adjust and tune these parameters as you wish. \n",
    "\n",
    "Let's use the ``print_param_info`` defined above to see what the trainable and frozen parameters look like. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Ybeyl20n3dYH"
   },
   "outputs": [],
   "source": [
    "config = LoraConfig(\n",
    "    r=8,\n",
    "    lora_alpha=16,\n",
    "    target_modules=[\n",
    "        \"q_proj\",\n",
    "        \"k_proj\",\n",
    "        \"v_proj\",\n",
    "        \"o_proj\",\n",
    "        \"gate_proj\",\n",
    "        \"up_proj\",\n",
    "        \"down_proj\",\n",
    "        \"lm_head\",\n",
    "    ],\n",
    "    bias=\"none\",\n",
    "    lora_dropout=0.05,\n",
    "    task_type=\"CAUSAL_LM\",\n",
    ")\n",
    "\n",
    "model = get_peft_model(model, config)\n",
    "print_param_info(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "X_FHi_VLyRgn"
   },
   "source": [
    "And reprinting the model architecture shows us the updated model with proper quantization and LoRA layers wrapping the original linear layers. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "IaYMWak4yRgn"
   },
   "outputs": [],
   "source": [
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_0MOtwf3zdZp"
   },
   "source": [
    "### 6. Run QLoRA Fine-Tuning"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fEe0uWYSyRgo"
   },
   "source": [
    "Now with the dataset processed and tokenized, and with the model prepared, we are ready to begin running the fine-tuning. This following cell will configure the trainer object with various default parameters. \n",
    "\n",
    "On an 1x A100-80GB system, this cell can take several hours to complete as-written. Depending on your hardware and patience, you may need to adjust certain parameters to achieve reasonable training times. Notably, we set the `max_steps` to 1000 and the checkpoint and evaluation to every 50 steps; you may reduce the number of steps and/or make less frequent checkpoints if you would like to reduce the training time. \n",
    "\n",
    "For your convenience, a progress bar is generated, as well as checkpointing for the training and validation errors. If the validation error begins increasing, you may be running into issues with model overfitting. At this point, you may interrupt the kernel to stop training, and pass the appropriate ``checkpoint-xx`` to Step 7. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "jq0nX33BmfaC"
   },
   "outputs": [],
   "source": [
    "# Parallelization is possible if system is multi-GPU\n",
    "if torch.cuda.device_count() > 1: \n",
    "    model.is_parallelizable = True\n",
    "    model.model_parallel = True\n",
    "\n",
    "tokenizer.pad_token = tokenizer.eos_token\n",
    "\n",
    "# Training configs\n",
    "trainer = transformers.Trainer(\n",
    "    model=model,\n",
    "    train_dataset=tokenized_train_ds,\n",
    "    eval_dataset=tokenized_val_ds,\n",
    "    args=transformers.TrainingArguments(\n",
    "        output_dir=\"./mistral-code-instruct\",\n",
    "        warmup_steps=5,\n",
    "        per_device_train_batch_size=2,\n",
    "        gradient_checkpointing=True,\n",
    "        gradient_accumulation_steps=4,\n",
    "        max_steps=500,\n",
    "        learning_rate=2.5e-5,\n",
    "        logging_steps=50,\n",
    "        bf16=True if (QUANTIZATION != \"8bit\") else False,\n",
    "        fp16=True if (QUANTIZATION == \"8bit\") else False,\n",
    "        optim=\"paged_adamw_8bit\",\n",
    "        logging_dir=\"./logs\",\n",
    "        save_strategy=\"steps\",\n",
    "        save_steps=50,\n",
    "        evaluation_strategy=\"steps\", \n",
    "        eval_steps=50,\n",
    "        report_to=\"none\",\n",
    "        do_eval=True,\n",
    "    ),\n",
    "    data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),\n",
    ")\n",
    "\n",
    "# Silencing warnings. If using for inference, consider re-enabling.\n",
    "model.config.use_cache = False \n",
    "\n",
    "# Train! \n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0D57XqcsyRgo"
   },
   "source": [
    "### 7. Evaluate the Fine-Tuned Model\n",
    "\n",
    "Good news, the model is now fine-tuned to your dataset! \n",
    "\n",
    "If you find you are running low on VRAM, you may consider restarting the kernel at this point. The PEFT library functionality saves only the QLoRA adapters in the checkpoints by default, and so the original weights need to be reloaded. Restarting the kernel may prevent any out-of-memory headaches when loading the base model again on top of this customized model we just fine-tuned. \n",
    "\n",
    "In case you restarted the kernel, let's redefine everything again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import transformers\n",
    "from datasets import load_dataset\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig\n",
    "from peft import prepare_model_for_kbit_training, LoraConfig, get_peft_model, PeftModel\n",
    "\n",
    "# Pre-define quantization configs\n",
    "\n",
    "################## 4bit ##################\n",
    "bb_config_4b_eval = BitsAndBytesConfig(\n",
    "    load_in_4bit=True,\n",
    "    bnb_4bit_use_double_quant=True,\n",
    "    bnb_4bit_quant_type=\"nf4\",\n",
    "    bnb_4bit_compute_dtype=torch.bfloat16\n",
    ")\n",
    "##########################################\n",
    "\n",
    "################## 8bit ##################\n",
    "bb_config_8b_eval = BitsAndBytesConfig(\n",
    "    load_in_8bit=True,\n",
    ")\n",
    "##########################################\n",
    "\n",
    "def quantization_config_eval(quantization):\n",
    "    if quantization == \"8bit\":\n",
    "        return bb_config_8b_eval\n",
    "    else:\n",
    "        return bb_config_4b_eval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "SKSnF016yRgp"
   },
   "outputs": [],
   "source": [
    "model_id = \"mistralai/Mistral-7B-v0.1\"\n",
    "hf_api_token = os.environ['HUGGING_FACE_HUB_TOKEN']\n",
    "\n",
    "if QUANTIZATION == \"none\":\n",
    "    base_model = AutoModelForCausalLM.from_pretrained(\n",
    "        model_id,\n",
    "        token=hf_api_token, \n",
    "        device_map=\"auto\",\n",
    "        trust_remote_code=True,\n",
    "    ).to(\"cuda\")\n",
    "else: \n",
    "    base_model = AutoModelForCausalLM.from_pretrained(\n",
    "        model_id,\n",
    "        token=hf_api_token, \n",
    "        quantization_config=quantization_config_eval(QUANTIZATION),\n",
    "        device_map=\"auto\",\n",
    "        trust_remote_code=True,\n",
    "    )\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(\n",
    "    model_id,\n",
    "    model_max_length=512,\n",
    "    padding_side=\"left\",\n",
    "    add_eos_token=True)\n",
    "\n",
    "tokenizer.pad_token = tokenizer.eos_token"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_BxOhAiqyRgp"
   },
   "source": [
    "Now, we can merge our updated model weights from the QLoRA training with the original weights of the base Mistral model. Make sure you choose the best performing model checkpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "GwsiqhWuyRgp"
   },
   "outputs": [],
   "source": [
    "ft_model = PeftModel.from_pretrained(base_model, \"mistral-code-instruct/checkpoint-500\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UUehsaVNyRgp"
   },
   "source": [
    "We are ready to use this merged model for inference! Let's go ahead and try a similar prime factorization programming question to what we had asked above, and see if our fine-tuned Mistral model achieves better quality responses."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_prompt = f\"\"\"<s>\n",
    "For a given integer n, print out all its prime factors one on each line. \n",
    "n = 30\n",
    "[INST]\n",
    "\"\"\"\n",
    "\n",
    "input_ids = tokenizer(eval_prompt, return_tensors=\"pt\", truncation=True).input_ids.cuda()\n",
    "outputs = ft_model.generate(input_ids=input_ids, max_new_tokens=256, do_sample=True, top_p=0.9,temperature=0.5)\n",
    "\n",
    "print(f\"Prompt:\\n{eval_prompt}\\n\")\n",
    "print(f\"\\nGenerated response:\\n{tokenizer.batch_decode(outputs.detach().cpu().numpy(), skip_special_tokens=True)[0][len(eval_prompt):]}\")\n",
    "print('''\\nGround truth:\\ndef print_prime_factors(n): \n",
    "  for i in range(2, n + 1):\n",
    "    while n % i == 0:\n",
    "      print(i)\n",
    "      n //= i\n",
    "print_prime_factors(n)''')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "30 can be factored into the following primes: 2, 3, 5. Because we fine tune on generating code snippets and not answering the question posed to the LLM, you may in some cases see 'hallucinatory' returned answers that may not align perfectly with the actual ouput of the generated code snippet, so be sure to examine the generated code snippet rather than solely relying on the returned outputted response or answer. Feel free to spin up a sandbox environment to evaluate any generated code. \n",
    "\n",
    "Check out the ``Generated response`` output and compare it with the ``Ground truth`` code. Try out the ``Generated response`` yourself in a sandbox environment. Could you be underfitting? Overfitting? Or does the code work as intended? \n",
    "\n",
    "If so, nice! Using QLoRA fine-tuning, we can now generate comprehensible and accurate code that accomplishes what the out-of-the-box baseline Mistral model was unable to achieve. Now, feel free to adjust the hyperparameters, bring in your own custom data, or customize this fine-tuning workflow to improve model performance for your particular use case. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 7. Merge and Save the Fine-tuned Model\n",
    "\n",
    "Now, we are ready to save the fine tuned model weights to the base model. Let's save this in under `models`, which you have already mounted to your host system for easy access. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ft_model.save_pretrained(\"/project/models/finetuned\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "gpuClass": "standard",
  "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}