mistral-finetune.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "XIyP_0r6zuVc"
   },
   "source": [
    "<!-- Banner Image -->\n",
    "<img src=\"https://uohmivykqgnnbiouffke.supabase.co/storage/v1/object/public/landingpage/brevdevnotebooks.png\" width=\"100%\">\n",
    "\n",
    "<!-- Links -->\n",
    "<center>\n",
    "  <a href=\"https://console.brev.dev\" style=\"color: #06b6d4;\">Console</a> •\n",
    "  <a href=\"https://brev.dev\" style=\"color: #06b6d4;\">Docs</a> •\n",
    "  <a href=\"/\" style=\"color: #06b6d4;\">Templates</a> •\n",
    "  <a href=\"https://discord.gg/NVDyv7TUgJ\" style=\"color: #06b6d4;\">Discord</a>\n",
    "</center>\n",
    "\n",
    "# Fine-tuning Mistral 7B using QLoRA 🤙\n",
    "\n",
    "Welcome!\n",
    "\n",
    "In this notebook and tutorial, we will fine-tune the [Mistral 7B](https://github.com/mistralai/mistral-src) model - which outperforms Llama 2 13B on all tested benchmarks.\n",
    "\n",
    "**Watch an accompanying video walk-through (but for using your own data) [here](https://youtu.be/kmkcNVvEz-k?si=Ogt1wRFNqYI6zXfw&t=1)!** If you'd like to see that notebook instead, click [here](https://github.com/brevdev/notebooks/blob/main/mistral-finetune-own-data.ipynb).\n",
    "\n",
    "This tutorial will use QLoRA, a fine-tuning method that combines quantization and LoRA. For more information about what those are and how they work, see [this post](https://brev.dev/blog/how-qlora-works).\n",
    "\n",
    "In this notebook, we will load the large model in 4bit using `bitsandbytes` (`Mistral-7B-v0.1`) and use LoRA to train using the PEFT library from Hugging Face 🤗.\n",
    "\n",
    "Note that if you ever have trouble importing something from Huggingface, you may need to run `huggingface-cli login` in a shell. To open a shell in Jupyter Lab, click on 'Launcher' (or the '+' if it's not there) next to the notebook tab at the top of the screen. Under \"Other\", click \"Terminal\" and then run the command.\n",
    "\n",
    "### Help us make this tutorial better! Please provide feedback on the [Discord channel](https://discord.gg/y9428NwTh3) or on [X](https://x.com/harperscarroll)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "SSBw-KpkyRga"
   },
   "source": [
    "#### Before we begin: A note on OOM errors\n",
    "\n",
    "If you get an error like this: `OutOfMemoryError: CUDA out of memory`, tweak your parameters to make the model less computationally intensive. I will help guide you through that in this guide, and if you have any additional questions you can reach out on the [Discord channel](https://discord.gg/y9428NwTh3) or on [X](https://x.com/harperscarroll).\n",
    "\n",
    "To re-try after you tweak your parameters, open a Terminal ('Launcher' or '+' in the nav bar above -> Other -> Terminal) and run the command `nvidia-smi`. Then find the process ID `PID` under `Processes` and run the command `kill [PID]`. You will need to re-start your notebook from the beginning. (There may be a better way to do this... if so please do let me know!)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "hWI-uRLEyRgb"
   },
   "source": [
    "## Let's begin!\n",
    "\n",
    "I used a GPU and dev environment from [brev.dev](https://brev.dev). The whole thing cost me $1 using a 1xA10G 24GB. Click the badge below to get your preconfigured instance:\n",
    "\n",
    "[![](https://uohmivykqgnnbiouffke.supabase.co/storage/v1/object/public/landingpage/brevdeploynavy.svg)](https://console.brev.dev/environment/new?instance=A10G:g5.xlarge&diskStorage=256&name=mistral-finetune&file=https://github.com/brevdev/notebooks/raw/main/mistral-finetune.ipynb&python=3.10&cuda=12.1.1)\n",
    "\n",
    "A single A10G (as linked) or L4 should be enough for this dataset; anything with >= 24GB GPU Memory. You may need more GPUs and/or Memory if your sequence max_length is larger than 512.\n",
    "\n",
    "Once you've checked out your machine and landed in your instance page, select the specs you'd like (I used **Python 3.10 and CUDA 12.1.1**; these should be preconfigured for you if you use the badge above) and click the \"Build\" button to build your verb container. Give this a few minutes.\n",
    "\n",
    "A few minutes after your model has started Running, click the 'Notebook' button on the top right of your screen once it illuminates (you may need to refresh the screen). You will be taken to a Jupyter Lab environment, where you can upload this Notebook.\n",
    "\n",
    "\n",
    "Note: You can connect your cloud credits (AWS or GCP) by clicking \"Org: \" on the top right, and in the panel that slides over, click \"Connect AWS\" or \"Connect GCP\" under \"Connect your cloud\" and follow the instructions linked to attach your credentials.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "FuXIFTFapAMI"
   },
   "outputs": [],
   "source": [
    "# You only need to run this once per machine\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 -U datasets scipy ipywidgets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "05H5MIfjyRgc"
   },
   "source": [
    "### 0. Accelerator\n",
    "\n",
    "Set up the Accelerator. I'm not sure if we really need this for a QLoRA given its [description](https://huggingface.co/docs/accelerate/v0.19.0/en/usage_guides/fsdp) (I have to read more about it) but it seems it can't hurt, and it's helpful to have the code for future reference. You can always comment out the accelerator if you want to try without."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "TEzYBadkyRgd"
   },
   "outputs": [],
   "source": [
    "from accelerate import FullyShardedDataParallelPlugin, Accelerator\n",
    "from torch.distributed.fsdp.fully_sharded_data_parallel import FullOptimStateDictConfig, FullStateDictConfig\n",
    "\n",
    "fsdp_plugin = FullyShardedDataParallelPlugin(\n",
    "    state_dict_config=FullStateDictConfig(offload_to_cpu=True, rank0_only=False),\n",
    "    optim_state_dict_config=FullOptimStateDictConfig(offload_to_cpu=True, rank0_only=False),\n",
    ")\n",
    "\n",
    "accelerator = Accelerator(fsdp_plugin=fsdp_plugin)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QcE4NTeFyRgd"
   },
   "source": [
    "### 1. Load Dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "FCc64bfnmd3j"
   },
   "source": [
    "Let's load a meaning representation dataset, and fine-tune Mistral on that. This is a great fine-tuning dataset as it teaches the model a unique form of desired output on which the base model performs poorly out-of-the box, so it's helpful to easily and inexpensively gauge whether the fine-tuned model has learned well. (Sources: [here](https://ragntune.com/blog/gpt3.5-vs-llama2-finetuning) and [here](https://www.anyscale.com/blog/fine-tuning-is-for-form-not-facts)) (In contrast, if you fine-tune on a fact-based dataset, the model may already do quite well on that, and gauging learning is less obvious / may be more computationally expensive.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "s6f4z8EYmcJ6"
   },
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "train_dataset = load_dataset('gem/viggo', split='train')\n",
    "eval_dataset = load_dataset('gem/viggo', split='validation')\n",
    "test_dataset = load_dataset('gem/viggo', split='test')"
   ]
  },
  {
   "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 Base Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MJ-5idQwzvg-"
   },
   "source": [
    "Let's now load Mistral - `mistralai/Mistral-7B-v0.1` - using 4-bit quantization!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "E0Nl5mWL0k2T"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig\n",
    "\n",
    "base_model_id = \"mistralai/Mistral-7B-v0.1\"\n",
    "bnb_config = 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",
    "model = AutoModelForCausalLM.from_pretrained(base_model_id, quantization_config=bnb_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UjNdXolqyRgf"
   },
   "source": [
    "### 3. Tokenization\n",
    "\n",
    "Set up the tokenizer. Add padding on the left as it [makes training use less memory](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",
    "    base_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": "WLvc85zDyRgg"
   },
   "source": [
    "Setup the tokenize function to make labels and input_ids the same. This is basically what [self-supervised fine-tuning is](https://neptune.ai/blog/self-supervised-learning):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "1hFsEFp5yRgg"
   },
   "outputs": [],
   "source": [
    "def tokenize(prompt):\n",
    "    result = tokenizer(\n",
    "        prompt,\n",
    "        truncation=True,\n",
    "        max_length=512,\n",
    "        padding=\"max_length\",\n",
    "    )\n",
    "    result[\"labels\"] = result[\"input_ids\"].copy()\n",
    "    return result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "tJtsbrr6yRgg"
   },
   "source": [
    "And convert each sample into a prompt that I found from [this notebook](https://github.com/samlhuillier/viggo-finetune/blob/main/llama/fine-tune-code-llama.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "6z9rvnoDyRgg"
   },
   "outputs": [],
   "source": [
    "def generate_and_tokenize_prompt(data_point):\n",
    "    full_prompt =f\"\"\"Given a target sentence construct the underlying meaning representation of the input sentence as a single function with attributes and attribute values.\n",
    "This function should describe the target string accurately and the function must be one of the following ['inform', 'request', 'give_opinion', 'confirm', 'verify_attribute', 'suggest', 'request_explanation', 'recommend', 'request_attribute'].\n",
    "The attributes must be one of the following: ['name', 'exp_release_date', 'release_year', 'developer', 'esrb', 'rating', 'genres', 'player_perspective', 'has_multiplayer', 'platforms', 'available_on_steam', 'has_linux_release', 'has_mac_release', 'specifier']\n",
    "\n",
    "### Target sentence:\n",
    "{data_point[\"target\"]}\n",
    "\n",
    "### Meaning representation:\n",
    "{data_point[\"meaning_representation\"]}\n",
    "\"\"\"\n",
    "    return tokenize(full_prompt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "WHnKLcq4yRgg"
   },
   "source": [
    "Reformat the prompt and tokenize each sample:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "q2qYeNA2yRgh"
   },
   "outputs": [],
   "source": [
    "tokenized_train_dataset = train_dataset.map(generate_and_tokenize_prompt)\n",
    "tokenized_val_dataset = eval_dataset.map(generate_and_tokenize_prompt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TQL796OayRgh"
   },
   "source": [
    "Check that `input_ids` is padded on the left with the `eos_token` (2) and there is an `eos_token` 2 added to the end, and the prompt starts with a `bos_token` (1).\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "-CRrG-SkyRgh"
   },
   "outputs": [],
   "source": [
    "print(tokenized_train_dataset[4]['input_ids'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "nBk4Qp_vyRgh"
   },
   "source": [
    "Check that a sample has the max length, i.e. 512."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "V2_zCXcJyRgh"
   },
   "outputs": [],
   "source": [
    "print(len(tokenized_train_dataset[4]['input_ids']))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7fi9wEZYyRgh"
   },
   "source": [
    "#### How does the base model do?"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Vxbl4ACsyRgi"
   },
   "source": [
    "Let's grab a test input (`meaning_representation`) and desired output (`target`) pair to see how the base model does on it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "k_VRZDh9yRgi",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print(\"Target Sentence: \" + test_dataset[1]['target'])\n",
    "print(\"Meaning Representation: \" + test_dataset[1]['meaning_representation'] + \"\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gOxnx-cAyRgi"
   },
   "outputs": [],
   "source": [
    "eval_prompt = \"\"\"Given a target sentence construct the underlying meaning representation of the input sentence as a single function with attributes and attribute values.\n",
    "This function should describe the target string accurately and the function must be one of the following ['inform', 'request', 'give_opinion', 'confirm', 'verify_attribute', 'suggest', 'request_explanation', 'recommend', 'request_attribute'].\n",
    "The attributes must be one of the following: ['name', 'exp_release_date', 'release_year', 'developer', 'esrb', 'rating', 'genres', 'player_perspective', 'has_multiplayer', 'platforms', 'available_on_steam', 'has_linux_release', 'has_mac_release', 'specifier']\n",
    "\n",
    "### Target sentence:\n",
    "Earlier, you stated that you didn't have strong feelings about PlayStation's Little Big Adventure. Is your opinion true for all games which don't have multiplayer?\n",
    "\n",
    "### Meaning representation:\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NidIuFXMyRgi"
   },
   "outputs": [],
   "source": [
    "# Re-init the tokenizer so it doesn't add padding or eos token\n",
    "eval_tokenizer = AutoTokenizer.from_pretrained(\n",
    "    base_model_id,\n",
    "    add_bos_token=True,\n",
    ")\n",
    "\n",
    "model_input = eval_tokenizer(eval_prompt, return_tensors=\"pt\").to(\"cuda\")\n",
    "\n",
    "model.eval()\n",
    "with torch.no_grad():\n",
    "    print(eval_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."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "AapDoyfAyRgi"
   },
   "source": [
    "### 4. Set Up LoRA"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Mp2gMi1ZzGET"
   },
   "source": [
    "Now, to start our fine-tuning, we have to apply some preprocessing to the model to prepare it for training. For that use the `prepare_model_for_kbit_training` method from PEFT."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "a9EUEDAl0ss3"
   },
   "outputs": [],
   "source": [
    "from peft import prepare_model_for_kbit_training\n",
    "\n",
    "model.gradient_checkpointing_enable()\n",
    "model = prepare_model_for_kbit_training(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "gkIcwsSU01EB"
   },
   "outputs": [],
   "source": [
    "def print_trainable_parameters(model):\n",
    "    \"\"\"\n",
    "    Prints the number of trainable parameters in the model.\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": [
    "Let's print the model to examine its layers, as we will apply QLoRA to all the linear layers of the model. Those layers are `q_proj`, `k_proj`, `v_proj`, `o_proj`, `gate_proj`, `up_proj`, `down_proj`, and `lm_head`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XshGNsbxyRgj",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "I6mTLuQJyRgj"
   },
   "source": [
    "Here we define the LoRA config.\n",
    "\n",
    "`r` is the rank of the low-rank matrix used in the adapters, which thus controls the number of parameters trained. A higher rank will allow for more expressivity, but there is a compute tradeoff.\n",
    "\n",
    "`alpha` is the scaling factor for the learned weights. The weight matrix is scaled by `alpha/r`, and thus a higher value for `alpha` assigns more weight to the LoRA activations.\n",
    "\n",
    "The values used in the QLoRA paper were `r=64` and `lora_alpha=16`, and these are said to generalize well, but we will use `r=8` and `lora_alpha=16` so that we have more emphasis on the new fine-tuned data while also reducing computational complexity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Ybeyl20n3dYH"
   },
   "outputs": [],
   "source": [
    "from peft import LoraConfig, get_peft_model\n",
    "\n",
    "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,  # Conventional\n",
    "    task_type=\"CAUSAL_LM\",\n",
    ")\n",
    "\n",
    "model = get_peft_model(model, config)\n",
    "print_trainable_parameters(model)\n",
    "\n",
    "# Apply the accelerator. You can comment this out to remove the accelerator.\n",
    "model = accelerator.prepare_model(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "X_FHi_VLyRgn"
   },
   "source": [
    "See how the model looks different now, with the LoRA adapters added:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "IaYMWak4yRgn"
   },
   "outputs": [],
   "source": [
    "print(model)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-9KNTJZkyRgn"
   },
   "source": [
    "\n",
    "Let's use Weights & Biases to track our training metrics. You'll need to apply an API key when prompted. Feel free to skip this if you'd like, and just comment out the `wandb` parameters in the `Trainer` definition below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "DDqUNyIoyRgo"
   },
   "outputs": [],
   "source": [
    "!pip install -q wandb -U\n",
    "\n",
    "import wandb, os\n",
    "wandb.login()\n",
    "\n",
    "wandb_project = \"viggo-finetune\"\n",
    "if len(wandb_project) > 0:\n",
    "    os.environ[\"WANDB_PROJECT\"] = wandb_project"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_0MOtwf3zdZp"
   },
   "source": [
    "### 5. Run Training!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fEe0uWYSyRgo"
   },
   "source": [
    "I used 500 steps, but I found the model should have trained for longer as it had not converged by then, so I upped the steps to 1000 below.\n",
    "\n",
    "A note on training. You can set the `max_steps` to be high initially, and examine at what step your model's performance starts to degrade. There is where you'll find a sweet spot for how many steps to perform. For example, say you start with 1000 steps, and find that at around 500 steps the model starts overfitting - the validation loss goes up (bad) while the training loss goes down significantly, meaning the model is learning the training set really well, but is unable to generalize to new datapoints. Therefore, 500 steps would be your sweet spot, so you would use the `checkpoint-500` model repo in your output dir (`mistral-finetune-viggo`) as your final model in step 6 below.\n",
    "\n",
    "You can interrupt the process via Kernel -> Interrupt Kernel in the top nav bar once you realize you didn't need to train anymore."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "c_L1131GyRgo"
   },
   "outputs": [],
   "source": [
    "if torch.cuda.device_count() > 1: # If more than 1 GPU\n",
    "    model.is_parallelizable = True\n",
    "    model.model_parallel = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "jq0nX33BmfaC"
   },
   "outputs": [],
   "source": [
    "import transformers\n",
    "from datetime import datetime\n",
    "\n",
    "project = \"viggo-finetune\"\n",
    "base_model_name = \"mistral\"\n",
    "run_name = base_model_name + \"-\" + project\n",
    "output_dir = \"./\" + run_name\n",
    "\n",
    "tokenizer.pad_token = tokenizer.eos_token\n",
    "\n",
    "trainer = transformers.Trainer(\n",
    "    model=model,\n",
    "    train_dataset=tokenized_train_dataset,\n",
    "    eval_dataset=tokenized_val_dataset,\n",
    "    args=transformers.TrainingArguments(\n",
    "        output_dir=output_dir,\n",
    "        warmup_steps=5,\n",
    "        per_device_train_batch_size=2,\n",
    "        gradient_checkpointing=True,\n",
    "        gradient_accumulation_steps=4,\n",
    "        max_steps=1000,\n",
    "        learning_rate=2.5e-5, # Want about 10x smaller than the Mistral learning rate\n",
    "        logging_steps=50,\n",
    "        bf16=True,\n",
    "        optim=\"paged_adamw_8bit\",\n",
    "        logging_dir=\"./logs\",        # Directory for storing logs\n",
    "        save_strategy=\"steps\",       # Save the model checkpoint every logging step\n",
    "        save_steps=50,                # Save checkpoints every 50 steps\n",
    "        evaluation_strategy=\"steps\", # Evaluate the model every logging step\n",
    "        eval_steps=50,               # Evaluate and save checkpoints every 50 steps\n",
    "        do_eval=True,                # Perform evaluation at the end of training\n",
    "        report_to=\"wandb\",           # Comment this out if you don't want to use weights & baises\n",
    "        run_name=f\"{run_name}-{datetime.now().strftime('%Y-%m-%d-%H-%M')}\"          # Name of the W&B run (optional)\n",
    "    ),\n",
    "    data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),\n",
    ")\n",
    "\n",
    "model.config.use_cache = False  # silence the warnings. Please re-enable for inference!\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "0D57XqcsyRgo"
   },
   "source": [
    "### 6. Drum Roll... Try the Trained Model!\n",
    "\n",
    "It's a good idea to kill the current process so that you don't run out of memory loading the base model again on top of the model we just trained. Go to `Kernel > Restart Kernel` or kill the process via the Terminal (`nvidia smi` > `kill [PID]`). \n",
    "\n",
    "By default, the PEFT library will only save the QLoRA adapters, so we need to first load the base Mistral model from the Huggingface Hub:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "SKSnF016yRgp"
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig\n",
    "\n",
    "base_model_id = \"mistralai/Mistral-7B-v0.1\"\n",
    "bnb_config = 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",
    "base_model = AutoModelForCausalLM.from_pretrained(\n",
    "    base_model_id,  # Mistral, same as before\n",
    "    quantization_config=bnb_config,  # Same quantization config as before\n",
    "    device_map=\"auto\",\n",
    "    trust_remote_code=True,\n",
    "    use_auth_token=True\n",
    ")\n",
    "\n",
    "eval_tokenizer = AutoTokenizer.from_pretrained(\n",
    "    base_model_id,\n",
    "    add_bos_token=True,\n",
    "    trust_remote_code=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "_BxOhAiqyRgp"
   },
   "source": [
    "Now load the QLoRA adapter from the appropriate checkpoint directory, i.e. the best performing model checkpoint:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "GwsiqhWuyRgp"
   },
   "outputs": [],
   "source": [
    "from peft import PeftModel\n",
    "\n",
    "ft_model = PeftModel.from_pretrained(base_model, \"mistral-viggo-finetune/checkpoint-1000\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "lX39ibolyRgp"
   },
   "source": [
    "and run your inference!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UUehsaVNyRgp"
   },
   "source": [
    "Let's try the same `eval_prompt` and thus `model_input` as above, and see if the new finetuned model performs better."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lMkVNEUvyRgp"
   },
   "outputs": [],
   "source": [
    "eval_prompt = \"\"\"Given a target sentence construct the underlying meaning representation of the input sentence as a single function with attributes and attribute values.\n",
    "This function should describe the target string accurately and the function must be one of the following ['inform', 'request', 'give_opinion', 'confirm', 'verify_attribute', 'suggest', 'request_explanation', 'recommend', 'request_attribute'].\n",
    "The attributes must be one of the following: ['name', 'exp_release_date', 'release_year', 'developer', 'esrb', 'rating', 'genres', 'player_perspective', 'has_multiplayer', 'platforms', 'available_on_steam', 'has_linux_release', 'has_mac_release', 'specifier']\n",
    "\n",
    "### Target sentence:\n",
    "Earlier, you stated that you didn't have strong feelings about PlayStation's Little Big Adventure. Is your opinion true for all games which don't have multiplayer?\n",
    "\n",
    "### Meaning representation:\n",
    "\"\"\"\n",
    "\n",
    "model_input = tokenizer(eval_prompt, return_tensors=\"pt\").to(\"cuda\")\n",
    "\n",
    "ft_model.eval()\n",
    "with torch.no_grad():\n",
    "    print(eval_tokenizer.decode(ft_model.generate(**model_input, max_new_tokens=100)[0], skip_special_tokens=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VCJnpZoayRgq"
   },
   "source": [
    "### Sweet... it worked! The fine-tuned model now understands the meaning representation!\n",
    "\n",
    "I hope you enjoyed this tutorial on fine-tuning Mistral. If you have any questions, feel free to reach out to me on [X](https://x.com/harperscarroll) or on the [Discord channel](https://discord.gg/y9428NwTh3).\n",
    "\n",
    "🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙 🤙"
   ]
  }
 ],
 "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
}