transformers_fine_tune.ipynb

{
 "cells": [
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## Fine tuning LLM\n",
    "This notebook shows how to fine tune LLM model. Fine-tuning is advanced technique which can modify model parameters to perform better in specific task. "
   ],
   "id": "cb478451d99f9bd3"
  },
  {
   "metadata": {
    "jupyter": {
     "is_executing": true
    }
   },
   "cell_type": "code",
   "source": [
    "%pip install \\\n",
    "datasets \\\n",
    "transformers \\\n",
    "evaluate \\\n",
    "torch \\\n",
    "torchdata \\\n",
    "rouge_score \\\n",
    "loralib \\\n",
    "peft"
   ],
   "id": "4012afec7007a957",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-09-15T15:40:25.359916Z",
     "start_time": "2024-09-15T15:40:19.453478Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import torch\n",
    "from datasets import load_dataset\n",
    "from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, GenerationConfig, TrainingArguments, Trainer"
   ],
   "id": "cc3d6045b7e78c13",
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-09-15 17:40:22.676461: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:485] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
      "2024-09-15 17:40:22.720989: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:8454] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
      "2024-09-15 17:40:22.734283: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1452] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
      "2024-09-15 17:40:22.808849: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.\n",
      "To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2024-09-15 17:40:23.996240: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT\n"
     ]
    }
   ],
   "execution_count": 1
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-09-15T15:40:56.720710Z",
     "start_time": "2024-09-15T15:40:25.376086Z"
    }
   },
   "cell_type": "code",
   "source": [
    "model_name='google/flan-t5-base'\n",
    "model = AutoModelForSeq2SeqLM.from_pretrained(model_name, torch_dtype=torch.bfloat16)\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "\n",
    "dataset_train = load_dataset(\"gretelai/synthetic_text_to_sql\", split='train[0%:80%]')\n",
    "dataset_valid = load_dataset(\"gretelai/synthetic_text_to_sql\", split='train[80%:100%]')\n",
    "dataset_test = load_dataset(\"gretelai/synthetic_text_to_sql\", split='test')\n",
    "\n",
    "dataset_train"
   ],
   "id": "3014341f65aa129",
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/michal/anaconda3/lib/python3.11/site-packages/huggingface_hub/file_download.py:1150: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['id', 'domain', 'domain_description', 'sql_complexity', 'sql_complexity_description', 'sql_task_type', 'sql_task_type_description', 'sql_prompt', 'sql_context', 'sql', 'sql_explanation'],\n",
       "    num_rows: 80000\n",
       "})"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 2
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "### Check results with zero-shot",
   "id": "11edd6923d6074a2"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-09-15T15:41:48.073784Z",
     "start_time": "2024-09-15T15:40:56.813683Z"
    }
   },
   "cell_type": "code",
   "source": [
    "example_indexes = [37, 67]\n",
    "for i, index in enumerate(example_indexes):\n",
    "    schema = dataset_test[index]['sql_context']\n",
    "    instruction = dataset_test[index]['sql_prompt']\n",
    "    sql = dataset_test[index]['sql']\n",
    "\n",
    "    prompt = f\"\"\"\n",
    "Translate schema and description into SQL query. Respond only with SQL query without any additional characters.\n",
    "\n",
    "Schema: {schema}\n",
    "\n",
    "Instruction: {instruction}\n",
    "\n",
    "SQL query:\n",
    "    \"\"\"\n",
    "\n",
    "    inputs = tokenizer(prompt, return_tensors='pt')\n",
    "    output = tokenizer.decode(\n",
    "        model.generate(\n",
    "            inputs[\"input_ids\"],\n",
    "            max_new_tokens=50,\n",
    "        )[0],\n",
    "        skip_special_tokens=True\n",
    "    )\n",
    "\n",
    "    print('Example ', i + 1)\n",
    "    print(f'Task:\\n{instruction}')\n",
    "    print(f'Schema:\\n{schema}')\n",
    "    print(f'CORRECT SQL:\\n{sql}')\n",
    "    print(f'MODEL GENERATION SQL - ZERO SHOT:\\n{output}\\n')"
   ],
   "id": "e221d8bba828406",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Example  1\n",
      "Task:\n",
      "What is the number of medical supplies distributed by each organization, in East Africa, for the last 3 years, and the total cost of the supplies?\n",
      "Schema:\n",
      "CREATE TABLE medical_supplies (supply_id INT, organization_id INT, location VARCHAR(255), supply_type VARCHAR(255), supply_cost DECIMAL(10,2), distribution_date DATE); INSERT INTO medical_supplies VALUES (1, 1, 'Country A', 'Medicine', 5000, '2020-01-01'); INSERT INTO medical_supplies VALUES (2, 1, 'Country A', 'Medical Equipment', 7000, '2021-01-01'); INSERT INTO medical_supplies VALUES (3, 2, 'Country B', 'Vaccines', 10000, '2021-01-01'); INSERT INTO medical_supplies VALUES (4, 2, 'Country B', 'First Aid Kits', 8000, '2020-01-01');\n",
      "CORRECT SQL:\n",
      "SELECT organization_id, location as region, COUNT(*) as number_of_supplies, SUM(supply_cost) as total_supply_cost FROM medical_supplies WHERE location = 'East Africa' AND distribution_date >= DATE_SUB(CURRENT_DATE, INTERVAL 3 YEAR) GROUP BY organization_id, location;\n",
      "MODEL GENERATION SQL - ZERO SHOT:\n",
      "SELECT TABLE_NUMBER_OF_MEDICAL_SURFACES FROM DISTRICTIONS AS T1, T2, T3, T4, T5, T6 \n",
      "\n",
      "Example  2\n",
      "Task:\n",
      "List the names of all parks in urban areas\n",
      "Schema:\n",
      "CREATE TABLE parks (park_id INT, area_id INT, park_name TEXT);CREATE TABLE areas (area_id INT, area_type TEXT);\n",
      "CORRECT SQL:\n",
      "SELECT p.park_name FROM parks p INNER JOIN areas a ON p.area_id = a.area_id WHERE a.area_type = 'urban';\n",
      "MODEL GENERATION SQL - ZERO SHOT:\n",
      "SELECT TABLE parks (park_id INT, area_id INT, park_name TEXT); SELECT TABLE areas (area_id INT, area_type TEXT\n",
      "\n"
     ]
    }
   ],
   "execution_count": 3
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "### Conduct model fine tuning process\n",
    "Fine tune model using HuggingFace Trainer class"
   ],
   "id": "47d4cfe91775f7a8"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-09-15T15:43:12.004644Z",
     "start_time": "2024-09-15T15:41:48.377598Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# Preprocess dataset\n",
    "def tokenize(row):\n",
    "    prompt = f\"\"\"\n",
    "Translate schema and description into SQL query. Respond only with SQL query without any additional characters.\n",
    "\n",
    "Schema: {row[\"sql_context\"]}\n",
    "\n",
    "Instruction: {row[\"sql_prompt\"]}\n",
    "\n",
    "SQL query:\n",
    "    \"\"\"\n",
    "    row['input_ids'] = tokenizer(prompt, padding=\"max_length\", truncation=True, return_tensors=\"pt\").input_ids.squeeze(0)\n",
    "    row['labels'] = tokenizer(row[\"sql\"], padding=\"max_length\", truncation=True,return_tensors=\"pt\").input_ids.squeeze(0)\n",
    "\n",
    "    return row\n",
    "columns = dataset_train.column_names\n",
    "tokenized_datasets_train = dataset_train.map(tokenize)\n",
    "tokenized_datasets_train = tokenized_datasets_train.remove_columns(columns)\n",
    "print(\"Train:\")\n",
    "print(tokenized_datasets_train.shape)\n",
    "print(tokenized_datasets_train)\n",
    "\n",
    "tokenized_datasets_valid = dataset_valid.map(tokenize)\n",
    "tokenized_datasets_valid = tokenized_datasets_valid.remove_columns(columns)\n",
    "print(\"Valid:\")\n",
    "print(tokenized_datasets_valid.shape)\n",
    "print(tokenized_datasets_valid)\n",
    "\n",
    "tokenized_datasets_test = dataset_test.map(tokenize)\n",
    "tokenized_datasets_test = tokenized_datasets_test.remove_columns(columns)\n",
    "print(\"Test:\")\n",
    "print(tokenized_datasets_test.shape)\n",
    "print(tokenized_datasets_test)"
   ],
   "id": "f3ca509043a05b33",
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Map:   0%|          | 0/80000 [00:00<?, ? examples/s]"
      ],
      "application/vnd.jupyter.widget-view+json": {
       "version_major": 2,
       "version_minor": 0,
       "model_id": "81f7f87752b84cd49b8b7030a7de9c91"
      }
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train:\n",
      "(80000, 2)\n",
      "Dataset({\n",
      "    features: ['input_ids', 'labels'],\n",
      "    num_rows: 80000\n",
      "})\n",
      "Valid:\n",
      "(20000, 2)\n",
      "Dataset({\n",
      "    features: ['input_ids', 'labels'],\n",
      "    num_rows: 20000\n",
      "})\n",
      "Test:\n",
      "(5851, 2)\n",
      "Dataset({\n",
      "    features: ['input_ids', 'labels'],\n",
      "    num_rows: 5851\n",
      "})\n"
     ]
    }
   ],
   "execution_count": 4
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-09-15T15:43:12.077982Z",
     "start_time": "2024-09-15T15:43:12.074466Z"
    }
   },
   "cell_type": "code",
   "source": "print(tokenized_datasets_train[0])",
   "id": "484f850f9f734dd0",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'input_ids': [30355, 15, 26622, 11, 4210, 139, 12558, 11417, 5, 23483, 163, 28, 12558, 11417, 406, 136, 1151, 2850, 5, 10248, 51, 9, 10, 205, 4386, 6048, 332, 17098, 1085, 6075, 41, 7, 4529, 6075, 834, 23, 26, 3, 13777, 6, 564, 3, 3463, 4, 382, 6, 1719, 3, 3463, 4, 382, 3670, 3, 14750, 24203, 3388, 5647, 1085, 6075, 41, 7, 4529, 6075, 834, 23, 26, 6, 564, 6, 1719, 61, 3, 21712, 5078, 134, 4077, 6, 3, 31, 18300, 531, 15, 31, 6, 3, 31, 22969, 31, 201, 4743, 6, 3, 31, 683, 152, 15, 3931, 31, 6, 3, 31, 22081, 31, 3670, 205, 4386, 6048, 332, 17098, 14592, 834, 7, 4529, 41, 7, 4529, 834, 23, 26, 3, 13777, 6, 1085, 6075, 834, 23, 26, 3, 13777, 6, 2908, 17833, 6, 1048, 834, 5522, 309, 6048, 3670, 3, 14750, 24203, 3388, 5647, 14592, 834, 7, 4529, 41, 7, 4529, 834, 23, 26, 6, 1085, 6075, 834, 23, 26, 6, 2908, 6, 1048, 834, 5522, 61, 3, 21712, 5078, 134, 4077, 6, 1914, 5864, 6, 3, 31, 1755, 2658, 14772, 14772, 31, 201, 4743, 6, 1914, 4261, 6, 3, 31, 19818, 18930, 357, 14772, 31, 201, 6918, 6, 3547, 8003, 6, 3, 31, 1755, 2658, 14772, 14772, 31, 3670, 21035, 10, 363, 19, 8, 792, 2908, 13, 14592, 1916, 57, 284, 1085, 6075, 6, 3, 14504, 57, 1085, 6075, 58, 12558, 11417, 10, 3, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], 'labels': [3, 23143, 14196, 1085, 6075, 834, 23, 26, 6, 564, 6, 180, 6122, 599, 23439, 61, 38, 792, 834, 23439, 21680, 14592, 834, 7, 4529, 3, 15355, 3162, 1085, 6075, 9191, 14592, 834, 7, 4529, 5, 7, 4529, 6075, 834, 23, 26, 3274, 1085, 6075, 5, 7, 4529, 6075, 834, 23, 26, 350, 4630, 6880, 272, 476, 1085, 6075, 834, 23, 26, 6, 564, 4674, 11300, 272, 476, 792, 834, 23439, 309, 25067, 117, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]}\n"
     ]
    }
   ],
   "execution_count": 5
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-09-15T15:43:12.285927Z",
     "start_time": "2024-09-15T15:43:12.215184Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# Perform fine tuning\n",
    "\n",
    "output_dir = f'./text-to-sql-translation'\n",
    "\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=output_dir,\n",
    "    learning_rate=1e-5,\n",
    "    num_train_epochs=10,\n",
    "    weight_decay=0.01,\n",
    "    logging_steps=1\n",
    ")\n",
    "\n",
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=tokenized_datasets_train,\n",
    "    eval_dataset=tokenized_datasets_valid\n",
    ")"
   ],
   "id": "b7c4833094681f53",
   "outputs": [],
   "execution_count": 6
  },
  {
   "metadata": {
    "jupyter": {
     "is_executing": true
    },
    "ExecuteTime": {
     "start_time": "2024-09-15T15:52:07.092939Z"
    }
   },
   "cell_type": "code",
   "source": "trainer.train()",
   "id": "17c6b2c6f8b14c6d",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "### Test fine-tuned model ",
   "id": "acc71a8d33f26b0e"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": "tuned_model = AutoModelForSeq2SeqLM.from_pretrained(\"./text-to-sql-translation\", torch_dtype=torch.bfloat16)",
   "id": "3fae60d77b11052f"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "index = 78\n",
    "schema = dataset_test[index]['sql_context']\n",
    "instruction = dataset_test[index]['sql_prompt']\n",
    "sql = dataset_test[index]['sql']\n",
    "\n",
    "prompt = f\"\"\"\n",
    "Translate schema and description into SQL query. Respond only with SQL query without any additional characters.\n",
    "\n",
    "Schema: {schema}\n",
    "\n",
    "Instruction: {instruction}\n",
    "\n",
    "SQL query:\n",
    "    \"\"\"\n",
    "\n",
    "inputs = tokenizer(prompt, return_tensors='pt')\n",
    "output = tokenizer.decode(\n",
    "    model.generate(\n",
    "        inputs[\"input_ids\"],\n",
    "        max_new_tokens=50,\n",
    "    )[0],\n",
    "    skip_special_tokens=True\n",
    ")\n",
    "\n",
    "tuned_output = tokenizer.decode(\n",
    "    tuned_model.generate(\n",
    "        inputs[\"input_ids\"],\n",
    "        max_new_tokens=50,\n",
    "    )[0],\n",
    "    skip_special_tokens=True\n",
    ")\n",
    "\n",
    "print(f'Task:\\n{instruction}')\n",
    "print(f'Schema:\\n{schema}')\n",
    "print(f'CORRECT SQL:\\n{sql}')\n",
    "print(f'MODEL GENERATION SQL - ZERO SHOT:\\n{output}\\n')\n",
    "print(f'MODEL GENERATION SQL - FINE TUNED:\\n{tuned_output}\\n')"
   ],
   "id": "6997abd886d09ab3"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "### PEFT / LoRA\n",
    "Fine tune model using PEFT/LoRA techniques"
   ],
   "id": "40e61d2676583405"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "from peft import LoraConfig, get_peft_model, TaskType\n",
    "\n",
    "lora_config = LoraConfig(\n",
    "    r=32, # Rank\n",
    "    lora_alpha=32,\n",
    "    target_modules=[\"q\", \"v\"],\n",
    "    lora_dropout=0.05,\n",
    "    bias=\"none\",\n",
    "    task_type=TaskType.SEQ_2_SEQ_LM # FLAN-T5\n",
    ")\n",
    "\n",
    "peft_model = get_peft_model(model, lora_config)"
   ],
   "id": "1c3ee569f1425f26"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "output_dir = f'./peft-text-to-sql-translation'\n",
    "\n",
    "peft_training_args = TrainingArguments(\n",
    "    output_dir=output_dir,\n",
    "    auto_find_batch_size=True,\n",
    "    learning_rate=1e-3, # Higher learning rate than full fine-tuning.\n",
    "    num_train_epochs=1,\n",
    "    logging_steps=1,\n",
    "    max_steps=1\n",
    ")\n",
    "\n",
    "peft_trainer = Trainer(\n",
    "    model=peft_model,\n",
    "    args=peft_training_args,\n",
    "    train_dataset=tokenized_datasets_train,\n",
    ")"
   ],
   "id": "3c3c218c0f486f5e"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "source": [
    "peft_trainer.train()\n",
    "\n",
    "peft_model_path=\"./peft-text-to-sql-translation\"\n",
    "\n",
    "peft_trainer.model.save_pretrained(peft_model_path)\n",
    "tokenizer.save_pretrained(peft_model_path)"
   ],
   "id": "4f10a5cecc9d75a0",
   "outputs": [],
   "execution_count": null
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "from peft import PeftModel, PeftConfig\n",
    "\n",
    "peft_model_base = AutoModelForSeq2SeqLM.from_pretrained(\"google/flan-t5-base\", torch_dtype=torch.bfloat16)\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"google/flan-t5-base\")\n",
    "\n",
    "peft_model = PeftModel.from_pretrained(peft_model_base,\n",
    "                                       './peft-text-to-sql-translation',\n",
    "                                       torch_dtype=torch.bfloat16,\n",
    "                                       is_trainable=False)"
   ],
   "id": "4f41c8b54d686fa2"
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "### Evaluate PEFT fine tuned model",
   "id": "fdec21227907a7f"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "index = 78\n",
    "schema = dataset_test[index]['sql_context']\n",
    "instruction = dataset_test[index]['sql_prompt']\n",
    "sql = dataset_test[index]['sql']\n",
    "\n",
    "prompt = f\"\"\"\n",
    "Translate schema and description into SQL query. Respond only with SQL query without any additional characters.\n",
    "\n",
    "Schema: {schema}\n",
    "\n",
    "Instruction: {instruction}\n",
    "\n",
    "SQL query:\n",
    "    \"\"\"\n",
    "\n",
    "inputs = tokenizer(prompt, return_tensors='pt')\n",
    "output = tokenizer.decode(\n",
    "    model.generate(\n",
    "        inputs[\"input_ids\"],\n",
    "        max_new_tokens=50,\n",
    "    )[0],\n",
    "    skip_special_tokens=True\n",
    ")\n",
    "\n",
    "tuned_output = tokenizer.decode(\n",
    "    tuned_model.generate(\n",
    "        inputs[\"input_ids\"],\n",
    "        max_new_tokens=50,\n",
    "    )[0],\n",
    "    skip_special_tokens=True\n",
    ")\n",
    "\n",
    "peft_output = tokenizer.decode(\n",
    "    tuned_model.generate(\n",
    "        inputs[\"input_ids\"],\n",
    "        max_new_tokens=50,\n",
    "    )[0],\n",
    "    skip_special_tokens=True\n",
    ")\n",
    "\n",
    "print(f'Task:\\n{instruction}')\n",
    "print(f'Schema:\\n{schema}')\n",
    "print(f'CORRECT SQL:\\n{sql}')\n",
    "print(f'MODEL GENERATION SQL - ZERO SHOT:\\n{output}\\n')\n",
    "print(f'MODEL GENERATION SQL - FINE TUNED:\\n{tuned_output}\\n')\n",
    "print(f'MODEL GENERATION SQL - PEFT/LoRA TUNED:\\n{peft_output}\\n')"
   ],
   "id": "5fa260e65671dcc7"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}