training/fine-tune-llms-in-2024-with-trl.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# How to Fine-Tune LLMs in 2024 with TRL\n",
    "\n",
    "Large Language Models or LLMs have seen a lot of progress in the last year. We went from now ChatGPT competitor to a whole zoo of LLMs, including Meta AI's [Llama 3](https://huggingface.co/blog/llama31), Mistrals [Mistral](https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.2) & [Mixtral](https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1) models, TII [Falcon](https://huggingface.co/tiiuae/falcon-40b), and many more. \n",
    "Those LLMs can be used for a variety of tasks, including chatbots, question answering, summarization without any additional training. However, if you want to customize a model for your application. You may need to fine-tune the model on your data to achieve higher quality results than prompting or saving cost by training smaller models more efficient model.\n",
    "\n",
    " \n",
    "This blog post walks you thorugh how to fine-tune open LLMs using Hugging Face [TRL](https://huggingface.co/docs/trl/index), [Transformers](https://huggingface.co/docs/transformers/index) & [datasets](https://huggingface.co/docs/datasets/index) in 2024. In the blog, we are going to:\n",
    "\n",
    "1. Define our use case \n",
    "2. Setup development environment\n",
    "3. Create and prepare the dataset\n",
    "4. Fine-tune LLM using `trl` and the `SFTTrainer` \n",
    "5. Test and evaluate the LLM\n",
    "6. Deploy the LLM for Production\n",
    "\n",
    "_Note: This blog was created to run on consumer size GPUs (24GB), e.g. NVIDIA A10G or RTX 4090/3090, but can be easily adapted to run on bigger GPUs._\n",
    "\n",
    "\n",
    "## 1. Define our use case \n",
    "\n",
    "When fine-tuning LLMs, it is important you know your use case and the task you want to solve. This will help you to choose the right model or help you to create a dataset to fine-tune your model. If you haven't defined your use case yet. You might want to go back to the drawing board. \n",
    "I want to mention that not all use cases require fine-tuning and it is always recommended to evaluate and try out already fine-tuned models or API-based models before fine-tuning your own model. \n",
    "\n",
    "As an example, we are going to use the following use case:\n",
    "\n",
    "> We want to fine-tune a model, which can generate SQL queries based on a natural language instruction, which can then be integrated into our BI tool. The goal is to reduce the time it takes to create a SQL query and make it easier for non-technical users to create SQL queries. \n",
    "\n",
    "Text to SQL can be a good use case for fine-tuning LLMs, as it is a complex task that requires a lot of (internal) knowledge about the data and the SQL language. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Setup development environment\n",
    "\n",
    "Our first step is to install Hugging Face Libraries and Pyroch, including trl, transformers and datasets. If you haven't heard of trl yet, don't worry. It is a new library on top of transformers and datasets, which makes it easier to fine-tune, rlhf, align open LLMs. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install Pytorch & other libraries\n",
    "%pip install \"torch==2.4.0\" tensorboard\n",
    "\n",
    "# Install Hugging Face libraries\n",
    "%pip install  --upgrade \\\n",
    "  \"transformers==4.44.2\" \\\n",
    "  \"datasets==2.21.0\" \\\n",
    "  \"accelerate==0.33.0\" \\\n",
    "  \"evaluate==0.4.2\" \\\n",
    "  \"bitsandbytes==0.43.3\" \\\n",
    "  \"trl==0.9.6\" \\\n",
    "  \"peft==0.12.0\" "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you are using a GPU with Ampere architecture (e.g. NVIDIA A10G or RTX 4090/3090) or newer you can use Flash attention. Flash Attention is a an method that reorders the attention computation and leverages classical techniques (tiling, recomputation) to significantly speed it up and reduce memory usage from quadratic to linear in sequence length. The TL;DR; accelerates training up to 3x. Learn more at [FlashAttention](https://github.com/Dao-AILab/flash-attention/tree/main).\n",
    "\n",
    "_Note: If your machine has less than 96GB of RAM and lots of CPU cores, reduce the number of `MAX_JOBS`. On the `g6.2xlarge` we used `4`._"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch; assert torch.cuda.get_device_capability()[0] >= 8, 'Hardware not supported for Flash Attention'\n",
    "# install flash-attn\n",
    "!pip install ninja packaging\n",
    "!MAX_JOBS=4 pip install flash-attn --no-build-isolation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "_Installing flash attention can take quite a bit of time (10-45 minutes)._"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will use the [Hugging Face Hub](https://huggingface.co/models) as a remote model versioning service. This means we will automatically push our model, logs and information to the Hub during training. You must register on the [Hugging Face](https://huggingface.co/join) for this. After you have an account, we will use the `login` util from the `huggingface_hub` package to log into our account and store our token (access key) on the disk.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Token is valid (permission: write).\n",
      "\u001b[1m\u001b[31mCannot authenticate through git-credential as no helper is defined on your machine.\n",
      "You might have to re-authenticate when pushing to the Hugging Face Hub.\n",
      "Run the following command in your terminal in case you want to set the 'store' credential helper as default.\n",
      "\n",
      "git config --global credential.helper store\n",
      "\n",
      "Read https://git-scm.com/book/en/v2/Git-Tools-Credential-Storage for more details.\u001b[0m\n",
      "Token has not been saved to git credential helper.\n",
      "Your token has been saved to /home/ubuntu/.cache/huggingface/token\n",
      "Login successful\n"
     ]
    }
   ],
   "source": [
    "from huggingface_hub import login\n",
    "\n",
    "login(\n",
    "  token=\"\", # ADD YOUR TOKEN HERE\n",
    "  add_to_git_credential=True\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Create and prepare the dataset\n",
    "\n",
    "Once you have determined that fine-tuning is the right solution we need to create a dataset to fine-tune our model. The dataset should be a diverse set of demonstrations of the task you want to solve. There are several ways to create such a dataset, including:\n",
    "* Using existing open-source datasets, e.g., [Spider](https://huggingface.co/datasets/spider)\n",
    "* Using LLMs to create synthetically datasets, e.g., [Alpaca](https://huggingface.co/datasets/tatsu-lab/alpaca)\n",
    "* Using Humans to create datasets, e.g., [Dolly](https://huggingface.co/datasets/databricks/databricks-dolly-15k).\n",
    "* Using a combination of the above methods, e.g., [Orca](https://huggingface.co/datasets/Open-Orca/OpenOrca)\n",
    "\n",
    "Each of the methods has its own advantages and disadvantages and depends on the budget, time, and quality requirements. For example, using an existing dataset is the easiest but might not be tailored to your specific use case, while using humans might be the most accurate but can be time-consuming and expensive. It is also possible to combine several methods to create an instruction dataset, as shown in [Orca: Progressive Learning from Complex Explanation Traces of GPT-4.](https://arxiv.org/abs/2306.02707)\n",
    "\n",
    "In our example we will use an already existing dataset called [sql-create-context](https://huggingface.co/datasets/b-mc2/sql-create-context), which contains samples of natural language instructions, schema definitions and the corresponding SQL query.\n",
    "\n",
    "With the latest release of `trl` we now support popular instruction and conversation dataset formats. This means we only need to convert our dataset to one of the supported formats and `trl` will take care of the rest. Those formats include:\n",
    "* conversational format\n",
    "```json\n",
    "{\"messages\": [{\"role\": \"system\", \"content\": \"You are...\"}, {\"role\": \"user\", \"content\": \"...\"}, {\"role\": \"assistant\", \"content\": \"...\"}]}\n",
    "{\"messages\": [{\"role\": \"system\", \"content\": \"You are...\"}, {\"role\": \"user\", \"content\": \"...\"}, {\"role\": \"assistant\", \"content\": \"...\"}]}\n",
    "{\"messages\": [{\"role\": \"system\", \"content\": \"You are...\"}, {\"role\": \"user\", \"content\": \"...\"}, {\"role\": \"assistant\", \"content\": \"...\"}]}\n",
    "```\n",
    "* instruction format\n",
    "\n",
    "```json\n",
    "{\"prompt\": \"<prompt text>\", \"completion\": \"<ideal generated text>\"}\n",
    "{\"prompt\": \"<prompt text>\", \"completion\": \"<ideal generated text>\"}\n",
    "{\"prompt\": \"<prompt text>\", \"completion\": \"<ideal generated text>\"}\n",
    "```"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In our example we are going to load our open-source dataset using the 🤗 Datasets library and then convert it into the the conversational format, where we include the schema definition in the system message for our assistant. We'll then save the dataset as jsonl file, which we can then use to fine-tune our model. We are randomly downsampling the dataset to only 10,000 samples.\n",
    "\n",
    "_Note: This step can be different for your use case. For example, if you have already a dataset from, e.g. working with OpenAI, you can skip this step and go directly to the fine-tuning step._"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "633ff23ca89d45d080ad6ee0928b3e33",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Downloading readme:   0%|          | 0.00/4.43k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5a19a3f0865645f29d87d77d98e3d045",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/12500 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{'content': 'You are an text to SQL query translator. Users will ask you questions in English and you will generate a SQL query based on the provided SCHEMA.\\nSCHEMA:\\nCREATE TABLE table_name_32 (date VARCHAR, attendance VARCHAR)', 'role': 'system'}, {'content': 'On what Date was the Attendance 73,405?', 'role': 'user'}, {'content': 'SELECT date FROM table_name_32 WHERE attendance = \"73,405\"', 'role': 'assistant'}]\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b9f50c05f4874e67916c2987fa067fdc",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Creating json from Arrow format:   0%|          | 0/10 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "85650bb79acf4da1add5f46cc088564b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Creating json from Arrow format:   0%|          | 0/3 [00:00<?, ?ba/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "1187015"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "# Convert dataset to OAI messages\n",
    "system_message = \"\"\"You are an text to SQL query translator. Users will ask you questions in English and you will generate a SQL query based on the provided SCHEMA.\n",
    "SCHEMA:\n",
    "{schema}\"\"\"\n",
    "\n",
    "def create_conversation(sample):\n",
    "  return {\n",
    "    \"messages\": [\n",
    "      {\"role\": \"system\", \"content\": system_message.format(schema=sample[\"context\"])},\n",
    "      {\"role\": \"user\", \"content\": sample[\"question\"]},\n",
    "      {\"role\": \"assistant\", \"content\": sample[\"answer\"]}\n",
    "    ]\n",
    "  }  \n",
    "\n",
    "# Load dataset from the hub\n",
    "dataset = load_dataset(\"b-mc2/sql-create-context\", split=\"train\")\n",
    "dataset = dataset.shuffle().select(range(12500))\n",
    "\n",
    "# Convert dataset to OAI messages\n",
    "dataset = dataset.map(create_conversation, remove_columns=dataset.features,batched=False)\n",
    "# split dataset into 10,000 training samples and 2,500 test samples\n",
    "dataset = dataset.train_test_split(test_size=2500/12500)\n",
    "\n",
    "print(dataset[\"train\"][345][\"messages\"])\n",
    "\n",
    "# save datasets to disk \n",
    "dataset[\"train\"].to_json(\"train_dataset.json\", orient=\"records\")\n",
    "dataset[\"test\"].to_json(\"test_dataset.json\", orient=\"records\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Fine-tune LLM using `trl` and the `SFTTrainer` \n",
    "\n",
    "We are now ready to fine-tune our model. We will use the [SFTTrainer](https://huggingface.co/docs/trl/sft_trainer) from `trl` to fine-tune our model. The `SFTTrainer` makes it straightfoward to supervise fine-tune open LLMs. The `SFTTrainer` is a subclass of the `Trainer` from the `transformers` library and supports all the same features, including logging, evaluation, and checkpointing, but adds additiional quality of life features, including:\n",
    "* Dataset formatting, including conversational and instruction format\n",
    "* Training on completions only, ignoring prompts\n",
    "* Packing datasets for more efficient training\n",
    "* PEFT (parameter-efficient fine-tuning) support including Q-LoRA\n",
    "* Preparing the model and tokenizer for conversational fine-tuning (e.g. adding special tokens)\n",
    " \n",
    "We will use the dataset formatting, packing and PEFT features in our example. As peft method we will use [QLoRA](https://arxiv.org/abs/2305.14314) a technique to reduce the memory footprint of large language models during finetuning, without sacrificing performance by using quantization. If you want to learn more about QLoRA and how it works, check out [Making LLMs even more accessible with bitsandbytes, 4-bit quantization and QLoRA](https://huggingface.co/blog/4bit-transformers-bitsandbytes) blog post.\n",
    "\n",
    "Now, lets get started! 🚀 \n",
    "\n",
    "First, we need to load our dataset from disk. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "# Load jsonl data from disk\n",
    "dataset = load_dataset(\"json\", data_files=\"train_dataset.json\", split=\"train\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we will load our LLM. For our use case we are going to use Llama 3.1 8B. \n",
    "But we can easily swap out the model for another model, e.g. [Mistral](https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.2) or [Mixtral](https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1) models, TII [Falcon](https://huggingface.co/tiiuae/falcon-40b), or any other LLMs by changing our `model_id` variable. We will use bitsandbytes to quantize our model to 4-bit.\n",
    "\n",
    "_Note: Be aware the bigger the model the more memory it will require. In our example we will use the 8B version, which can be tuned on 24GB GPUs. If you have a smaller GPU._\n",
    "\n",
    "\n",
    "Correctly, preparing the LLM and Tokenizer for training chat/conversational models is crucial. We need to add new special tokens to the tokenizer and model and teach to understand the different roles in a conversation. In `trl` we have a convinient method called [setup_chat_format](https://huggingface.co/docs/trl/main/en/sft_trainer#add-special-tokens-for-chat-format), which:\n",
    "* Adds special tokens to the tokenizer, e.g. `<|im_start|>` and `<|im_end|>`, to indicate the start and end of a conversation.\n",
    "* Resizes the model’s embedding layer to accommodate the new tokens.\n",
    "* Sets the `chat_template` of the tokenizer, which is used to format the input data into a chat-like format. The default is `chatml` from OpenAI."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9c4b2fd01614421186848976aa440952",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/4 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import torch\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig\n",
    "from trl import setup_chat_format\n",
    "\n",
    "# Hugging Face model id\n",
    "model_id = \"meta-llama/Meta-Llama-3.1-8B\" # or `mistralai/Mistral-7B-v0.1`\n",
    "\n",
    "# BitsAndBytesConfig int-4 config\n",
    "bnb_config = BitsAndBytesConfig(\n",
    "    load_in_4bit=True, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type=\"nf4\", bnb_4bit_compute_dtype=torch.bfloat16\n",
    ")\n",
    "\n",
    "# Load model and tokenizer\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    model_id,\n",
    "    device_map=\"auto\",\n",
    "    attn_implementation=\"flash_attention_2\",\n",
    "    torch_dtype=torch.bfloat16,\n",
    "    quantization_config=bnb_config\n",
    ")\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
    "tokenizer.padding_side = 'right' # to prevent warnings\n",
    "\n",
    "# # set chat template to OAI chatML, remove if you start from a fine-tuned model\n",
    "model, tokenizer = setup_chat_format(model, tokenizer)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `SFTTrainer`  supports a native integration with `peft`, which makes it super easy to efficiently tune LLMs using, e.g. QLoRA. We only need to create our `LoraConfig` and provide it to the trainer. Our `LoraConfig` parameters are defined based on the [qlora paper](https://arxiv.org/pdf/2305.14314.pdf) and sebastian's [blog post](https://magazine.sebastianraschka.com/p/practical-tips-for-finetuning-llms)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from peft import LoraConfig\n",
    "\n",
    "# LoRA config based on QLoRA paper & Sebastian Raschka experiment\n",
    "peft_config = LoraConfig(\n",
    "        lora_alpha=128,\n",
    "        lora_dropout=0.05,\n",
    "        r=256,\n",
    "        bias=\"none\",\n",
    "        target_modules=\"all-linear\",\n",
    "        task_type=\"CAUSAL_LM\", \n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before we can start our training we need to define the hyperparameters (`TrainingArguments`) we want to use."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import TrainingArguments\n",
    "\n",
    "args = TrainingArguments(\n",
    "    output_dir=\"code-llama-3-1-8b-text-to-sql\", # directory to save and repository id\n",
    "    num_train_epochs=3,                     # number of training epochs\n",
    "    per_device_train_batch_size=1,          # batch size per device during training\n",
    "    gradient_accumulation_steps=8,          # number of steps before performing a backward/update pass\n",
    "    gradient_checkpointing=True,            # use gradient checkpointing to save memory\n",
    "    optim=\"adamw_torch_fused\",              # use fused adamw optimizer\n",
    "    logging_steps=10,                       # log every 10 steps\n",
    "    save_strategy=\"epoch\",                  # save checkpoint every epoch\n",
    "    learning_rate=2e-4,                     # learning rate, based on QLoRA paper\n",
    "    bf16=True,                              # use bfloat16 precision\n",
    "    tf32=True,                              # use tf32 precision\n",
    "    max_grad_norm=0.3,                      # max gradient norm based on QLoRA paper\n",
    "    warmup_ratio=0.03,                      # warmup ratio based on QLoRA paper\n",
    "    lr_scheduler_type=\"constant\",           # use constant learning rate scheduler\n",
    "    push_to_hub=True,                       # push model to hub\n",
    "    report_to=\"tensorboard\",                # report metrics to tensorboard\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We now have every building block we need to create our `SFTTrainer` to start then training our model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/envs/pytorch/lib/python3.10/site-packages/huggingface_hub/utils/_deprecation.py:100: FutureWarning: Deprecated argument(s) used in '__init__': max_seq_length, packing, dataset_kwargs. Will not be supported from version '1.0.0'.\n",
      "\n",
      "Deprecated positional argument(s) used in SFTTrainer, please use the SFTConfig to set these arguments instead.\n",
      "  warnings.warn(message, FutureWarning)\n",
      "/opt/conda/envs/pytorch/lib/python3.10/site-packages/trl/trainer/sft_trainer.py:192: UserWarning: You passed a `packing` argument to the SFTTrainer, the value you passed will override the one in the `SFTConfig`.\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/pytorch/lib/python3.10/site-packages/trl/trainer/sft_trainer.py:280: UserWarning: You passed a `max_seq_length` argument to the SFTTrainer, the value you passed will override the one in the `SFTConfig`.\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/pytorch/lib/python3.10/site-packages/trl/trainer/sft_trainer.py:366: UserWarning: You passed a `dataset_kwargs` argument to the SFTTrainer, the value you passed will override the one in the `SFTConfig`.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "from trl import SFTTrainer\n",
    "\n",
    "max_seq_length = 2048 # max sequence length for model and packing of the dataset\n",
    "\n",
    "trainer = SFTTrainer(\n",
    "    model=model,\n",
    "    args=args,\n",
    "    train_dataset=dataset,\n",
    "    peft_config=peft_config,\n",
    "    max_seq_length=max_seq_length,\n",
    "    tokenizer=tokenizer,\n",
    "    packing=True,\n",
    "    dataset_kwargs={\n",
    "        \"add_special_tokens\": False,  # We template with special tokens\n",
    "        \"append_concat_token\": False, # No need to add additional separator token\n",
    "    }\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Start training our model by calling the `train()` method on our `Trainer` instance. This will start the training loop and train our model for 3 epochs. Since we are using a PEFT method, we will only save the adapted model weights and not the full model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`.\n",
      "/opt/conda/envs/pytorch/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py:600: UserWarning: torch.utils.checkpoint: the use_reentrant parameter should be passed explicitly. In version 2.4 we will raise an exception if use_reentrant is not passed. use_reentrant=False is recommended, but if you need to preserve the current default behavior, you can pass use_reentrant=True. Refer to docs for more details on the differences between the two variants.\n",
      "  return fn(*args, **kwargs)\n",
      "The input hidden states seems to be silently casted in float32, this might be related to the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in torch.bfloat16.\n",
      "/opt/conda/envs/pytorch/lib/python3.10/site-packages/torch/utils/checkpoint.py:295: FutureWarning: `torch.cpu.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cpu', args...)` instead.\n",
      "  with torch.enable_grad(), device_autocast_ctx, torch.cpu.amp.autocast(**ctx.cpu_autocast_kwargs):  # type: ignore[attr-defined]\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='186' max='186' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [186/186 2:19:28, Epoch 2/3]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>10</td>\n",
       "      <td>0.912000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>20</td>\n",
       "      <td>0.635400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>30</td>\n",
       "      <td>0.605700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>40</td>\n",
       "      <td>0.583800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>50</td>\n",
       "      <td>0.568800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>60</td>\n",
       "      <td>0.554400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>70</td>\n",
       "      <td>0.508300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>80</td>\n",
       "      <td>0.484600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>90</td>\n",
       "      <td>0.475800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>100</td>\n",
       "      <td>0.484700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>110</td>\n",
       "      <td>0.481500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>120</td>\n",
       "      <td>0.478700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>130</td>\n",
       "      <td>0.436100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>140</td>\n",
       "      <td>0.391300</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>150</td>\n",
       "      <td>0.396100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>160</td>\n",
       "      <td>0.396200</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>170</td>\n",
       "      <td>0.404600</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>180</td>\n",
       "      <td>0.408300</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/envs/pytorch/lib/python3.10/site-packages/peft/utils/save_and_load.py:232: UserWarning: Setting `save_embedding_layers` to `True` as the embedding layer has been resized during finetuning.\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/pytorch/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py:600: UserWarning: torch.utils.checkpoint: the use_reentrant parameter should be passed explicitly. In version 2.4 we will raise an exception if use_reentrant is not passed. use_reentrant=False is recommended, but if you need to preserve the current default behavior, you can pass use_reentrant=True. Refer to docs for more details on the differences between the two variants.\n",
      "  return fn(*args, **kwargs)\n",
      "/opt/conda/envs/pytorch/lib/python3.10/site-packages/torch/utils/checkpoint.py:295: FutureWarning: `torch.cpu.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cpu', args...)` instead.\n",
      "  with torch.enable_grad(), device_autocast_ctx, torch.cpu.amp.autocast(**ctx.cpu_autocast_kwargs):  # type: ignore[attr-defined]\n",
      "/opt/conda/envs/pytorch/lib/python3.10/site-packages/peft/utils/save_and_load.py:232: UserWarning: Setting `save_embedding_layers` to `True` as the embedding layer has been resized during finetuning.\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/pytorch/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py:600: UserWarning: torch.utils.checkpoint: the use_reentrant parameter should be passed explicitly. In version 2.4 we will raise an exception if use_reentrant is not passed. use_reentrant=False is recommended, but if you need to preserve the current default behavior, you can pass use_reentrant=True. Refer to docs for more details on the differences between the two variants.\n",
      "  return fn(*args, **kwargs)\n",
      "/opt/conda/envs/pytorch/lib/python3.10/site-packages/torch/utils/checkpoint.py:295: FutureWarning: `torch.cpu.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cpu', args...)` instead.\n",
      "  with torch.enable_grad(), device_autocast_ctx, torch.cpu.amp.autocast(**ctx.cpu_autocast_kwargs):  # type: ignore[attr-defined]\n",
      "/opt/conda/envs/pytorch/lib/python3.10/site-packages/peft/utils/save_and_load.py:232: UserWarning: Setting `save_embedding_layers` to `True` as the embedding layer has been resized during finetuning.\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/pytorch/lib/python3.10/site-packages/peft/utils/save_and_load.py:232: UserWarning: Setting `save_embedding_layers` to `True` as the embedding layer has been resized during finetuning.\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/pytorch/lib/python3.10/site-packages/peft/utils/save_and_load.py:232: UserWarning: Setting `save_embedding_layers` to `True` as the embedding layer has been resized during finetuning.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f5631d039bde4c6e87f2716bf20a7795",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "events.out.tfevents.1724931649.ip-172-31-27-157.2377.0:   0%|          | 0.00/10.1k [00:00<?, ?B/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# start training, the model will be automatically saved to the hub and the output directory\n",
    "trainer.train()\n",
    "\n",
    "# save model \n",
    "trainer.save_model()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The training with Flash Attention for 3 epochs with a dataset of 10k samples took 02:05:58 on a `g6.2xlarge`. The instance costs `1,212$/h` which brings us to a total cost of only `1.8$`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# free the memory again\n",
    "del model\n",
    "del trainer\n",
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Merge LoRA adapter in to the original model\n",
    "\n",
    "When using QLoRA, we only train adapters and not the full model. This means when saving the model during training we only save the adapter weights and not the full model. If you want to save the full model, which makes it easier to use with Text Generation Inference you can merge the adapter weights into the model weights using the `merge_and_unload` method and then save the model with the `save_pretrained` method. This will save a default model, which can be used for inference.\n",
    "\n",
    "_Note: This requires > 30GB CPU Memory._"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "07e8fc0add14407c8c32c88abe998e78",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/4 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#### COMMENT IN TO MERGE PEFT AND BASE MODEL ####\n",
    "from peft import AutoPeftModelForCausalLM\n",
    "\n",
    "# Load PEFT model on CPU\n",
    "model = AutoPeftModelForCausalLM.from_pretrained(\n",
    "    args.output_dir,\n",
    "    torch_dtype=torch.float16,\n",
    "    low_cpu_mem_usage=True,\n",
    ")  \n",
    "# Merge LoRA and base model and save\n",
    "merged_model = model.merge_and_unload()\n",
    "merged_model.save_pretrained(args.output_dir,safe_serialization=True, max_shard_size=\"2GB\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Test Model and run Inference\n",
    "\n",
    "After the training is done we want to evaluate and test our model. We will load different samples from the original dataset and evaluate the model on those samples, using a simple loop and accuracy as our metric. \n",
    "\n",
    "_Note: Evaluating Generative AI models is not a trivial task since 1 input can have multiple correct outputs. If you want to learn more about evaluating generative models, check out [Evaluate LLMs and RAG a practical example using Langchain and Hugging Face](https://www.philschmid.de/evaluate-llm) blog post._\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/envs/pytorch/lib/python3.10/site-packages/torchvision/io/image.py:13: UserWarning: Failed to load image Python extension: '/opt/conda/envs/pytorch/lib/python3.10/site-packages/torchvision/image.so: undefined symbol: _ZN3c1017RegisterOperatorsD1Ev'If you don't plan on using image functionality from `torchvision.io`, you can ignore this warning. Otherwise, there might be something wrong with your environment. Did you have `libjpeg` or `libpng` installed before building `torchvision` from source?\n",
      "  warn(\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "89b45d7199c04dbea843232c087488fd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/4 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import torch\n",
    "from transformers import AutoTokenizer, pipeline, AutoModelForCausalLM\n",
    "\n",
    "model_id = \"./code-llama-3-1-8b-text-to-sql\"\n",
    "\n",
    "# Load Model with PEFT adapter\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "  model_id,\n",
    "  device_map=\"auto\",\n",
    "  torch_dtype=torch.float16\n",
    ")\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
    "# load into pipeline\n",
    "pipe = pipeline(\"text-generation\", model=model, tokenizer=tokenizer)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let’s load our test dataset try to generate an instruction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "3844eccf972047eb82384102948ca553",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Generating train split: 0 examples [00:00, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/envs/pytorch/lib/python3.10/site-packages/transformers/generation/configuration_utils.py:567: UserWarning: `do_sample` is set to `False`. However, `temperature` is set to `0.1` -- this flag is only used in sample-based generation modes. You should set `do_sample=True` or unset `temperature`.\n",
      "  warnings.warn(\n",
      "/opt/conda/envs/pytorch/lib/python3.10/site-packages/transformers/generation/configuration_utils.py:572: UserWarning: `do_sample` is set to `False`. However, `top_p` is set to `0.1` -- this flag is only used in sample-based generation modes. You should set `do_sample=True` or unset `top_p`.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Query:\n",
      "What is the Highest first elected year that has a district of 06.0 6, and a committee of economic matters?\n",
      "Original Answer:\n",
      "SELECT MAX(first_elected) FROM table_name_99 WHERE district = \"06.0 6\" AND committee = \"economic matters\"\n",
      "Generated Answer:\n",
      "SELECT MAX(first_elected) FROM table_name_99 WHERE district = \"06.0 6\" AND committee = \"economic matters\"\n"
     ]
    }
   ],
   "source": [
    "from datasets import load_dataset \n",
    "from random import randint\n",
    "\n",
    "\n",
    "# Load our test dataset\n",
    "eval_dataset = load_dataset(\"json\", data_files=\"test_dataset.json\", split=\"train\")\n",
    "rand_idx = randint(0, len(eval_dataset))\n",
    "\n",
    "# Test on sample \n",
    "prompt = pipe.tokenizer.apply_chat_template(eval_dataset[rand_idx][\"messages\"][:2], tokenize=False, add_generation_prompt=True)\n",
    "outputs = pipe(prompt, max_new_tokens=256, do_sample=False, temperature=0.1, top_k=50, top_p=0.1, eos_token_id=pipe.tokenizer.eos_token_id, pad_token_id=pipe.tokenizer.pad_token_id)\n",
    "\n",
    "print(f\"Query:\\n{eval_dataset[rand_idx]['messages'][1]['content']}\")\n",
    "print(f\"Original Answer:\\n{eval_dataset[rand_idx]['messages'][2]['content']}\")\n",
    "print(f\"Generated Answer:\\n{outputs[0]['generated_text'][len(prompt):].strip()}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Nice! Our model was able to generate a SQL query based on the natural language instruction. Lets evaluate our model on the full 2,500 samples of our test dataset. \n",
    "_Note: As mentioned above, evaluating generative models is not a trivial task. In our example we used the accuracy of the generated SQL based on the ground truth SQL query as our metric. An alternative way could be to automatically execute the generated SQL query and compare the results with the ground truth. This would be a more accurate metric but requires more work to setup._"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  1%|          | 9/1000 [00:12<21:30,  1.30s/it]You seem to be using the pipelines sequentially on GPU. In order to maximize efficiency please use a dataset\n",
      "100%|██████████| 1000/1000 [24:27<00:00,  1.47s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy: 80.00%\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "from tqdm import tqdm\n",
    "\n",
    "\n",
    "def evaluate(sample):\n",
    "    prompt = pipe.tokenizer.apply_chat_template(sample[\"messages\"][:2], tokenize=False, add_generation_prompt=True)\n",
    "    outputs = pipe(prompt, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95, eos_token_id=pipe.tokenizer.eos_token_id, pad_token_id=pipe.tokenizer.pad_token_id)\n",
    "    predicted_answer = outputs[0]['generated_text'][len(prompt):].strip()\n",
    "    if predicted_answer == sample[\"messages\"][2][\"content\"]:\n",
    "        return 1 \n",
    "    else:\n",
    "        return 0\n",
    "\n",
    "success_rate = []\n",
    "number_of_eval_samples = 1000\n",
    "# iterate over eval dataset and predict\n",
    "for s in tqdm(eval_dataset.shuffle().select(range(number_of_eval_samples))):\n",
    "    success_rate.append(evaluate(s))\n",
    "\n",
    "# compute accuracy\n",
    "accuracy = sum(success_rate)/len(success_rate)\n",
    "\n",
    "print(f\"Accuracy: {accuracy*100:.2f}%\")  \n",
    "        "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We evaluated our model on 1000 samples from the evaluation dataset and got an accuracy of 80.00%, which took ~25 minutes. \n",
    "This is quite good, but as mentioned you need to take this metric with a grain of salt. It would be better if we could evaluate our model by running the qureies against a real database and compare the results. Since there might be different \"correct\" SQL queries for the same instruction. There are also several ways on how we could improve the performance by using few-shot learning, using RAG, Self-healing to generate the SQL query."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Deploy the LLM for Production\n",
    "\n",
    "You can now deploy your model to production. For deploying open LLMs into production we recommend using [Text Generation Inference (TGI)](https://github.com/huggingface/text-generation-inference). TGI is a purpose-built solution for deploying and serving Large Language Models (LLMs). TGI enables high-performance text generation using Tensor Parallelism and continous batching for the most popular open LLMs, including Llama, Mistral, Mixtral, StarCoder, T5 and more. Text Generation Inference is used by companies as IBM, Grammarly, Uber, Deutsche Telekom, and many more. There are several ways to deploy your model, including:\n",
    "\n",
    " * [Deploy LLMs with Hugging Face Inference Endpoints](https://huggingface.co/blog/inference-endpoints-llm)\n",
    " * [Hugging Face LLM Inference Container for Amazon SageMaker](https://huggingface.co/blog/sagemaker-huggingface-llm)\n",
    "\n",
    "If you have docker installed you can use the following command to start the inference server. \n",
    "\n",
    "_Note: Make sure that you have enough GPU memory to run the container. Restart kernel to remove all allocated GPU memory from the notebook._ "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Unable to find image 'ghcr.io/huggingface/text-generation-inference:2.2.0' locally\n",
      "2.2.0: Pulling from huggingface/text-generation-inference\n",
      "aece8493d397: Already exists\n",
      "45f7ea5367fe: Already exists\n",
      "3d97a47c3c73: Already exists\n",
      "12cd4d19752f: Already exists\n",
      "da5a484f9d74: Already exists\n",
      "4f4fb700ef54: Already exists\n",
      "43566b48e5d6: Already exists\n",
      "f165933352a8: Already exists\n",
      "f166ffc7c7b4: Already exists\n",
      "58165ae83a0e: Already exists\n",
      "074d930e1b90: Already exists\n",
      "1033b2636622: Already exists\n",
      "e0aa534acffe: Already exists\n",
      "130989d28b48: Already exists\n",
      "a65ea9ebfaba: Already exists\n",
      "7225b2c46f88: Already exists\n",
      "43154e73908f: Already exists\n",
      "8f400e318724: Already exists\n",
      "f694acf6c40f: Already exists\n",
      "44fc79164bc4: Already exists\n",
      "8bc7c142e917: Already exists\n",
      "021f7d48bdcb: Already exists\n",
      "c9d01f7d10cc: Already exists\n",
      "400740bc31be: Already exists\n",
      "bd4b49ea4512: Already exists\n",
      "141228b9bdde: Already exists\n",
      "4f4fb700ef54: Already exists\n",
      "34d4a7457184: Already exists\n",
      "66e724dff43a: Already exists\n",
      "25c75c242d08: Already exists\n",
      "6a4be63c7e70: Already exists\n",
      "b2d83f4bca52: Already exists\n",
      "373c47aa4b50: Already exists\n",
      "4f4fb700ef54: Already exists\n",
      "Digest: sha256:d39d513f13727ffa9b6a4d0e949f36413b944aabc9a236c0aa2986c929906769\n",
      "Status: Downloaded newer image for ghcr.io/huggingface/text-generation-inference:2.2.0\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "42be7f00ddeb0a3214920a09a5ea303d8eb034942d7020155b6a6761fca87193\n"
     ]
    }
   ],
   "source": [
    "%%bash \n",
    "# model=$PWD/{args.output_dir} # path to model\n",
    "model=$(pwd)/code-llama-3-1-8b-text-to-sql # path to model\n",
    "num_shard=1             # number of shards\n",
    "max_input_length=1024   # max input length\n",
    "max_total_tokens=2048   # max total tokens\n",
    "\n",
    "docker run -d --name tgi --gpus all -ti -p 8080:80 \\\n",
    "  -e MODEL_ID=/workspace \\\n",
    "  -e NUM_SHARD=$num_shard \\\n",
    "  -e MAX_INPUT_LENGTH=$max_input_length \\\n",
    "  -e MAX_TOTAL_TOKENS=$max_total_tokens \\\n",
    "  -v $model:/workspace \\\n",
    "  ghcr.io/huggingface/text-generation-inference:2.2.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Once your container is running you can send requests using the `openai` or `huggingface_hub` sdk. Here we ll use the `openai` sdk to send a request to our inference server. If you don't have the `openai` sdk installed you can install it using `pip install openai`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Query:\n",
      "Name the first elected for kentucky 1\n",
      "Original Answer:\n",
      "SELECT first_elected FROM table_2668378_5 WHERE district = \"Kentucky 1\"\n",
      "Generated Answer:\n",
      "SELECT first_elected FROM table_2668378_5 WHERE district = \"Kentucky 1\"\n"
     ]
    }
   ],
   "source": [
    "from openai import OpenAI\n",
    "from datasets import load_dataset\n",
    "from random import randint\n",
    "\n",
    "# create client \n",
    "client = OpenAI(base_url=\"http://localhost:8080/v1\",api_key=\"-\")\n",
    "\n",
    "# Load our test dataset\n",
    "eval_dataset = load_dataset(\"json\", data_files=\"test_dataset.json\", split=\"train\")\n",
    "rand_idx = randint(0, len(eval_dataset))\n",
    "\n",
    "# Take a random sample from the dataset and remove the last message and send it to the model\n",
    "response = client.chat.completions.create(\n",
    "\tmodel=\"code-llama-3-1-8b-text-to-sql\",\n",
    "\tmessages=eval_dataset[rand_idx][\"messages\"][:2],\n",
    "\tstream=False, # no streaming\n",
    "\tmax_tokens=1024,\n",
    ")\n",
    "response = response.choices[0].message.content\n",
    "\n",
    "# Print results\n",
    "print(f\"Query:\\n{eval_dataset[rand_idx]['messages'][1]['content']}\")\n",
    "print(f\"Original Answer:\\n{eval_dataset[rand_idx]['messages'][2]['content']}\")\n",
    "print(f\"Generated Answer:\\n{response}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Awesome, Don't forget to stop your container once you are done."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tgi\n"
     ]
    }
   ],
   "source": [
    "!docker stop tgi"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "2d58e898dde0263bc564c6968b04150abacfd33eed9b19aaa8e45c040360e146"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}