LLM_finetuning_detailed.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1cc557a5-5a55-4bae-8555-a2b655abfa4a",
   "metadata": {},
   "source": [
    "SPDX-License-Identifier: Apache-2.0\n",
    "Copyright (c) 2023, Rahul Unnikrishnan Nair <rahul.unnikrishnan.nair@intel.com>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "f9ad2a68-5e1a-491d-9ad8-e0d8eaf067c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# if not installed, install the bigDL library for xpus:\n",
    "#pip install --pre --upgrade bigdl-llm[xpu]==2.4.0b20231028 -f https://developer.intel.com/ipex-whl-stable-xpu"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "21c23f99-bfcf-476b-91d5-4f2ec1178707",
   "metadata": {},
   "source": [
    "# Streamlining Text-to-SQL Tasks: Fine-Tuning LLMs with Quantized LoRA on Intel\n",
    "\n",
    "## Applying Quantized LoRA and Hugging Face Transformers on Intel Architecture\n",
    "\n",
    "This Jupyter Notebook is designed to facilitate the fine-tuning of Language Models for specialized tasks such as Text-to-SQL conversions. The notebook is suitable for AI engineers and practitioners looking to leverage the power of Intel hardware for fine-tuning models efficiently.tuning of Language Models for specialized tasks such as Text-to-SQL conversions. The notebook is suitable for AI engineers and practioners looking to leverage the power of Intel hardware for fine-tuning models efficiently.\n",
    "\n",
    "### What you will learn with this Notebook\n",
    "- Fine-tune a Language Model with either a pre-existing dataset or a custom dataset tailored to your needs.\n",
    "- Gain insights into the fine-tuning process, including how to manipulate various training parameters to optimize your model's performance.\n",
    "- Test different configurations and observe the results in real-time.\n",
    "\n",
    "### Hardware Compatibility\n",
    "- The notebook is compatible with 4th Generation Intel® Xeon® Scalable Processors, ensuring high performance for AI-related tasks.\n",
    "- Optimizations are in place for the Intel® Data Center GPU Max Series, offering advanced AI acceleration capabilities.\n",
    "\n",
    "### Fine-Tuning Methodology\n",
    "We use the Quantized Low Rank Adapter (QLoRA) method for fine-tuning, which allows for a foundational model to be enhanced with task-specific adapters efficiently. This method is chosen for its computational efficiency and the flexibility it provides in quickly adapting to different tasks, like Text-to-SQL translation.\n",
    "\n",
    "LoRA (Low-Rank Adaptation) works by inserting small trainable adapter modules between the layers of a large pretrained model like an LLM.These adapters are low-rank matrices of shape (r x d) where r is a small rank hyperparameter and d is the hidden dimension size. For example, r may be 64 while d is 1024 or larger.The low-rank structure comes from factorizing the adapter matrix into two smaller matrices L1 and L2 of shape (r x d) and (d x r). The adapter then computes L1 * L2. This factorization allows the adapter to have far fewer trainable parameters compared to the full d x d dimensions, yet still adapt the function of each layer.\n",
    "\n",
    "LoRA adapters are scaled by a parameter α during training to control their capacity. Dropout is also added to regularize them. During finetuning, only the adapter parameters L1 and L2 are updated, while the original pretrained model weights remain fixed. This allows efficiently optimizing just the small adapters rather than all the parameters.\n",
    "\n",
    "QLoRA builds on top of LoRA by first quantizing the pretrained model weights before adding adapters. For example, the weights may be quantized to 4 bits rather than 32-bit floats. This quantization shrinks the model size in memory, providing further savings. QLoRA keeps the LoRA adapters in full precision for accuracy. During training, the quantized weights are temporarily dequantized on-the-fly to compute the forward and backward passes. But only the adapter gradients are accumulated, not the base model.\n",
    "\n",
    "For a deeper understanding of LoRA, readers are encouraged to refer to the [paper](https://arxiv.org/abs/2106.09685), \"LoRA: Low-Rank Adaptation of Large Language Models\", which details the methodology and its application to models as large as GPT-3. The subsequent development, QLoRA, extends these principles with quantization techniques as detailed in the [paper](https://arxiv.org/abs/2305.14314/) \"Quantized LoRA: Scaling Up Low-Rank Adaptation to Large Language Models\".\n",
    "\n",
    "In summary, LoRA uses low-rank adapters to minimize new parameters, while QLoRA adds quantization of the base model for further memory savings. Together they enable highly efficient finetuning. With these optimizations, this notebook demonstrates how to fine-tune a language model on Intel architecture using Hugging Face Transformers and the [BigDL](https://github.com/intel-analytics/BigDL) library for efficient Text-to-SQL task adaptation.\n",
    "\n",
    "Let's begin our journey into fine-tuning a model that can adeptly translate natural language queries into SQL statements."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6362a562-f3ce-4dce-9678-ce317e554a04",
   "metadata": {},
   "source": [
    "## Initialization\n",
    "Import all the necessary libraries required for the fine-tuning process."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "fc3338e7-c919-458d-8d76-09f8aea1bca0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Defaulting to user installation because normal site-packages is not writeable\n",
      "Collecting peft\n",
      "  Using cached peft-0.5.0-py3-none-any.whl (85 kB)\n",
      "Installing collected packages: peft\n",
      "Successfully installed peft-0.5.0\n"
     ]
    }
   ],
   "source": [
    "!pip install peft --no-deps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "31eb9cf2-abcf-48f8-918b-37a18f85ac7c",
   "metadata": {},
   "outputs": [
    {
     "ename": "ImportError",
     "evalue": "cannot import name 'is_torch_xpu_available' from 'transformers.utils' (/home/ue4799802b0f9e4bc38f5f2b87003ee2/.local/lib/python3.9/site-packages/transformers/utils/__init__.py)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mImportError\u001b[0m                               Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[9], line 27\u001b[0m\n\u001b[1;32m     25\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mdatasets\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Dataset\n\u001b[1;32m     26\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mbigdl\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mllm\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mtransformers\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m AutoModelForCausalLM\n\u001b[0;32m---> 27\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mbigdl\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mllm\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mtransformers\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mqlora\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m (\n\u001b[1;32m     28\u001b[0m     get_peft_model,\n\u001b[1;32m     29\u001b[0m     prepare_model_for_kbit_training \u001b[38;5;28;01mas\u001b[39;00m prepare_model,\n\u001b[1;32m     30\u001b[0m )\n\u001b[1;32m     31\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mpeft\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m LoraConfig\n\u001b[1;32m     32\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mtransformers\u001b[39;00m\n",
      "File \u001b[0;32m~/.local/lib/python3.9/site-packages/bigdl/llm/transformers/qlora.py:213\u001b[0m\n\u001b[1;32m    209\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mpatch_prepare_ipex\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;241m*\u001b[39margs):\n\u001b[1;32m    210\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mtuple\u001b[39m(args)\n\u001b[0;32m--> 213\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtransformers\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mutils\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m (\n\u001b[1;32m    214\u001b[0m     requires_backends,\n\u001b[1;32m    215\u001b[0m     is_sagemaker_mp_enabled,\n\u001b[1;32m    216\u001b[0m     is_accelerate_available,\n\u001b[1;32m    217\u001b[0m     is_torch_xpu_available,\n\u001b[1;32m    218\u001b[0m     is_sagemaker_dp_enabled,\n\u001b[1;32m    219\u001b[0m     is_torch_tpu_available,\n\u001b[1;32m    220\u001b[0m     is_torch_npu_available)\n\u001b[1;32m    221\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtransformers\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mutils\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mgeneric\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m strtobool\n\u001b[1;32m    222\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mtransformers\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mutils\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m cached_property\n",
      "\u001b[0;31mImportError\u001b[0m: cannot import name 'is_torch_xpu_available' from 'transformers.utils' (/home/ue4799802b0f9e4bc38f5f2b87003ee2/.local/lib/python3.9/site-packages/transformers/utils/__init__.py)"
     ]
    }
   ],
   "source": [
    "import logging\n",
    "import os\n",
    "import sys\n",
    "from math import ceil\n",
    "from typing import Optional, Tuple\n",
    "import warnings\n",
    "\n",
    "warnings.filterwarnings(\n",
    "    \"ignore\", category=UserWarning, module=\"intel_extension_for_pytorch\"\n",
    ")\n",
    "warnings.filterwarnings(\n",
    "    \"ignore\", category=UserWarning, module=\"torchvision.io.image\", lineno=13\n",
    ")\n",
    "warnings.filterwarnings('ignore', message='You are using the default legacy behaviour')\n",
    "warnings.filterwarnings('ignore', category=UserWarning, message='.*Parameter.*')\n",
    "warnings.filterwarnings('ignore', category=FutureWarning, message='This implementation of AdamW is deprecated')\n",
    "os.environ['TOKENIZERS_PARALLELISM'] = 'false'\n",
    "logging.getLogger(\"transformers\").setLevel(logging.ERROR)\n",
    "logging.getLogger(\"bigdl\").setLevel(logging.ERROR)\n",
    "\n",
    "\n",
    "import torch\n",
    "import intel_extension_for_pytorch as ipex\n",
    "from datasets import load_dataset\n",
    "from datasets import Dataset\n",
    "from bigdl.llm.transformers import AutoModelForCausalLM\n",
    "from bigdl.llm.transformers.qlora import (\n",
    "    get_peft_model,\n",
    "    prepare_model_for_kbit_training as prepare_model,\n",
    ")\n",
    "from peft import LoraConfig\n",
    "import transformers\n",
    "from transformers import (\n",
    "    DataCollatorForSeq2Seq,\n",
    "    LlamaTokenizer,\n",
    "    AutoTokenizer,\n",
    "    Trainer,\n",
    "    TrainingArguments,\n",
    ")\n",
    "transformers.logging.set_verbosity_error()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "131864cb-ce5d-405b-8886-5d0f1f487c30",
   "metadata": {},
   "source": [
    "## Configuration\n",
    "\n",
    "Set up basic configurations for fine-tuning. These include the base model to use, data paths, and device settings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "88077cc2-8fcf-4128-ac44-9fb2bb327398",
   "metadata": {},
   "outputs": [],
   "source": [
    "# TODO(rahul): Move these to a config file later\n",
    "BASE_MODEL = \"openlm-research/open_llama_3b_v2\"\n",
    "DATA_PATH = \"b-mc2/sql-create-context\"\n",
    "MODEL_PATH = \"./final_model\"\n",
    "ADAPTER_PATH = \"./lora_adapters\"\n",
    "DEVICE = torch.device(\"xpu\" if torch.xpu.is_available() else \"cpu\")\n",
    "LORA_CONFIG = LoraConfig(\n",
    "    r=16,\n",
    "    lora_alpha=32,\n",
    "    target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\"],\n",
    "    lora_dropout=0.05,\n",
    "    bias=\"none\",\n",
    "    task_type=\"CAUSAL_LM\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d797b621-fd06-4ae9-a883-d7d15f16d6c4",
   "metadata": {},
   "source": [
    "## Prompt Engineering for Text-to-SQL Conversion\n",
    "\n",
    "In the realm of fine-tuning language models for specialized tasks, the design of the prompt is a pivotal aspect. The function `generate_prompt_sql` plays a crucial role in this process, particularly for the task of converting natural language questions into SQL queries.\n",
    "\n",
    "### The Importance of a Well-Defined Prompt\n",
    "\n",
    "A well-defined prompt is instrumental for several reasons:\n",
    "\n",
    "- **Clarity of Task**: It communicates to the model the exact nature of the task it is expected to perform. In this case, the model is directed to generate SQL queries based on given questions and contextual information about the database.\n",
    "\n",
    "- **Consistent Format**: Consistency in the prompt structure allows for uniform training examples. This helps the model to understand and learn the pattern of the input-to-output relationship, which is vital for generating correct SQL queries.\n",
    "\n",
    "- **Inclusion of Context**: The function ensures that the model is provided with the database context or schema, a critical piece of information required to formulate valid SQL statements.\n",
    "\n",
    "- **Result Conditioning**: By incorporating an optional expected output, the model can be fine-tuned more effectively, guiding it towards the desired output format.\n",
    "\n",
    "### Crafting the Prompt\n",
    "\n",
    "The `generate_prompt_sql` function is crafted to encapsulate the input question, the relevant database context, and the expected output in a structured and concise manner. This structure is not just a facilitator for the model's learning process but also serves as a debugging tool, allowing for easier inspection of how the model processes and responds to the input.\n",
    "\n",
    "With this function, we are not just fine-tuning a model; we are engineering a precise and efficient interaction between the user's natural language queries and the model's SQL generation capabilities. This forms the bedrock of a reliable text-to-SQL conversion system, tailored to deliver accurate and useful SQL queries in response to diverse and complex questions.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d8ecd7df-b7ce-48b0-ba9f-04b7ec0c1cf4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_prompt_sql(input_question, context, output=\"\"):\n",
    "    \"\"\"\n",
    "    Generates a prompt for fine-tuning the LLM model for text-to-SQL tasks.\n",
    "\n",
    "    Parameters:\n",
    "        input_question (str): The input text or question to be converted to SQL.\n",
    "        context (str): The schema or context in which the SQL query operates.\n",
    "        output (str, optional): The expected SQL query as the output.\n",
    "\n",
    "    Returns:\n",
    "        str: A formatted string serving as the prompt for the fine-tuning task.\n",
    "    \"\"\"\n",
    "    return f\"\"\"You are a powerful text-to-SQL model. Your job is to answer questions about a database. You are given a question and context regarding one or more tables. \n",
    "\n",
    "You must output the SQL query that answers the question.\n",
    "\n",
    "### Input:\n",
    "{input_question}\n",
    "\n",
    "### Context:\n",
    "{context}\n",
    "\n",
    "### Response:\n",
    "{output}\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3777656b-74f4-4ebc-b741-a5d50bc6e79a",
   "metadata": {},
   "source": [
    "## FineTuner\n",
    "\n",
    "The `FineTuner` class encapsulates the entire process of fine-tuning large language models (LLMs) for specialized tasks such as text-to-SQL conversion. Its design adheres to best practices in machine learning and leverages the powerful features of pre-trained models, quantization techniques, and advanced tokenization processes.\n",
    "\n",
    "## Model Loading and Configuration\n",
    "When initializing the `FineTuner`, we load the base model specified by the `base_model_id`. Crucially, we use the `load_in_low_bit` parameter with the value \"nf4\", which allows the model to load in a 4-bit format, significantly reducing memory footprint while maintaining performance. The `optimize_model` flag is set to `False` to ensure that the original architecture of the model is preserved.\n",
    "\n",
    "The model is also set to use a `torch_dtype` of `torch.float16`, enabling mixed-precision training for faster computation and lower memory usage. By specifying `modules_to_not_convert`, we can control which parts of the model are kept in full precision, which is critical for maintaining certain functionalities, like the `lm_head`, at maximum fidelity.\n",
    "\n",
    "## Tokenization Strategy\n",
    "The tokenization process is tailored to the type of model being fine-tuned. For instance, if we are working with a Llama model, we utilize a `LlamaTokenizer` to ensure compatibility with the model's expected input format. For other models, a generic `AutoTokenizer` is used.\n",
    "\n",
    "We configure the tokenizer to pad from the left side (`padding_side=\"left\"`) and set the pad token ID to 0, which is a common practice for certain language models that are sensitive to the position and order of tokens.\n",
    "\n",
    "## Data Tokenization and Preparation\n",
    "The `tokenize_batch` method is where the fine-tuner ingests raw text data and converts it into a format suitable for training the model. This method handles the addition of end-of-sequence tokens, truncation to a specified `cutoff_len`, and conditioning on the input for training.\n",
    "\n",
    "## Dataset Handling\n",
    "`prepare_data` manages the splitting of data into training and validation sets, applying the `tokenize_batch` transformation to each entry. This ensures that our datasets are ready for input into the model, with all necessary tokenization applied.\n",
    "\n",
    "## Training Process\n",
    "Finally, the `train_model` method orchestrates the training process, setting up the `Trainer` with the correct datasets, training arguments, and data collator. It ensures that caching is disabled (`use_cache = False`) for the model during training, which can be an essential step for some LLMs that have large memory footprints.\n",
    "\n",
    "The fine-tuning process is encapsulated within the `finetune` method, which strings together all the previous steps into a coherent pipeline, from model setup to training execution.\n",
    "\n",
    "By abstracting the fine-tuning process into a class with clear methods for each step, we enable a modular and understandable approach to enhancing LLMs for specific tasks. This modularity not only makes the codebase maintainable but also provides clear points for customization and optimization for different hardware setups, such as Intel® XPU.\n",
    "\n",
    "The key steps to using QLoRA for efficient finetuning are:\n",
    "\n",
    "- Load a pretrained model like LLaMA and initialize it in low precision mode by setting load_in_low_bit=\"nf4\". This will load the model with weights quantized to 4-bit NormalFloat (nf4).\n",
    "- Prepare the quantized model for finetuning with prepare_model(model). This handles quantizing the model weights into blocks and computing quantization constants.\n",
    "- Add LoRA adapters to the model using get_peft_model(model, config). The config defines the LoRA hyperparameters like adapter size, dropout, etc.\n",
    "- Finetune the model by passing gradients only through the adapters. The base model weights remain fixed in low precision.\n",
    "\n",
    "Looking at the code:\n",
    "\n",
    "- AutoModelForCausalLM is loaded with load_in_low_bit=\"nf4\" to initialize in 4-bit.\n",
    "- The model is prepared via prepare_model() which quantizes weights into blocks.\n",
    "- LoRA adapters are added with get_peft_model() using the provided config.\n",
    "- Trainer finetunes the model, with gradients only flowing into the adapters.\n",
    "\n",
    "So in summary, we leverage QLoRA in BigDL to load the base LLM in low precision, inject adapters, and efficiently finetune by optimizing just the adapters end-to-end while keeping the base model fixed. This unlocks huge memory savings, allowing us to adapt giant models on ordinary GPUs.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4d8f4cb8-0da5-4572-bd07-bdae1db897a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "class FineTuner:\n",
    "    \"\"\"A class to handle the fine-tuning of LLM models.\"\"\"\n",
    "\n",
    "    def __init__(self, base_model_id: str, model_path: str, device: torch.device):\n",
    "        \"\"\"\n",
    "        Initialize the FineTuner with base model, model path, and device.\n",
    "\n",
    "        Parameters:\n",
    "            base_model_id (str): Id of pre-trained model to use for fine-tuning.\n",
    "            model_path (str): Path to save the fine-tuned model.\n",
    "            device (torch.device): Device to run the model on.\n",
    "        \"\"\"\n",
    "        self.base_model_id = base_model_id\n",
    "        self.model_path = model_path\n",
    "        self.device = device\n",
    "\n",
    "    def setup_models(self):\n",
    "        \"\"\"Downloads the pre-trained model and tokenizer based on the given base model ID.\"\"\"\n",
    "        try:\n",
    "            self.model = AutoModelForCausalLM.from_pretrained(\n",
    "                self.base_model_id,\n",
    "                load_in_low_bit=\"nf4\",\n",
    "                optimize_model=False,\n",
    "                torch_dtype=torch.float16,\n",
    "                modules_to_not_convert=[\"lm_head\"],\n",
    "            )\n",
    "            # Choose the appropriate tokenizer based on the model name\n",
    "            if 'llama' in self.base_model_id.lower():\n",
    "                self.tokenizer = LlamaTokenizer.from_pretrained(self.base_model_id)\n",
    "            else:\n",
    "                self.tokenizer = AutoTokenizer.from_pretrained(self.base_model_id)\n",
    "            self.tokenizer.pad_token_id = 0\n",
    "            self.tokenizer.padding_side = \"left\"\n",
    "\n",
    "        except Exception as e:\n",
    "            logging.error(f\"Error in downloading models: {e}\")\n",
    "\n",
    "    def tokenize_batch(\n",
    "        self, data_points, add_eos_token=True, train_on_inputs=False, cutoff_len=512\n",
    "    ) -> dict:\n",
    "        \"\"\"\n",
    "        Tokenizes a batch of SQL related data points consisting of questions, context, and answers.\n",
    "\n",
    "        Parameters:\n",
    "            data_points (dict): A batch from the dataset containing 'question', 'context', and 'answer'.\n",
    "            add_eos_token (bool): Whether to add an EOS token at the end of each tokenized sequence.\n",
    "            cutoff_len (int): The maximum length for each tokenized sequence.\n",
    "\n",
    "        Returns:\n",
    "            dict: A dictionary containing tokenized 'input_ids', 'attention_mask', and 'labels'.\n",
    "        \"\"\"\n",
    "        try:\n",
    "            question = data_points[\"question\"]\n",
    "            context = data_points[\"context\"]\n",
    "            answer = data_points[\"answer\"]\n",
    "            if train_on_inputs:\n",
    "                user_prompt = generate_prompt_sql(question, context)\n",
    "                tokenized_user_prompt = self.tokenizer(\n",
    "                    user_prompt,\n",
    "                    truncation=True,\n",
    "                    max_length=cutoff_len,\n",
    "                    padding=False,\n",
    "                    return_tensors=None,\n",
    "                )\n",
    "                user_prompt_len = len(tokenized_user_prompt[\"input_ids\"])\n",
    "                if add_eos_token:\n",
    "                    user_prompt_len -= 1\n",
    "\n",
    "            combined_text = generate_prompt_sql(question, context, answer)\n",
    "            tokenized = self.tokenizer(\n",
    "                combined_text,\n",
    "                truncation=True,\n",
    "                max_length=cutoff_len,\n",
    "                padding=False,\n",
    "                return_tensors=None,\n",
    "            )\n",
    "            if (\n",
    "                tokenized[\"input_ids\"][-1] != self.tokenizer.eos_token_id\n",
    "                and add_eos_token\n",
    "                and len(tokenized[\"input_ids\"]) < cutoff_len\n",
    "            ):\n",
    "                tokenized[\"input_ids\"].append(self.tokenizer.eos_token_id)\n",
    "                tokenized[\"attention_mask\"].append(1)\n",
    "            tokenized[\"labels\"] = tokenized[\"input_ids\"].copy()\n",
    "            if train_on_inputs:\n",
    "                tokenized[\"labels\"] = [-100] * user_prompt_len + tokenized[\"labels\"][\n",
    "                    user_prompt_len:\n",
    "                ]\n",
    "\n",
    "            return tokenized\n",
    "        except Exception as e:\n",
    "            logging.error(\n",
    "                f\"Error in batch tokenization: {e}, Line: {e.__traceback__.tb_lineno}\"\n",
    "            )\n",
    "            raise e\n",
    "\n",
    "    def prepare_data(self, data, val_set_size=100) -> Dataset:\n",
    "        \"\"\"Prepare training and validation datasets.\"\"\"\n",
    "        try:\n",
    "            train_val_split = data[\"train\"].train_test_split(\n",
    "                test_size=val_set_size, shuffle=True, seed=42\n",
    "            )\n",
    "            train_data = train_val_split[\"train\"].shuffle().map(self.tokenize_batch)\n",
    "            val_data = train_val_split[\"test\"].shuffle().map(self.tokenize_batch)\n",
    "            return train_data, val_data\n",
    "        except Exception as e:\n",
    "            logging.error(\n",
    "                f\"Error in preparing data: {e}, Line: {e.__traceback__.tb_lineno}\"\n",
    "            )\n",
    "            raise e\n",
    "\n",
    "    def train_model(self, train_data, val_data, training_args):\n",
    "        \"\"\"\n",
    "        Fine-tune the model with the given training and validation data.\n",
    "\n",
    "        Parameters:\n",
    "            train_data (Dataset): Training data.\n",
    "            val_data (Optional[Dataset]): Validation data.\n",
    "            training_args (TrainingArguments): Training configuration.\n",
    "        \"\"\"\n",
    "        try:\n",
    "            self.model = self.model.to(DEVICE)\n",
    "            self.model = prepare_model(self.model)\n",
    "            self.model = get_peft_model(self.model, LORA_CONFIG)\n",
    "            trainer = Trainer(\n",
    "                model=self.model,\n",
    "                train_dataset=train_data,\n",
    "                eval_dataset=val_data,\n",
    "                args=training_args,\n",
    "                data_collator=DataCollatorForSeq2Seq(\n",
    "                    self.tokenizer,\n",
    "                    pad_to_multiple_of=8,\n",
    "                    return_tensors=\"pt\",\n",
    "                    padding=True,\n",
    "                ),\n",
    "            )\n",
    "            self.model.config.use_cache = False\n",
    "            trainer.train()\n",
    "            self.model.save_pretrained(self.model_path)\n",
    "        except Exception as e:\n",
    "            logging.error(f\"Error in model training: {e}\")\n",
    "\n",
    "    def finetune(self, data_path, training_args):\n",
    "        \"\"\"\n",
    "        Execute the fine-tuning pipeline.\n",
    "\n",
    "        Parameters:\n",
    "            data_path (str): Path to the data for fine-tuning.\n",
    "            training_args (TrainingArguments): Training configuration.\n",
    "        \"\"\"\n",
    "        try:\n",
    "            self.setup_models()\n",
    "            data = load_dataset(data_path)\n",
    "            train_data, val_data = self.prepare_data(data)\n",
    "            self.train_model(train_data, val_data, training_args)\n",
    "        except KeyboardInterrupt:\n",
    "            print(\"Interrupt received, saving model...\")\n",
    "            self.model.save_pretrained(f\"{self.model_path}_interrupted\")\n",
    "            print(f\"Model saved to {self.model_path}_interrupted\")\n",
    "            sys.exit(0)\n",
    "        except Exception as e:\n",
    "            logging.error(f\"Error in fintuning: {e}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34af4187-d362-49b8-bdbc-e8d2a7ae0ac0",
   "metadata": {},
   "source": [
    "The `lets_finetune` function serves as the orchestrator for the fine-tuning process, providing a high-level interface to initiate training with a set of defined parameters. It allows users to specify the device for training, the model to fine-tune, batch size, warm-up steps, learning rate, and the maximum number of steps to take during training. By configuring these parameters, users can tailor the fine-tuning process to fit the specific needs of their dataset and computational resources.\n",
    "\n",
    "Upon invocation, the function sets up the fine-tuning environment, initializes the `FineTuner` class with the selected model and device, and defines the training arguments that control the behavior of the training loop. These arguments include strategies for saving checkpoints, evaluation frequency, learning rate scheduling, precision training settings, and more. With bf16 precision and efficient batching, the function is geared towards achieving a balance between training speed and memory usage.\n",
    "\n",
    "Once the setup is complete, the `finetuner.finetune` method is called to begin the actual fine-tuning task using the provided dataset. If any errors occur during this process, they are logged for troubleshooting. This function is a crucial component of the fine-tuning workflow, encapsulating the complexity of the training setup and execution.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "67a39122-eafc-483c-ad81-9c55de15e936",
   "metadata": {},
   "outputs": [],
   "source": [
    "def lets_finetune(device=DEVICE, model=BASE_MODEL, per_device_train_batch_size=2, warmup_steps=100, learning_rate=3e-4, max_steps=200):\n",
    "    if device != torch.device(\"cpu\"):\n",
    "        print(f\"Finetuning on device: {ipex.xpu.get_device_name()}\")\n",
    "    print(f\"Using model: {model}\")\n",
    "    print(f\"per device batch size: {per_device_train_batch_size}\")\n",
    "    print(f\"warmup steps: {warmup_steps}\")\n",
    "    print(f\"learning rate: {learning_rate}\")\n",
    "    print(f\"max steps: {max_steps}\")\n",
    "    try:\n",
    "        finetuner = FineTuner(\n",
    "            base_model_id=model, model_path=MODEL_PATH, device=device\n",
    "        )\n",
    "        training_args = TrainingArguments(\n",
    "            per_device_train_batch_size=per_device_train_batch_size,\n",
    "            gradient_accumulation_steps=2,\n",
    "            warmup_steps=warmup_steps,\n",
    "            save_steps=max_steps // 4,\n",
    "            save_strategy=\"steps\",\n",
    "            eval_steps=max_steps // 4,\n",
    "            evaluation_strategy=\"steps\",\n",
    "            max_steps=max_steps,\n",
    "            learning_rate=learning_rate,\n",
    "            # num_train_epochs=2,\n",
    "            max_grad_norm=0.3,\n",
    "            bf16=True,\n",
    "            lr_scheduler_type=\"cosine\",\n",
    "            load_best_model_at_end=True,\n",
    "            ddp_find_unused_parameters=False,\n",
    "            group_by_length=True,\n",
    "            save_total_limit=3,\n",
    "            logging_steps=max_steps // 10,\n",
    "            optim=\"adamw_hf\",\n",
    "            output_dir=ADAPTER_PATH,\n",
    "            logging_dir=\"./logs\",\n",
    "            report_to= [],\n",
    "        )\n",
    "        finetuner.finetune(DATA_PATH, training_args)\n",
    "    except Exception as e:\n",
    "        logging.error(f\"Error occurred: {e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c497cbaf-994b-474b-92f1-fb33fca6b81f",
   "metadata": {},
   "source": [
    "## Fine-Tuning the Model\n",
    "Now it's time to actually fine-tune the model. The `lets_finetune` function below takes care of this. It initializes a FineTuner object with the configurations you've set or left as default.\n",
    "\n",
    "### What Does It Do?\n",
    "- Initializes the FineTuner object with the base model and other configurations.\n",
    "- Sets up training arguments like batch size, learning rate, evaluation steps, etc.\n",
    "- Starts the fine-tuning process using the data and configurations provided."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "1fa45c6d-fac1-4331-8dd5-3ebf9cceed6d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finetuning on device: Intel(R) Arc(TM) A770M Graphics\n",
      "Using model: openlm-research/open_llama_3b_v2\n",
      "per device batch size: 2\n",
      "warmup steps: 100\n",
      "learning rate: 0.0003\n",
      "max steps: 200\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b3933616045b468b8bbb7164f7e3afed",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/78477 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "608afa5cf1e240d1a234438be968b350",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/100 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 2.0321, 'learning_rate': 5.9999999999999995e-05, 'epoch': 0.0}\n",
      "{'loss': 1.0263, 'learning_rate': 0.00011999999999999999, 'epoch': 0.0}\n",
      "{'eval_loss': 0.682393491268158, 'eval_runtime': 6.5451, 'eval_samples_per_second': 15.279, 'eval_steps_per_second': 1.986, 'epoch': 0.0}\n",
      "{'loss': 0.6114, 'learning_rate': 0.00017999999999999998, 'epoch': 0.0}\n",
      "{'loss': 0.5986, 'learning_rate': 0.00023999999999999998, 'epoch': 0.0}\n",
      "{'loss': 0.4158, 'learning_rate': 0.0003, 'epoch': 0.01}\n",
      "{'eval_loss': 0.6204714179039001, 'eval_runtime': 6.7285, 'eval_samples_per_second': 14.862, 'eval_steps_per_second': 1.932, 'epoch': 0.01}\n",
      "{'loss': 0.6238, 'learning_rate': 0.0002713525491562421, 'epoch': 0.01}\n",
      "{'loss': 0.4599, 'learning_rate': 0.0001963525491562421, 'epoch': 0.01}\n",
      "{'eval_loss': 0.5696825981140137, 'eval_runtime': 6.7261, 'eval_samples_per_second': 14.867, 'eval_steps_per_second': 1.933, 'epoch': 0.01}\n",
      "{'loss': 0.5081, 'learning_rate': 0.0001036474508437579, 'epoch': 0.01}\n",
      "{'loss': 0.5237, 'learning_rate': 2.8647450843757897e-05, 'epoch': 0.01}\n",
      "{'loss': 0.3939, 'learning_rate': 0.0, 'epoch': 0.01}\n",
      "{'eval_loss': 0.5195469856262207, 'eval_runtime': 6.7291, 'eval_samples_per_second': 14.861, 'eval_steps_per_second': 1.932, 'epoch': 0.01}\n",
      "{'train_runtime': 395.2213, 'train_samples_per_second': 2.024, 'train_steps_per_second': 0.506, 'train_loss': 0.7193752455711365, 'epoch': 0.01}\n"
     ]
    }
   ],
   "source": [
    "lets_finetune()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39613af7-3408-4412-8ded-aa6116759f06",
   "metadata": {},
   "source": [
    "### Test Your Fine-Tuned Language Model\n",
    "\n",
    "Congratulations on successfully fine-tuning your Language Model for Text-to-SQL tasks! It's now time to put your model to the test. You have two models at your disposal: the base model, which is the original pre-trained model without fine-tuning, and the fine-tuned model, which has been optimized with LoRA for your specific Text-to-SQL tasks.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00cb638e-1464-4596-8809-186834b3f277",
   "metadata": {},
   "source": [
    "### TextToSQLGenerator: Generating SQL Queries from Text Prompts\n",
    "\n",
    "The `TextToSQLGenerator` class is a pivotal component of our Text-to-SQL translation tool. It encapsulates the logic for generating SQL queries from natural language prompts, leveraging the power of large language models. Here's what you need to know about its implementation:\n",
    "\n",
    "#### Initialization and Configuration:\n",
    "\n",
    "Upon initialization, the `TextToSQLGenerator` class allows the selection of the underlying model. By setting the `use_lora` parameter to `True`, you can enable the fine-tuned LoRA model for inference. Otherwise, the base pre-trained model is used by default.\n",
    "\n",
    "The constructor of the class is designed to be robust and user-friendly. It automatically selects the correct tokenizer based on the model ID, with a special case for the 'llama' models. The model is then loaded with options optimized for performance on CPUs, such as `low_cpu_mem_usage` and `load_in_4bit`, which are beneficial for environments with memory constraints.\n",
    "\n",
    "For LoRA models, the constructor also takes care of loading the LoRA-specific checkpoints, ensuring that the fine-tuned parameters are utilized during inference.\n",
    "\n",
    "#### Generating SQL Queries:\n",
    "\n",
    "The `generate` method is where the actual translation occurs. Given a text prompt, the method encodes the prompt using the tokenizer, ensuring that it fits within the model's maximum length constraints. It then performs inference to generate the SQL query, carefully managing the GPU or CPU memory and leveraging half-precision computation when available for faster inference.\n",
    "\n",
    "The method is equipped with parameters like `temperature` and `num_beams` to control the creativity and quality of the generated queries, ensuring they are both diverse and close to the expected outputs.\n",
    "\n",
    "#### Usage:\n",
    "\n",
    "With `TextToSQLGenerator`, users can easily compare the performance of the base model and the LoRA-enhanced model by instantiating two objects of the class with different `use_lora` settings. This facilitates a side-by-side comparison of the generated SQL queries, demonstrating the effectiveness of the fine-tuning process.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10a0e504-46b7-4c6e-8d80-b28798607ed6",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TextToSQLGenerator:\n",
    "    \"\"\"Handles SQL query generation for a given text prompt.\"\"\"\n",
    "\n",
    "    def __init__(self, use_lora=False):\n",
    "        \"\"\"\n",
    "        Initialize the InferenceModel class.\n",
    "        Parameters:\n",
    "            use_lora (bool, optional): Whether to use LoRA model. Defaults to False.\n",
    "        \"\"\"\n",
    "        try:\n",
    "            # Choose the appropriate tokenizer based on the model name\n",
    "            if 'llama' in self.base_model_id.lower():\n",
    "                self.tokenizer = LlamaTokenizer.from_pretrained(self.base_model_id)\n",
    "            else:\n",
    "                self.tokenizer = AutoTokenizer.from_pretrained(self.base_model_id)\n",
    "            self.model = AutoModelForCausalLM.from_pretrained(\n",
    "                BASE_MODEL,\n",
    "                low_cpu_mem_usage=True,\n",
    "                load_in_4bit=True,\n",
    "                optimize_model=False,\n",
    "                use_cache=True,\n",
    "            )\n",
    "            if use_lora:\n",
    "                self.model = PeftModel.from_pretrained(self.model, LORA_CHECKPOINT)\n",
    "        except Exception as e:\n",
    "            logging.error(f\"Exception occurred during model initialization: {e}\")\n",
    "            raise\n",
    "            \n",
    "        self.model.to(DEVICE)\n",
    "        self.max_length = 512\n",
    "        self.tokenizer.pad_token_id = 0\n",
    "        self.tokenizer.padding_side = \"left\"\n",
    "\n",
    "    def generate(self, prompt, **kwargs):\n",
    "        \"\"\"Generates an SQL query based on the given prompt.\n",
    "        Parameters:\n",
    "            prompt (str): The SQL prompt.\n",
    "        Returns:\n",
    "            str: The generated SQL query.\n",
    "        \"\"\"\n",
    "        try:\n",
    "            encoded_prompt = self.tokenizer(\n",
    "                prompt,\n",
    "                truncation=True,\n",
    "                max_length=self.max_length,\n",
    "                padding=False,\n",
    "                return_tensors=\"pt\",\n",
    "            ).input_ids.to(DEVICE)\n",
    "            with torch.no_grad():\n",
    "                with torch.xpu.amp.autocast():\n",
    "                    outputs = self.model.generate(\n",
    "                        input_ids=encoded_prompt,\n",
    "                        do_sample=False,\n",
    "                        max_length=self.max_length,\n",
    "                        temperature=0.3,\n",
    "                        num_beams=5,\n",
    "                        repetition_penalty=1.2,\n",
    "                    )\n",
    "            generated = self.tokenizer.decode(outputs[0], skip_special_tokens=True)\n",
    "            return generated\n",
    "        except Exception as e:\n",
    "            logging.error(f\"Exception occurred during query generation: {e}\")\n",
    "            raise\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dfb1c247-f6e2-4ccd-bdc6-26f41ea63c55",
   "metadata": {},
   "source": [
    "\n",
    "#### Steps to Follow:\n",
    "1. **Select or Input a Natural Language Query**: Start with a natural language question or prompt that you want to translate into an SQL query. You can either choose from the preloaded sample data or input your own.\n",
    "\n",
    "2. **Generate the SQL Query Using the Base Model**: Use the base model to generate an SQL query based on your natural language prompt. This will serve as a baseline for comparison.\n",
    "\n",
    "3. **Generate the SQL Query Using the Fine-Tuned Model**: Next, generate the SQL query using the fine-tuned model. Observe the differences and improvements in the SQL query generated by the fine-tuned model.\n",
    "\n",
    "4. **Compare the Outputs**: Look at the SQL queries generated by both models. Does the fine-tuned model capture the intent of the natural language prompt more accurately? Is the SQL query syntactically correct and optimized?\n",
    "\n",
    "5. **Iterate and Refine**: Fine-tuning is an iterative process. If the outputs aren't as expected, consider revising the training parameters or dataset and fine-tune again.\n",
    "\n",
    "#### Try with Sample Data:\n",
    "The following code snippet automatically processes a few examples from the loaded dataset. Observe how the base model and the fine-tuned model perform on the same prompts.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8417146-0c28-4996-81a4-c4b0857d81e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "TEST_DATA = \"./test_data/sample_test_data.json\"\n",
    "base_model = TextToSQLGenerator()\n",
    "finetuned_model = TextToSQLGenerator(use_lora=True)\n",
    "sample_data = load_dataset(\"json\", data_files=TEST_DATA)[\"train\"]\n",
    "for row in sample_data:\n",
    "    try:\n",
    "        prompt = generate_prompt_sql(row[\"question\"], context=row[\"context\"])\n",
    "        print(\"Using base model...\")\n",
    "        output = base_model.generate(prompt)\n",
    "        print(f\"\\n\\tbot response: {output}\\n\")\n",
    "\n",
    "        print(\"Using finetuned model...\")\n",
    "        output = finetuned_model.generate(prompt)\n",
    "        print(f\"\\n\\tbot response: {output}\\n\")\n",
    "    except Exception as e:\n",
    "        logging.error(f\"Exception occurred during sample processing: {e}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "61837bc3-7281-4bed-84fb-f797228d2f28",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "question: What was the score of the game wehre the deportes savio were the home team?\n",
      "context: CREATE TABLE table_name_78 (score VARCHAR, home VARCHAR)\n",
      "\n",
      "\n",
      "question: In which championship did the winner have a score of −9 (66-69-73-71=279)?\n",
      "context: CREATE TABLE table_name_83 (championship VARCHAR, winning_score VARCHAR)\n",
      "\n",
      "\n",
      "question: During the championship where the winning score was −9 (66-69-73-71=279)?, who was the runner-up?\n",
      "context: CREATE TABLE table_name_32 (runner_s__up VARCHAR, winning_score VARCHAR)\n",
      "\n",
      "\n",
      "question: What was the winning score in the year 2002?\n",
      "context: CREATE TABLE table_name_28 (winning_score VARCHAR, year VARCHAR)\n",
      "\n",
      "\n",
      "question: What is the fewest goals of CD Alcoyano with more than 25 points?\n",
      "context: CREATE TABLE table_name_57 (goals_against INTEGER, club VARCHAR, points VARCHAR)\n",
      "\n",
      "\n",
      "question: What is the fewest number of wins that had fewer than 7 draws and more than 30 played?\n",
      "context: CREATE TABLE table_name_26 (wins INTEGER, draws VARCHAR, played VARCHAR)\n",
      "\n",
      "\n",
      "question: What is the fewest points that has more than 29 goals?\n",
      "context: CREATE TABLE table_name_42 (points INTEGER, goal_difference INTEGER)\n",
      "\n",
      "\n",
      "question: What date have highest 0-4-4 forney locomotive with number larger than 20 and works number larger than 23754?\n",
      "context: CREATE TABLE table_name_90 (date INTEGER, works_number VARCHAR, type VARCHAR, number VARCHAR)\n",
      "\n",
      "\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "569834efd4ba465190da369eec9a1db9",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(Textarea(value='', description='Question:', layout=Layout(height='100px', width='95%'), placeho…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import ipywidgets as widgets\n",
    "from IPython.display import display, clear_output\n",
    "import json\n",
    "\n",
    "import ipywidgets as widgets\n",
    "from IPython.display import display, clear_output\n",
    "from datasets import load_dataset\n",
    "\n",
    "def format_and_print_json_objects(file_path):\n",
    "    with open(file_path, 'r') as file:\n",
    "        json_object_str = ''\n",
    "        json_objects = []\n",
    "\n",
    "        for line in file:\n",
    "            json_object_str += line.strip()\n",
    "            if line.strip() == '}':\n",
    "                try:\n",
    "                    json_object = json.loads(json_object_str)\n",
    "                    json_objects.append(json_object)\n",
    "                    json_object_str = ''\n",
    "                except json.JSONDecodeError as e:\n",
    "                    logging.error(\"JSON Decode Error:\", e)\n",
    "                    json_object_str = ''\n",
    "                    continue\n",
    "\n",
    "    # Now we have our list of JSON objects, let's format and print them\n",
    "    for entry in json_objects:\n",
    "        question = entry['question']\n",
    "        context = entry['context']\n",
    "        print(f\"question: {question}\\ncontext: {context}\\n\\n\")\n",
    "\n",
    "format_and_print_json_objects(TEST_DATA)\n",
    "\n",
    "custom_query_input = widgets.Textarea(\n",
    "    value=\"\",\n",
    "    placeholder=\"Type your natural language query here...\",\n",
    "    description=\"Question:\",\n",
    "    layout=widgets.Layout(width=\"95%\", height=\"100px\"),\n",
    "    style={\"description_width\": \"initial\"}\n",
    ")\n",
    "\n",
    "custom_context_input = widgets.Textarea(\n",
    "    value=\"\",\n",
    "    placeholder=\"Type your SQL schema here...\",\n",
    "    description=\"Context:\",\n",
    "    layout=widgets.Layout(width=\"95%\", height=\"100px\"),\n",
    "    style={\"description_width\": \"initial\"}\n",
    ")\n",
    "\n",
    "def generate_custom_query(b):\n",
    "    query = custom_query_input.value\n",
    "    context = custom_context_input.value\n",
    "    prompt = generate_prompt_sql(query, context=context)\n",
    "    print(\"Using base model...\")\n",
    "    base_output = base_model.generate(prompt)  # base_model should be defined and loaded\n",
    "    print(f\"\\n\\tBase model SQL: {base_output}\\n\")\n",
    "    print(\"Using fine-tuned model...\")\n",
    "    finetuned_output = finetuned_model.generate(prompt)  # finetuned_model should be defined and loaded\n",
    "    print(f\"\\n\\tFine-tuned model SQL: {finetuned_output}\\n\")\n",
    "\n",
    "generate_custom_query_btn = widgets.Button(\n",
    "    description=\"Generate SQL from Custom Query\",\n",
    "    button_style=\"primary\",\n",
    ")\n",
    "generate_custom_query_btn.on_click(generate_custom_query)\n",
    "output_area = widgets.Output()\n",
    "\n",
    "ui = widgets.VBox([\n",
    "    custom_query_input,\n",
    "    custom_context_input,\n",
    "    generate_custom_query_btn,\n",
    "    output_area\n",
    "])\n",
    "\n",
    "display(ui)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe2d2048-0bf5-407d-9fee-f24fb38b60b8",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch-gpu",
   "language": "python",
   "name": "pytorch-gpu"
  },
  "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.9.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}