bart.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T11:27:00.809120Z",
     "iopub.status.busy": "2024-11-24T11:27:00.808459Z",
     "iopub.status.idle": "2024-11-24T11:27:19.459522Z",
     "shell.execute_reply": "2024-11-24T11:27:19.458809Z",
     "shell.execute_reply.started": "2024-11-24T11:27:00.809076Z"
    },
    "trusted": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/oumar/Projects/venv/lib64/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "from transformers import BartForConditionalGeneration, BartTokenizer\n",
    "from datasets import load_dataset\n",
    "import torch\n",
    "import pandas as pd\n",
    "import json\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "\n",
    "from peft import get_peft_config, get_peft_model, get_peft_model_state_dict, LoraConfig, TaskType\n",
    "from transformers import default_data_collator, get_linear_schedule_with_warmup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T11:27:19.461706Z",
     "iopub.status.busy": "2024-11-24T11:27:19.460706Z",
     "iopub.status.idle": "2024-11-24T11:27:19.950532Z",
     "shell.execute_reply": "2024-11-24T11:27:19.949672Z",
     "shell.execute_reply.started": "2024-11-24T11:27:19.461666Z"
    },
    "trusted": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Path to dataset files: /home/oumar/.cache/kagglehub/datasets/shahrukhkhan/wikisql/versions/2\n"
     ]
    }
   ],
   "source": [
    "import kagglehub\n",
    "\n",
    "# Download latest version\n",
    "path = kagglehub.dataset_download(\"shahrukhkhan/wikisql\")\n",
    "\n",
    "print(\"Path to dataset files:\", path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T11:27:19.952372Z",
     "iopub.status.busy": "2024-11-24T11:27:19.952112Z",
     "iopub.status.idle": "2024-11-24T11:27:22.638443Z",
     "shell.execute_reply": "2024-11-24T11:27:22.637542Z",
     "shell.execute_reply.started": "2024-11-24T11:27:19.952346Z"
    },
    "trusted": true
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>question</th>\n",
       "      <th>answer</th>\n",
       "      <th>id</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Which sum of week that had an attendance large...</td>\n",
       "      <td>SELECT SUM Week FROM table WHERE Attendance &gt; ...</td>\n",
       "      <td>bd28c679552f44cfad548ab32a60c5d1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>WHAT IS THE POINTS WITH 55 TRIES?</td>\n",
       "      <td>SELECT Points FROM table WHERE Tries for = 55</td>\n",
       "      <td>f848519ba4e44968b135294e2d707922</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>What is 2nd Member, when Assembled is \"30 Marc...</td>\n",
       "      <td>SELECT 2nd member FROM table WHERE Assembled =...</td>\n",
       "      <td>e1153462ca3642fb999f28f476fb9d88</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>What is the title of the king who left office ...</td>\n",
       "      <td>SELECT Title FROM table WHERE Left office = 98...</td>\n",
       "      <td>e89fd713b65f4b6d94d9755674fac9ce</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>What is the address for the assistant principa...</td>\n",
       "      <td>SELECT Address FROM table WHERE Assistant Prin...</td>\n",
       "      <td>9839ab63525c463fa69f4e66d7d8d2bf</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                            question  \\\n",
       "0  Which sum of week that had an attendance large...   \n",
       "1                  WHAT IS THE POINTS WITH 55 TRIES?   \n",
       "2  What is 2nd Member, when Assembled is \"30 Marc...   \n",
       "3  What is the title of the king who left office ...   \n",
       "4  What is the address for the assistant principa...   \n",
       "\n",
       "                                              answer  \\\n",
       "0  SELECT SUM Week FROM table WHERE Attendance > ...   \n",
       "1      SELECT Points FROM table WHERE Tries for = 55   \n",
       "2  SELECT 2nd member FROM table WHERE Assembled =...   \n",
       "3  SELECT Title FROM table WHERE Left office = 98...   \n",
       "4  SELECT Address FROM table WHERE Assistant Prin...   \n",
       "\n",
       "                                 id  \n",
       "0  bd28c679552f44cfad548ab32a60c5d1  \n",
       "1  f848519ba4e44968b135294e2d707922  \n",
       "2  e1153462ca3642fb999f28f476fb9d88  \n",
       "3  e89fd713b65f4b6d94d9755674fac9ce  \n",
       "4  9839ab63525c463fa69f4e66d7d8d2bf  "
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dataset = load_dataset(path)\n",
    "data_train = dataset['train'].shuffle(seed=42).select(range(20000))\n",
    "data_test = dataset['test'].shuffle(seed=42).select(range(100))\n",
    "data_valid = dataset['validation']\n",
    "\n",
    "data_train = pd.DataFrame(data_train)\n",
    "data_test = pd.DataFrame(data_test)\n",
    "data_valid = pd.DataFrame(data_valid)\n",
    "\n",
    "data_train.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T11:27:22.639710Z",
     "iopub.status.busy": "2024-11-24T11:27:22.639448Z",
     "iopub.status.idle": "2024-11-24T11:27:22.644368Z",
     "shell.execute_reply": "2024-11-24T11:27:22.643245Z",
     "shell.execute_reply.started": "2024-11-24T11:27:22.639685Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "def preprocess_data(data):\n",
    "    data = data.dropna()\n",
    "    data = data.drop_duplicates()\n",
    "    data = data.drop(columns=[\"id\"])\n",
    "    data = data.reset_index(drop=True)\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T11:27:22.645766Z",
     "iopub.status.busy": "2024-11-24T11:27:22.645417Z",
     "iopub.status.idle": "2024-11-24T11:27:22.716234Z",
     "shell.execute_reply": "2024-11-24T11:27:22.715405Z",
     "shell.execute_reply.started": "2024-11-24T11:27:22.645729Z"
    },
    "trusted": true
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>question</th>\n",
       "      <th>answer</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Which sum of week that had an attendance large...</td>\n",
       "      <td>SELECT SUM Week FROM table WHERE Attendance &gt; ...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>WHAT IS THE POINTS WITH 55 TRIES?</td>\n",
       "      <td>SELECT Points FROM table WHERE Tries for = 55</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>What is 2nd Member, when Assembled is \"30 Marc...</td>\n",
       "      <td>SELECT 2nd member FROM table WHERE Assembled =...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>What is the title of the king who left office ...</td>\n",
       "      <td>SELECT Title FROM table WHERE Left office = 98...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>What is the address for the assistant principa...</td>\n",
       "      <td>SELECT Address FROM table WHERE Assistant Prin...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                            question  \\\n",
       "0  Which sum of week that had an attendance large...   \n",
       "1                  WHAT IS THE POINTS WITH 55 TRIES?   \n",
       "2  What is 2nd Member, when Assembled is \"30 Marc...   \n",
       "3  What is the title of the king who left office ...   \n",
       "4  What is the address for the assistant principa...   \n",
       "\n",
       "                                              answer  \n",
       "0  SELECT SUM Week FROM table WHERE Attendance > ...  \n",
       "1      SELECT Points FROM table WHERE Tries for = 55  \n",
       "2  SELECT 2nd member FROM table WHERE Assembled =...  \n",
       "3  SELECT Title FROM table WHERE Left office = 98...  \n",
       "4  SELECT Address FROM table WHERE Assistant Prin...  "
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data_train = preprocess_data(data_train)\n",
    "data_test = preprocess_data(data_test)\n",
    "data_valid = preprocess_data(data_valid)\n",
    "\n",
    "data_train.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T11:27:22.717854Z",
     "iopub.status.busy": "2024-11-24T11:27:22.717461Z",
     "iopub.status.idle": "2024-11-24T11:27:22.725616Z",
     "shell.execute_reply": "2024-11-24T11:27:22.724652Z",
     "shell.execute_reply.started": "2024-11-24T11:27:22.717797Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "questions = data_train['question'].tolist()\n",
    "sql_queries = data_train['answer'].tolist()\n",
    "\n",
    "test_questions = data_test['question'].tolist()\n",
    "test_sql_queries = data_test['answer'].tolist()\n",
    "\n",
    "valid_questions = data_valid['question'].tolist()\n",
    "valid_sql_queries = data_valid['answer'].tolist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T11:27:22.726898Z",
     "iopub.status.busy": "2024-11-24T11:27:22.726590Z",
     "iopub.status.idle": "2024-11-24T11:27:37.199621Z",
     "shell.execute_reply": "2024-11-24T11:27:37.198562Z",
     "shell.execute_reply.started": "2024-11-24T11:27:22.726872Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "tokenizer = BartTokenizer.from_pretrained('facebook/bart-large')\n",
    "\n",
    "# OOV token\n",
    "tokenizer.pad_token = tokenizer.unk_token\n",
    "tokenizer.pad_token_id = tokenizer.unk_token_id\n",
    "\n",
    "\n",
    "# padding side\n",
    "tokenizer.pad_padding_side = \"right\"\n",
    "\n",
    "tokenized_input = tokenizer(questions, padding=True, truncation=True, return_tensors=\"pt\")\n",
    "tokenized_output = tokenizer(sql_queries, padding=True, truncation=True, return_tensors=\"pt\")\n",
    "\n",
    "\n",
    "\n",
    "device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n",
    "\n",
    "model = BartForConditionalGeneration.from_pretrained('facebook/bart-large').to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T11:27:37.201757Z",
     "iopub.status.busy": "2024-11-24T11:27:37.200990Z",
     "iopub.status.idle": "2024-11-24T11:27:37.303606Z",
     "shell.execute_reply": "2024-11-24T11:27:37.302868Z",
     "shell.execute_reply.started": "2024-11-24T11:27:37.201713Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "# PEFT\n",
    "peft_config = LoraConfig(\n",
    "    task_type=TaskType.SEQ_2_SEQ_LM,\n",
    "    inference_mode=False,\n",
    "    r=8,\n",
    "    lora_alpha=32,\n",
    "    lora_dropout=0.1,\n",
    ")\n",
    "\n",
    "model = get_peft_model(model, peft_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T11:27:37.306520Z",
     "iopub.status.busy": "2024-11-24T11:27:37.306253Z",
     "iopub.status.idle": "2024-11-24T11:27:37.313756Z",
     "shell.execute_reply": "2024-11-24T11:27:37.312965Z",
     "shell.execute_reply.started": "2024-11-24T11:27:37.306495Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "# num of trainable parameters\n",
    "def count_parameters(model):\n",
    "    return sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "\n",
    "print(f'Number of trainable parameters {count_parameters(model):,}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T11:27:37.315872Z",
     "iopub.status.busy": "2024-11-24T11:27:37.315046Z",
     "iopub.status.idle": "2024-11-24T11:27:37.326460Z",
     "shell.execute_reply": "2024-11-24T11:27:37.325732Z",
     "shell.execute_reply.started": "2024-11-24T11:27:37.315812Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "class SQLDataset(Dataset):\n",
    "    def __init__(self, tokenizer, questions, sql_queries, max_length):\n",
    "        self.tokenizer = tokenizer\n",
    "        self.input_sequence = questions\n",
    "        self.target = sql_queries\n",
    "        self.max_length = max_length\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.input_sequence)\n",
    "    \n",
    "    def __getitem__(self, idx):\n",
    "        input_sequence = self.input_sequence[idx]\n",
    "        target = self.target[idx]\n",
    "        tokenized_input = self.tokenizer(input_sequence, padding=\"max_length\", truncation=True, return_tensors=\"pt\", max_length=self.max_length)\n",
    "        tokenized_output = self.tokenizer(target, padding=\"max_length\", truncation=True, return_tensors=\"pt\", max_length = self.max_length)\n",
    "\n",
    "        return {\n",
    "            \"input_ids\": tokenized_input.input_ids.squeeze(0),\n",
    "            \"attention_mask\": tokenized_input.attention_mask.squeeze(0),\n",
    "            \"labels\": tokenized_output.input_ids.squeeze(0),\n",
    "            \"decoder_attention_mask\": tokenized_output.attention_mask.squeeze(0),\n",
    "            \"query\": target\n",
    "        }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T11:27:37.327682Z",
     "iopub.status.busy": "2024-11-24T11:27:37.327419Z",
     "iopub.status.idle": "2024-11-24T11:27:37.339251Z",
     "shell.execute_reply": "2024-11-24T11:27:37.338469Z",
     "shell.execute_reply.started": "2024-11-24T11:27:37.327659Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "train_dataset = SQLDataset(tokenizer, questions, sql_queries, 512)\n",
    "train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True, collate_fn=default_data_collator, pin_memory=True)\n",
    "\n",
    "test_dataset = SQLDataset(tokenizer, test_questions, test_sql_queries, 512)\n",
    "test_loader = DataLoader(test_dataset, batch_size=8, shuffle=True, collate_fn=default_data_collator, pin_memory=True)\n",
    "\n",
    "valid_dataset = SQLDataset(tokenizer, valid_questions, valid_sql_queries, 512)\n",
    "valid_loader = DataLoader(valid_dataset, batch_size=8, shuffle=True, collate_fn=default_data_collator, pin_memory=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T11:27:37.354272Z",
     "iopub.status.busy": "2024-11-24T11:27:37.354017Z",
     "iopub.status.idle": "2024-11-24T11:27:38.694018Z",
     "shell.execute_reply": "2024-11-24T11:27:38.692860Z",
     "shell.execute_reply.started": "2024-11-24T11:27:37.354248Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "from transformers import Trainer, TrainingArguments\n",
    "\n",
    "\n",
    "training_args = TrainingArguments(\n",
    "    output_dir='./results',          \n",
    "    num_train_epochs=5,             \n",
    "    per_device_train_batch_size=8,  \n",
    "    per_device_eval_batch_size=8,    \n",
    "    warmup_steps=500,                \n",
    "    weight_decay=0.01,               \n",
    "    logging_dir='./logs',           \n",
    "    logging_steps=10,                \n",
    "    eval_strategy=\"epoch\",     \n",
    "    report_to='none'                \n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T11:27:38.695608Z",
     "iopub.status.busy": "2024-11-24T11:27:38.695275Z",
     "iopub.status.idle": "2024-11-24T17:49:21.803752Z",
     "shell.execute_reply": "2024-11-24T17:49:21.802840Z",
     "shell.execute_reply.started": "2024-11-24T11:27:38.695572Z"
    },
    "trusted": true
   },
   "outputs": [],
   "source": [
    "trainer = Trainer(\n",
    "    model = model,\n",
    "    args = training_args,\n",
    "    train_dataset = train_dataset,\n",
    "    eval_dataset = valid_dataset\n",
    ")\n",
    "\n",
    "trainer.train()\n",
    "\n",
    "# Save model\n",
    "model.save_pretrained(\"sql_model\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-11-24T18:03:34.871870Z",
     "iopub.status.busy": "2024-11-24T18:03:34.870973Z",
     "iopub.status.idle": "2024-11-24T18:03:34.895217Z",
     "shell.execute_reply": "2024-11-24T18:03:34.894358Z",
     "shell.execute_reply.started": "2024-11-24T18:03:34.871816Z"
    },
    "trusted": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "BartForConditionalGeneration(\n",
       "  (model): BartModel(\n",
       "    (shared): BartScaledWordEmbedding(50265, 1024, padding_idx=1)\n",
       "    (encoder): BartEncoder(\n",
       "      (embed_tokens): BartScaledWordEmbedding(50265, 1024, padding_idx=1)\n",
       "      (embed_positions): BartLearnedPositionalEmbedding(1026, 1024)\n",
       "      (layers): ModuleList(\n",
       "        (0-11): 12 x BartEncoderLayer(\n",
       "          (self_attn): BartSdpaAttention(\n",
       "            (k_proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "            (v_proj): lora.Linear(\n",
       "              (base_layer): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (lora_dropout): ModuleDict(\n",
       "                (default): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (lora_A): ModuleDict(\n",
       "                (default): Linear(in_features=1024, out_features=8, bias=False)\n",
       "              )\n",
       "              (lora_B): ModuleDict(\n",
       "                (default): Linear(in_features=8, out_features=1024, bias=False)\n",
       "              )\n",
       "              (lora_embedding_A): ParameterDict()\n",
       "              (lora_embedding_B): ParameterDict()\n",
       "              (lora_magnitude_vector): ModuleDict()\n",
       "            )\n",
       "            (q_proj): lora.Linear(\n",
       "              (base_layer): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (lora_dropout): ModuleDict(\n",
       "                (default): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (lora_A): ModuleDict(\n",
       "                (default): Linear(in_features=1024, out_features=8, bias=False)\n",
       "              )\n",
       "              (lora_B): ModuleDict(\n",
       "                (default): Linear(in_features=8, out_features=1024, bias=False)\n",
       "              )\n",
       "              (lora_embedding_A): ParameterDict()\n",
       "              (lora_embedding_B): ParameterDict()\n",
       "              (lora_magnitude_vector): ModuleDict()\n",
       "            )\n",
       "            (out_proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "          )\n",
       "          (self_attn_layer_norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (activation_fn): GELUActivation()\n",
       "          (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "          (final_layer_norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "      )\n",
       "      (layernorm_embedding): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "    (decoder): BartDecoder(\n",
       "      (embed_tokens): BartScaledWordEmbedding(50265, 1024, padding_idx=1)\n",
       "      (embed_positions): BartLearnedPositionalEmbedding(1026, 1024)\n",
       "      (layers): ModuleList(\n",
       "        (0-11): 12 x BartDecoderLayer(\n",
       "          (self_attn): BartSdpaAttention(\n",
       "            (k_proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "            (v_proj): lora.Linear(\n",
       "              (base_layer): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (lora_dropout): ModuleDict(\n",
       "                (default): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (lora_A): ModuleDict(\n",
       "                (default): Linear(in_features=1024, out_features=8, bias=False)\n",
       "              )\n",
       "              (lora_B): ModuleDict(\n",
       "                (default): Linear(in_features=8, out_features=1024, bias=False)\n",
       "              )\n",
       "              (lora_embedding_A): ParameterDict()\n",
       "              (lora_embedding_B): ParameterDict()\n",
       "              (lora_magnitude_vector): ModuleDict()\n",
       "            )\n",
       "            (q_proj): lora.Linear(\n",
       "              (base_layer): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (lora_dropout): ModuleDict(\n",
       "                (default): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (lora_A): ModuleDict(\n",
       "                (default): Linear(in_features=1024, out_features=8, bias=False)\n",
       "              )\n",
       "              (lora_B): ModuleDict(\n",
       "                (default): Linear(in_features=8, out_features=1024, bias=False)\n",
       "              )\n",
       "              (lora_embedding_A): ParameterDict()\n",
       "              (lora_embedding_B): ParameterDict()\n",
       "              (lora_magnitude_vector): ModuleDict()\n",
       "            )\n",
       "            (out_proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "          )\n",
       "          (activation_fn): GELUActivation()\n",
       "          (self_attn_layer_norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (encoder_attn): BartSdpaAttention(\n",
       "            (k_proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "            (v_proj): lora.Linear(\n",
       "              (base_layer): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (lora_dropout): ModuleDict(\n",
       "                (default): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (lora_A): ModuleDict(\n",
       "                (default): Linear(in_features=1024, out_features=8, bias=False)\n",
       "              )\n",
       "              (lora_B): ModuleDict(\n",
       "                (default): Linear(in_features=8, out_features=1024, bias=False)\n",
       "              )\n",
       "              (lora_embedding_A): ParameterDict()\n",
       "              (lora_embedding_B): ParameterDict()\n",
       "              (lora_magnitude_vector): ModuleDict()\n",
       "            )\n",
       "            (q_proj): lora.Linear(\n",
       "              (base_layer): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "              (lora_dropout): ModuleDict(\n",
       "                (default): Dropout(p=0.1, inplace=False)\n",
       "              )\n",
       "              (lora_A): ModuleDict(\n",
       "                (default): Linear(in_features=1024, out_features=8, bias=False)\n",
       "              )\n",
       "              (lora_B): ModuleDict(\n",
       "                (default): Linear(in_features=8, out_features=1024, bias=False)\n",
       "              )\n",
       "              (lora_embedding_A): ParameterDict()\n",
       "              (lora_embedding_B): ParameterDict()\n",
       "              (lora_magnitude_vector): ModuleDict()\n",
       "            )\n",
       "            (out_proj): Linear(in_features=1024, out_features=1024, bias=True)\n",
       "          )\n",
       "          (encoder_attn_layer_norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "          (fc1): Linear(in_features=1024, out_features=4096, bias=True)\n",
       "          (fc2): Linear(in_features=4096, out_features=1024, bias=True)\n",
       "          (final_layer_norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "        )\n",
       "      )\n",
       "      (layernorm_embedding): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)\n",
       "    )\n",
       "  )\n",
       "  (lm_head): Linear(in_features=1024, out_features=50265, bias=False)\n",
       ")"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from transformers import AutoModelForSeq2SeqLM, AutoTokenizer\n",
    "\n",
    "model_tr = AutoModelForSeq2SeqLM.from_pretrained(\"nl2sql_model\")\n",
    "tokenizer_tr = AutoTokenizer.from_pretrained(\"nl2sql_model\")\n",
    "\n",
    "model_tr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/oumar/Projects/venv/lib64/python3.12/site-packages/transformers/generation/utils.py:1375: UserWarning: Using the model-agnostic default `max_length` (=20) to control the generation length. We recommend setting `max_new_tokens` to control the maximum length of the generation.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Question: What is the capital of France?\n",
      "SQL Query: SELECT Capital FROM table WHERE Country = france\n"
     ]
    }
   ],
   "source": [
    "question = \"What is the capital of France?\"\n",
    "inputs = tokenizer_tr(question, return_tensors=\"pt\")\n",
    "outputs = model_tr.generate(**inputs)\n",
    "sql_query = tokenizer_tr.decode(outputs[0], skip_special_tokens=True)\n",
    "\n",
    "print(f\"Question: {question}\")\n",
    "print(f\"SQL Query: {sql_query}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_queries(model, tokenizer, dataset, batch_size = 8):\n",
    "    model.eval()\n",
    "\n",
    "    queries = []\n",
    "    references = []\n",
    "    dataloader = DataLoader(dataset, batch_size=batch_size, pin_memory=True)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for batch in dataloader:\n",
    "            input_ids = batch['input_ids'].to(model.device)\n",
    "            attention_mask = batch['attention_mask'].to(model.device)\n",
    "\n",
    "            outputs = model.generate(input_ids, attention_mask=attention_mask, max_length=512, num_beams=5)\n",
    "            sql_query = [tokenizer.decode(ids, skip_special_tokens=True) for ids in outputs]\n",
    "            queries.extend(sql_query)\n",
    "            references.extend(batch['query'])\n",
    "    return queries, references"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "queries, references = generate_queries(model_tr, tokenizer_tr, test_dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading builder script: 100%|██████████| 5.94k/5.94k [00:00<00:00, 5.15MB/s]\n",
      "Downloading extra modules: 4.07kB [00:00, 1.73MB/s]                   \n",
      "Downloading extra modules: 100%|██████████| 3.34k/3.34k [00:00<00:00, 4.78MB/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'rouge1': np.float64(0.8909919454498688), 'rouge2': np.float64(0.750587669874629), 'rougeL': np.float64(0.8629469773551725), 'rougeLsum': np.float64(0.8621255664342888)}\n"
     ]
    }
   ],
   "source": [
    "import evaluate\n",
    "\n",
    "rouge = evaluate.load(\"rouge\")\n",
    "blue = evaluate.load(\"bleu\")\n",
    "\n",
    "rouge_score = rouge.compute(predictions=queries, references=references)\n",
    "print(rouge_score)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'bleu': 0.5980681639167751, 'precisions': [0.8743500866551126, 0.70521327014218, 0.5910041841004184, 0.4982497082847141], 'brevity_penalty': 0.9161990231034917, 'length_ratio': 0.9195219123505977, 'translation_length': 1154, 'reference_length': 1255}\n"
     ]
    }
   ],
   "source": [
    "blue_score = blue.compute(predictions=queries, references=references)\n",
    "print(blue_score)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Downloading builder script: 100%|██████████| 5.67k/5.67k [00:00<00:00, 2.49MB/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'exact_match': np.float64(0.21)}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "exact_match = evaluate.load(\"exact_match\")\n",
    "\n",
    "exact_match = exact_match.compute(predictions=queries, references=references)\n",
    "print(exact_match)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "{'bleu': 0.5980681639167751, 'precisions': [0.8743500866551126, 0.70521327014218, 0.5910041841004184, 0.4982497082847141], 'brevity_penalty': 0.9161990231034917, 'length_ratio': 0.9195219123505977, 'translation_length': 1154, 'reference_length': 1255}\n"
   ]
  }
 ],
 "metadata": {
  "kaggle": {
   "accelerator": "nvidiaTeslaT4",
   "dataSources": [
    {
     "databundleVersionId": 2452076,
     "datasetId": 1456187,
     "sourceId": 2409983,
     "sourceType": "datasetVersion"
    }
   ],
   "dockerImageVersionId": 30787,
   "isGpuEnabled": true,
   "isInternetEnabled": true,
   "language": "python",
   "sourceType": "notebook"
  },
  "kernelspec": {
   "display_name": "venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}