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": {
    "ExecuteTime": {
     "end_time": "2024-12-08T09:32:50.580695Z",
     "start_time": "2024-12-08T09:32:48.843999Z"
    }
   },
   "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": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: datasets in /home/michal/anaconda3/lib/python3.12/site-packages (3.1.0)\r\n",
      "Requirement already satisfied: transformers in /home/michal/anaconda3/lib/python3.12/site-packages (4.46.2)\r\n",
      "Requirement already satisfied: evaluate in /home/michal/anaconda3/lib/python3.12/site-packages (0.4.3)\r\n",
      "Requirement already satisfied: torch in /home/michal/anaconda3/lib/python3.12/site-packages (2.5.1)\r\n",
      "Requirement already satisfied: torchdata in /home/michal/anaconda3/lib/python3.12/site-packages (0.9.0)\r\n",
      "Requirement already satisfied: rouge_score in /home/michal/anaconda3/lib/python3.12/site-packages (0.1.2)\r\n",
      "Requirement already satisfied: loralib in /home/michal/anaconda3/lib/python3.12/site-packages (0.1.2)\r\n",
      "Requirement already satisfied: peft in /home/michal/anaconda3/lib/python3.12/site-packages (0.13.2)\r\n",
      "Requirement already satisfied: filelock in /home/michal/anaconda3/lib/python3.12/site-packages (from datasets) (3.13.1)\r\n",
      "Requirement already satisfied: numpy>=1.17 in /home/michal/anaconda3/lib/python3.12/site-packages (from datasets) (1.26.4)\r\n",
      "Requirement already satisfied: pyarrow>=15.0.0 in /home/michal/anaconda3/lib/python3.12/site-packages (from datasets) (16.1.0)\r\n",
      "Requirement already satisfied: dill<0.3.9,>=0.3.0 in /home/michal/anaconda3/lib/python3.12/site-packages (from datasets) (0.3.8)\r\n",
      "Requirement already satisfied: pandas in /home/michal/anaconda3/lib/python3.12/site-packages (from datasets) (1.5.3)\r\n",
      "Requirement already satisfied: requests>=2.32.2 in /home/michal/anaconda3/lib/python3.12/site-packages (from datasets) (2.32.3)\r\n",
      "Requirement already satisfied: tqdm>=4.66.3 in /home/michal/anaconda3/lib/python3.12/site-packages (from datasets) (4.66.5)\r\n",
      "Requirement already satisfied: xxhash in /home/michal/anaconda3/lib/python3.12/site-packages (from datasets) (3.5.0)\r\n",
      "Requirement already satisfied: multiprocess<0.70.17 in /home/michal/anaconda3/lib/python3.12/site-packages (from datasets) (0.70.16)\r\n",
      "Requirement already satisfied: fsspec<=2024.9.0,>=2023.1.0 in /home/michal/anaconda3/lib/python3.12/site-packages (from fsspec[http]<=2024.9.0,>=2023.1.0->datasets) (2024.6.1)\r\n",
      "Requirement already satisfied: aiohttp in /home/michal/anaconda3/lib/python3.12/site-packages (from datasets) (3.10.5)\r\n",
      "Requirement already satisfied: huggingface-hub>=0.23.0 in /home/michal/anaconda3/lib/python3.12/site-packages (from datasets) (0.26.2)\r\n",
      "Requirement already satisfied: packaging in /home/michal/anaconda3/lib/python3.12/site-packages (from datasets) (24.1)\r\n",
      "Requirement already satisfied: pyyaml>=5.1 in /home/michal/anaconda3/lib/python3.12/site-packages (from datasets) (6.0.1)\r\n",
      "Requirement already satisfied: regex!=2019.12.17 in /home/michal/anaconda3/lib/python3.12/site-packages (from transformers) (2024.9.11)\r\n",
      "Requirement already satisfied: safetensors>=0.4.1 in /home/michal/anaconda3/lib/python3.12/site-packages (from transformers) (0.4.5)\r\n",
      "Requirement already satisfied: tokenizers<0.21,>=0.20 in /home/michal/anaconda3/lib/python3.12/site-packages (from transformers) (0.20.3)\r\n",
      "Requirement already satisfied: typing-extensions>=4.8.0 in /home/michal/anaconda3/lib/python3.12/site-packages (from torch) (4.11.0)\r\n",
      "Requirement already satisfied: networkx in /home/michal/anaconda3/lib/python3.12/site-packages (from torch) (3.3)\r\n",
      "Requirement already satisfied: jinja2 in /home/michal/anaconda3/lib/python3.12/site-packages (from torch) (3.1.4)\r\n",
      "Requirement already satisfied: nvidia-cuda-nvrtc-cu12==12.4.127 in /home/michal/anaconda3/lib/python3.12/site-packages (from torch) (12.4.127)\r\n",
      "Requirement already satisfied: nvidia-cuda-runtime-cu12==12.4.127 in /home/michal/anaconda3/lib/python3.12/site-packages (from torch) (12.4.127)\r\n",
      "Requirement already satisfied: nvidia-cuda-cupti-cu12==12.4.127 in /home/michal/anaconda3/lib/python3.12/site-packages (from torch) (12.4.127)\r\n",
      "Requirement already satisfied: nvidia-cudnn-cu12==9.1.0.70 in /home/michal/anaconda3/lib/python3.12/site-packages (from torch) (9.1.0.70)\r\n",
      "Requirement already satisfied: nvidia-cublas-cu12==12.4.5.8 in /home/michal/anaconda3/lib/python3.12/site-packages (from torch) (12.4.5.8)\r\n",
      "Requirement already satisfied: nvidia-cufft-cu12==11.2.1.3 in /home/michal/anaconda3/lib/python3.12/site-packages (from torch) (11.2.1.3)\r\n",
      "Requirement already satisfied: nvidia-curand-cu12==10.3.5.147 in /home/michal/anaconda3/lib/python3.12/site-packages (from torch) (10.3.5.147)\r\n",
      "Requirement already satisfied: nvidia-cusolver-cu12==11.6.1.9 in /home/michal/anaconda3/lib/python3.12/site-packages (from torch) (11.6.1.9)\r\n",
      "Requirement already satisfied: nvidia-cusparse-cu12==12.3.1.170 in /home/michal/anaconda3/lib/python3.12/site-packages (from torch) (12.3.1.170)\r\n",
      "Requirement already satisfied: nvidia-nccl-cu12==2.21.5 in /home/michal/anaconda3/lib/python3.12/site-packages (from torch) (2.21.5)\r\n",
      "Requirement already satisfied: nvidia-nvtx-cu12==12.4.127 in /home/michal/anaconda3/lib/python3.12/site-packages (from torch) (12.4.127)\r\n",
      "Requirement already satisfied: nvidia-nvjitlink-cu12==12.4.127 in /home/michal/anaconda3/lib/python3.12/site-packages (from torch) (12.4.127)\r\n",
      "Requirement already satisfied: triton==3.1.0 in /home/michal/anaconda3/lib/python3.12/site-packages (from torch) (3.1.0)\r\n",
      "Requirement already satisfied: setuptools in /home/michal/anaconda3/lib/python3.12/site-packages (from torch) (75.1.0)\r\n",
      "Requirement already satisfied: sympy==1.13.1 in /home/michal/anaconda3/lib/python3.12/site-packages (from torch) (1.13.1)\r\n",
      "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /home/michal/anaconda3/lib/python3.12/site-packages (from sympy==1.13.1->torch) (1.3.0)\r\n",
      "Requirement already satisfied: urllib3>=1.25 in /home/michal/anaconda3/lib/python3.12/site-packages (from torchdata) (2.2.3)\r\n",
      "Requirement already satisfied: absl-py in /home/michal/anaconda3/lib/python3.12/site-packages (from rouge_score) (2.1.0)\r\n",
      "Requirement already satisfied: nltk in /home/michal/anaconda3/lib/python3.12/site-packages (from rouge_score) (3.9.1)\r\n",
      "Requirement already satisfied: six>=1.14.0 in /home/michal/anaconda3/lib/python3.12/site-packages (from rouge_score) (1.16.0)\r\n",
      "Requirement already satisfied: psutil in /home/michal/anaconda3/lib/python3.12/site-packages (from peft) (5.9.0)\r\n",
      "Requirement already satisfied: accelerate>=0.21.0 in /home/michal/anaconda3/lib/python3.12/site-packages (from peft) (1.1.1)\r\n",
      "Requirement already satisfied: aiohappyeyeballs>=2.3.0 in /home/michal/anaconda3/lib/python3.12/site-packages (from aiohttp->datasets) (2.4.0)\r\n",
      "Requirement already satisfied: aiosignal>=1.1.2 in /home/michal/anaconda3/lib/python3.12/site-packages (from aiohttp->datasets) (1.2.0)\r\n",
      "Requirement already satisfied: attrs>=17.3.0 in /home/michal/anaconda3/lib/python3.12/site-packages (from aiohttp->datasets) (24.2.0)\r\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in /home/michal/anaconda3/lib/python3.12/site-packages (from aiohttp->datasets) (1.4.0)\r\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /home/michal/anaconda3/lib/python3.12/site-packages (from aiohttp->datasets) (6.0.4)\r\n",
      "Requirement already satisfied: yarl<2.0,>=1.0 in /home/michal/anaconda3/lib/python3.12/site-packages (from aiohttp->datasets) (1.11.0)\r\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /home/michal/anaconda3/lib/python3.12/site-packages (from requests>=2.32.2->datasets) (3.3.2)\r\n",
      "Requirement already satisfied: idna<4,>=2.5 in /home/michal/anaconda3/lib/python3.12/site-packages (from requests>=2.32.2->datasets) (3.7)\r\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /home/michal/anaconda3/lib/python3.12/site-packages (from requests>=2.32.2->datasets) (2024.8.30)\r\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /home/michal/anaconda3/lib/python3.12/site-packages (from jinja2->torch) (2.1.3)\r\n",
      "Requirement already satisfied: click in /home/michal/anaconda3/lib/python3.12/site-packages (from nltk->rouge_score) (8.1.7)\r\n",
      "Requirement already satisfied: joblib in /home/michal/anaconda3/lib/python3.12/site-packages (from nltk->rouge_score) (1.4.2)\r\n",
      "Requirement already satisfied: python-dateutil>=2.8.1 in /home/michal/anaconda3/lib/python3.12/site-packages (from pandas->datasets) (2.9.0.post0)\r\n",
      "Requirement already satisfied: pytz>=2020.1 in /home/michal/anaconda3/lib/python3.12/site-packages (from pandas->datasets) (2024.1)\r\n",
      "Note: you may need to restart the kernel to use updated packages.\n"
     ]
    }
   ],
   "execution_count": 1
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-08T09:32:58.982191Z",
     "start_time": "2024-12-08T09:32:50.589945Z"
    }
   },
   "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-12-08 10:32:55.509478: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:477] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered\n",
      "WARNING: All log messages before absl::InitializeLog() is called are written to STDERR\n",
      "E0000 00:00:1733650375.580009   52899 cuda_dnn.cc:8310] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered\n",
      "E0000 00:00:1733650375.602028   52899 cuda_blas.cc:1418] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered\n",
      "2024-12-08 10:32:55.801899: 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"
     ]
    }
   ],
   "execution_count": 2
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-08T09:33:06.132178Z",
     "start_time": "2024-12-08T09:32:59.038584Z"
    }
   },
   "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": [
    {
     "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": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "execution_count": 3
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": "### Check results with zero-shot",
   "id": "11edd6923d6074a2"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-08T09:33:47.018080Z",
     "start_time": "2024-12-08T09:33:06.146060Z"
    }
   },
   "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": 4
  },
  {
   "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-12-08T09:33:47.148673Z",
     "start_time": "2024-12-08T09:33:47.060459Z"
    }
   },
   "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": [
    {
     "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": 5
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-08T09:33:47.521419Z",
     "start_time": "2024-12-08T09:33:47.514827Z"
    }
   },
   "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": 6
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-12-08T09:33:48.050141Z",
     "start_time": "2024-12-08T09:33:47.571175Z"
    }
   },
   "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": 7
  },
  {
   "metadata": {
    "jupyter": {
     "is_executing": true
    },
    "ExecuteTime": {
     "start_time": "2024-12-08T09:33:48.057928Z"
    }
   },
   "cell_type": "code",
   "source": "trainer.train()",
   "id": "17c6b2c6f8b14c6d",
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001B[34m\u001B[1mwandb\u001B[0m: \u001B[33mWARNING\u001B[0m The `run_name` is currently set to the same value as `TrainingArguments.output_dir`. If this was not intended, please specify a different run name by setting the `TrainingArguments.run_name` parameter.\n",
      "\u001B[34m\u001B[1mwandb\u001B[0m: Using wandb-core as the SDK backend.  Please refer to https://wandb.me/wandb-core for more information.\n",
      "\u001B[34m\u001B[1mwandb\u001B[0m: Currently logged in as: \u001B[33mmichalzarnecki88\u001B[0m (\u001B[33mrzarno\u001B[0m). Use \u001B[1m`wandb login --relogin`\u001B[0m to force relogin\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ],
      "text/html": [
       "Tracking run with wandb version 0.18.6"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ],
      "text/html": [
       "Run data is saved locally in <code>/home/michal/course-generative-ai-python/wandb/run-20241208_103349-brmx7khq</code>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ],
      "text/html": [
       "Syncing run <strong><a href='https://wandb.ai/rzarno/huggingface/runs/brmx7khq' target=\"_blank\">./text-to-sql-translation</a></strong> to <a href='https://wandb.ai/rzarno/huggingface' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/developer-guide' target=\"_blank\">docs</a>)<br/>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ],
      "text/html": [
       " View project at <a href='https://wandb.ai/rzarno/huggingface' target=\"_blank\">https://wandb.ai/rzarno/huggingface</a>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ],
      "text/html": [
       " View run at <a href='https://wandb.ai/rzarno/huggingface/runs/brmx7khq' target=\"_blank\">https://wandb.ai/rzarno/huggingface/runs/brmx7khq</a>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "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
}