ch4-sql-corrector.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "################# USUALLY PREINSTALLED IN COLAB #################\n",
    "# !pip install numpy pandas sqlalchemy sqlglot\n",
    "# !pip install torch transformers spacy \n",
    "\n",
    "!pip install Levenshtein\n",
    "!pip install accelerate\n",
    "!pip install bitsandbytes\n",
    "!pip install sentence-transformers\n",
    "!pip install spacy-transformers\n",
    "\n",
    "!python -m spacy download en_core_web_md\n",
    "!python -m spacy download en_core_web_trf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import pandas as pd\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Input data\n",
    "\n",
    "We had previously downloaded the [Young People Survey Dataset](https://www.kaggle.com/datasets/miroslavsabo/young-people-survey) from Kaggle and loaded it into a `sqlite` DB. Let's take a look at the data once more."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "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>Music</th>\n",
       "      <th>Slow_songs_or_fast_songs</th>\n",
       "      <th>Dance</th>\n",
       "      <th>Folk</th>\n",
       "      <th>Country</th>\n",
       "      <th>Classical_music</th>\n",
       "      <th>Musical</th>\n",
       "      <th>Pop</th>\n",
       "      <th>Rock</th>\n",
       "      <th>Metal_or_Hardrock</th>\n",
       "      <th>...</th>\n",
       "      <th>Age</th>\n",
       "      <th>Height</th>\n",
       "      <th>Weight</th>\n",
       "      <th>Number_of_siblings</th>\n",
       "      <th>Gender</th>\n",
       "      <th>Left_right_handed</th>\n",
       "      <th>Education</th>\n",
       "      <th>Only_child</th>\n",
       "      <th>Village_town</th>\n",
       "      <th>House_block_of_flats</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>5.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>5.0</td>\n",
       "      <td>5.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>...</td>\n",
       "      <td>20.0</td>\n",
       "      <td>163.0</td>\n",
       "      <td>48.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>female</td>\n",
       "      <td>right handed</td>\n",
       "      <td>college/bachelor degree</td>\n",
       "      <td>no</td>\n",
       "      <td>village</td>\n",
       "      <td>block of flats</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>4.0</td>\n",
       "      <td>4.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>5.0</td>\n",
       "      <td>4.0</td>\n",
       "      <td>...</td>\n",
       "      <td>19.0</td>\n",
       "      <td>163.0</td>\n",
       "      <td>58.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>female</td>\n",
       "      <td>right handed</td>\n",
       "      <td>college/bachelor degree</td>\n",
       "      <td>no</td>\n",
       "      <td>city</td>\n",
       "      <td>block of flats</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>5.0</td>\n",
       "      <td>5.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>4.0</td>\n",
       "      <td>5.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>5.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>...</td>\n",
       "      <td>20.0</td>\n",
       "      <td>176.0</td>\n",
       "      <td>67.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>female</td>\n",
       "      <td>right handed</td>\n",
       "      <td>secondary school</td>\n",
       "      <td>no</td>\n",
       "      <td>city</td>\n",
       "      <td>block of flats</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>5.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>...</td>\n",
       "      <td>22.0</td>\n",
       "      <td>172.0</td>\n",
       "      <td>59.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>female</td>\n",
       "      <td>right handed</td>\n",
       "      <td>college/bachelor degree</td>\n",
       "      <td>yes</td>\n",
       "      <td>city</td>\n",
       "      <td>house/bungalow</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>5.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>4.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>4.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>5.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>...</td>\n",
       "      <td>20.0</td>\n",
       "      <td>170.0</td>\n",
       "      <td>59.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>female</td>\n",
       "      <td>right handed</td>\n",
       "      <td>secondary school</td>\n",
       "      <td>no</td>\n",
       "      <td>village</td>\n",
       "      <td>house/bungalow</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5 rows × 150 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "   Music  Slow_songs_or_fast_songs  Dance  Folk  Country  Classical_music  \\\n",
       "0    5.0                       3.0    2.0   1.0      2.0              2.0   \n",
       "1    4.0                       4.0    2.0   1.0      1.0              1.0   \n",
       "2    5.0                       5.0    2.0   2.0      3.0              4.0   \n",
       "3    5.0                       3.0    2.0   1.0      1.0              1.0   \n",
       "4    5.0                       3.0    4.0   3.0      2.0              4.0   \n",
       "\n",
       "   Musical  Pop  Rock  Metal_or_Hardrock  ...   Age  Height  Weight  \\\n",
       "0      1.0  5.0   5.0                1.0  ...  20.0   163.0    48.0   \n",
       "1      2.0  3.0   5.0                4.0  ...  19.0   163.0    58.0   \n",
       "2      5.0  3.0   5.0                3.0  ...  20.0   176.0    67.0   \n",
       "3      1.0  2.0   2.0                1.0  ...  22.0   172.0    59.0   \n",
       "4      3.0  5.0   3.0                1.0  ...  20.0   170.0    59.0   \n",
       "\n",
       "   Number_of_siblings  Gender  Left_right_handed                Education  \\\n",
       "0                 1.0  female       right handed  college/bachelor degree   \n",
       "1                 2.0  female       right handed  college/bachelor degree   \n",
       "2                 2.0  female       right handed         secondary school   \n",
       "3                 1.0  female       right handed  college/bachelor degree   \n",
       "4                 1.0  female       right handed         secondary school   \n",
       "\n",
       "   Only_child  Village_town  House_block_of_flats  \n",
       "0          no       village        block of flats  \n",
       "1          no          city        block of flats  \n",
       "2          no          city        block of flats  \n",
       "3         yes          city        house/bungalow  \n",
       "4          no       village        house/bungalow  \n",
       "\n",
       "[5 rows x 150 columns]"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = pd.read_csv('data/young_people_survey.csv')\n",
    "df.columns = df.columns.str.replace('-','')\n",
    "df.columns = df.columns.str.replace(' ','_')\n",
    "df.columns = df.columns.str.replace('__','_')\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Checking the SQLite DB\n",
    "\n",
    "Let's connect to our `sqlite` file and look at its contents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Data loaded from CSV file!\n"
     ]
    }
   ],
   "source": [
    "from sqlalchemy import create_engine\n",
    "engine = create_engine(f\"sqlite:///mysqlitedb.db\")\n",
    "table_name = 'young_people_survey'\n",
    "\n",
    "try:\n",
    "    df = pd.read_csv('data/young_people_survey.csv')\n",
    "    df.columns = df.columns.str.replace('-','')\n",
    "    df.columns = df.columns.str.replace(' ','_')\n",
    "    df.columns = df.columns.str.replace('__','_')\n",
    "    df.to_sql(table_name, engine, index=False)\n",
    "    print('Data loaded from CSV file!')\n",
    "except ValueError as e:\n",
    "    err_msg = e.args[0]\n",
    "    if 'already exists' in err_msg:\n",
    "        print('Table already exists in SQLite DB')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CREATE TABLE young_people_survey (\n",
      "\t\"Music\" FLOAT, \n",
      "\t\"Slow songs or fast songs\" FLOAT, \n",
      "\t\"Dance\" FLOAT, \n",
      "\t\"Folk\" FLOAT, \n",
      "\t\"Country\" FLOAT, \n",
      "\t\"Classical music\" FLOAT, \n",
      "\t\"Musical\" FLOAT, \n",
      "\t\"Pop\" FLOAT, \n",
      "\t\"Rock\" FLOAT, \n",
      "\t\"Metal or Hardrock\" FLOAT, \n",
      "\t\"Punk\" FLOAT, \n",
      "\t\"Hiphop, Rap\" FLOAT, \n",
      "\t\"Reggae, Ska\" FLOAT, \n",
      "\t\"Swing, Jazz\" FLOAT, \n",
      "\t\"Rock n roll\" FLOAT, \n",
      "\t\"Alternative\" FLOAT, \n",
      "\t\"Latino\" FLOAT, \n",
      "\t\"Techno, Trance\" FLOAT, \n",
      "\t\"Opera\" FLOAT, \n",
      "\t\"Movies\" FLOAT, \n",
      "\t\"Horror\" FLOAT, \n",
      "\t\"Thriller\" FLOAT, \n",
      "\t\"Comedy\" FLOAT, \n",
      "\t\"Romantic\" FLOAT, \n",
      "\t\"Sci-fi\" FLOAT, \n",
      "\t\"War\" FLOAT, \n",
      "\t\"Fantasy/Fairy tales\" FLOAT, \n",
      "\t\"Animated\" FLOAT, \n",
      "\t\"Documentary\" FLOAT, \n",
      "\t\"Western\" FLOAT, \n",
      "\t\"Action\" FLOAT, \n",
      "\t\"History\" FLOAT, \n",
      "\t\"Psychology\" FLOAT, \n",
      "\t\"Politics\" FLOAT, \n",
      "\t\"Mathematics\" FLOAT, \n",
      "\t\"Physics\" FLOAT, \n",
      "\t\"Internet\" FLOAT, \n",
      "\t\"PC\" FLOAT, \n",
      "\t\"Economy Management\" FLOAT, \n",
      "\t\"Biology\" FLOAT, \n",
      "\t\"Chemistry\" FLOAT, \n",
      "\t\"Reading\" FLOAT, \n",
      "\t\"Geography\" FLOAT, \n",
      "\t\"Foreign languages\" FLOAT, \n",
      "\t\"Medicine\" FLOAT, \n",
      "\t\"Law\" FLOAT, \n",
      "\t\"Cars\" FLOAT, \n",
      "\t\"Art exhibitions\" FLOAT, \n",
      "\t\"Religion\" FLOAT, \n",
      "\t\"Countryside, outdoors\" FLOAT, \n",
      "\t\"Dancing\" FLOAT, \n",
      "\t\"Musical instruments\" FLOAT, \n",
      "\t\"Writing\" FLOAT, \n",
      "\t\"Passive sport\" FLOAT, \n",
      "\t\"Active sport\" FLOAT, \n",
      "\t\"Gardening\" FLOAT, \n",
      "\t\"Celebrities\" FLOAT, \n",
      "\t\"Shopping\" FLOAT, \n",
      "\t\"Science and technology\" FLOAT, \n",
      "\t\"Theatre\" FLOAT, \n",
      "\t\"Fun with friends\" FLOAT, \n",
      "\t\"Adrenaline sports\" FLOAT, \n",
      "\t\"Pets\" FLOAT, \n",
      "\t\"Flying\" FLOAT, \n",
      "\t\"Storm\" FLOAT, \n",
      "\t\"Darkness\" FLOAT, \n",
      "\t\"Heights\" FLOAT, \n",
      "\t\"Spiders\" FLOAT, \n",
      "\t\"Snakes\" BIGINT, \n",
      "\t\"Rats\" FLOAT, \n",
      "\t\"Ageing\" FLOAT, \n",
      "\t\"Dangerous dogs\" FLOAT, \n",
      "\t\"Fear of public speaking\" FLOAT, \n",
      "\t\"Smoking\" TEXT, \n",
      "\t\"Alcohol\" TEXT, \n",
      "\t\"Healthy eating\" FLOAT, \n",
      "\t\"Daily events\" FLOAT, \n",
      "\t\"Prioritising workload\" FLOAT, \n",
      "\t\"Writing notes\" FLOAT, \n",
      "\t\"Workaholism\" FLOAT, \n",
      "\t\"Thinking ahead\" FLOAT, \n",
      "\t\"Final judgement\" FLOAT, \n",
      "\t\"Reliability\" FLOAT, \n",
      "\t\"Keeping promises\" FLOAT, \n",
      "\t\"Loss of interest\" FLOAT, \n",
      "\t\"Friends versus money\" FLOAT, \n",
      "\t\"Funniness\" FLOAT, \n",
      "\t\"Fake\" FLOAT, \n",
      "\t\"Criminal damage\" FLOAT, \n",
      "\t\"Decision making\" FLOAT, \n",
      "\t\"Elections\" FLOAT, \n",
      "\t\"Self-criticism\" FLOAT, \n",
      "\t\"Judgment calls\" FLOAT, \n",
      "\t\"Hypochondria\" FLOAT, \n",
      "\t\"Empathy\" FLOAT, \n",
      "\t\"Eating to survive\" BIGINT, \n",
      "\t\"Giving\" FLOAT, \n",
      "\t\"Compassion to animals\" FLOAT, \n",
      "\t\"Borrowed stuff\" FLOAT, \n",
      "\t\"Loneliness\" FLOAT, \n",
      "\t\"Cheating in school\" FLOAT, \n",
      "\t\"Health\" FLOAT, \n",
      "\t\"Changing the past\" FLOAT, \n",
      "\t\"God\" FLOAT, \n",
      "\t\"Dreams\" BIGINT, \n",
      "\t\"Charity\" FLOAT, \n",
      "\t\"Number of friends\" BIGINT, \n",
      "\t\"Punctuality\" TEXT, \n",
      "\t\"Lying\" TEXT, \n",
      "\t\"Waiting\" FLOAT, \n",
      "\t\"New environment\" FLOAT, \n",
      "\t\"Mood swings\" FLOAT, \n",
      "\t\"Appearence and gestures\" FLOAT, \n",
      "\t\"Socializing\" FLOAT, \n",
      "\t\"Achievements\" FLOAT, \n",
      "\t\"Responding to a serious letter\" FLOAT, \n",
      "\t\"Children\" FLOAT, \n",
      "\t\"Assertiveness\" FLOAT, \n",
      "\t\"Getting angry\" FLOAT, \n",
      "\t\"Knowing the right people\" FLOAT, \n",
      "\t\"Public speaking\" FLOAT, \n",
      "\t\"Unpopularity\" FLOAT, \n",
      "\t\"Life struggles\" FLOAT, \n",
      "\t\"Happiness in life\" FLOAT, \n",
      "\t\"Energy levels\" FLOAT, \n",
      "\t\"Small - big dogs\" FLOAT, \n",
      "\t\"Personality\" FLOAT, \n",
      "\t\"Finding lost valuables\" FLOAT, \n",
      "\t\"Getting up\" FLOAT, \n",
      "\t\"Interests or hobbies\" FLOAT, \n",
      "\t\"Parents' advice\" FLOAT, \n",
      "\t\"Questionnaires or polls\" FLOAT, \n",
      "\t\"Internet usage\" TEXT, \n",
      "\t\"Finances\" FLOAT, \n",
      "\t\"Shopping centres\" FLOAT, \n",
      "\t\"Branded clothing\" FLOAT, \n",
      "\t\"Entertainment spending\" FLOAT, \n",
      "\t\"Spending on looks\" FLOAT, \n",
      "\t\"Spending on gadgets\" BIGINT, \n",
      "\t\"Spending on healthy eating\" FLOAT, \n",
      "\t\"Age\" FLOAT, \n",
      "\t\"Height\" FLOAT, \n",
      "\t\"Weight\" FLOAT, \n",
      "\t\"Number of siblings\" FLOAT, \n",
      "\t\"Gender\" TEXT, \n",
      "\t\"Left - right handed\" TEXT, \n",
      "\t\"Education\" TEXT, \n",
      "\t\"Only child\" TEXT, \n",
      "\t\"Village - town\" TEXT, \n",
      "\t\"House - block of flats\" TEXT\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "with engine.connect() as conn, conn.begin():\n",
    "    sqlite_master = pd.read_sql_query(\"SELECT * FROM sqlite_master\", conn)\n",
    "sqlite_master['sql_fmt'] = sqlite_master['sql'].apply(lambda z: [x.strip().strip(',').rsplit(' ', maxsplit=1) for x in z.split('\\n')[1:-1]])\n",
    "\n",
    "table_desc_dict = {}\n",
    "for _, row in sqlite_master.iterrows():\n",
    "    table_desc_dict[row['name']] = row['sql']\n",
    "\n",
    "schema = table_desc_dict[table_name]\n",
    "print(schema)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'young_people_survey':                            name   type  comment\n",
       " 0                       \"Music\"  FLOAT      NaN\n",
       " 1    \"Slow songs or fast songs\"  FLOAT      NaN\n",
       " 2                       \"Dance\"  FLOAT      NaN\n",
       " 3                        \"Folk\"  FLOAT      NaN\n",
       " 4                     \"Country\"  FLOAT      NaN\n",
       " ..                          ...    ...      ...\n",
       " 145       \"Left - right handed\"   TEXT      NaN\n",
       " 146                 \"Education\"   TEXT      NaN\n",
       " 147                \"Only child\"   TEXT      NaN\n",
       " 148            \"Village - town\"   TEXT      NaN\n",
       " 149    \"House - block of flats\"   TEXT      NaN\n",
       " \n",
       " [150 rows x 3 columns]}"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "table_desc_df = {}\n",
    "for _, row in sqlite_master.iterrows():\n",
    "    table_desc_df[row['name']] = pd.DataFrame(columns=['name', 'type'], data=row['sql_fmt'])\n",
    "    table_desc_df[row['name']]['comment'] = np.nan\n",
    "\n",
    "table_desc_df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### DDL Pruning and Prompt Generation\n",
    "\n",
    "As we saw in the previous chapter, we prune the table DDL so that the model doesn't get lost inside hundreds of column names and descriptions and only has to look at the most relevant ones to generate the SQL query."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/cloud-user/data/miniconda3/envs/myenv/lib/python3.12/site-packages/sentence_transformers/cross_encoder/CrossEncoder.py:11: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from tqdm.autonotebook import tqdm, trange\n"
     ]
    }
   ],
   "source": [
    "import spacy\n",
    "import torch\n",
    "import sqlglot\n",
    "import numpy as np\n",
    "import os, re, logging, pickle\n",
    "import torch.nn.functional as F\n",
    "from sentence_transformers import SentenceTransformer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def knn_(query: str, all_embs: torch.tensor, top_k: int, threshold: float) -> tuple[torch.tensor, torch.tensor]:\n",
    "    \n",
    "    \"\"\"\n",
    "    Get top most similar columns' embeddings to query using cosine similarity.\n",
    "    \"\"\"\n",
    "    encoder = SentenceTransformer(\"mixedbread-ai/mxbai-embed-large-v1\", device='cpu')\n",
    "    query_emb = encoder.encode(query, convert_to_tensor=True, device='cpu').unsqueeze(0)\n",
    "\n",
    "    similarity_scores = F.cosine_similarity(query_emb, all_embs)\n",
    "    top_results = torch.nonzero(similarity_scores > threshold).squeeze()\n",
    "\n",
    "    # if top_results is empty, return empty tensors\n",
    "    if top_results.numel() == 0:\n",
    "        return torch.tensor([]), torch.tensor([])\n",
    "\n",
    "    # if only 1 result is returned, we need to convert it to a tensor\n",
    "    elif top_results.numel() == 1:\n",
    "        return torch.tensor([similarity_scores[top_results]]), torch.tensor([top_results])\n",
    "    else:\n",
    "        top_k_scores, top_k_indices = torch.topk(similarity_scores[top_results], k=min(top_k, top_results.numel()))\n",
    "        return top_k_scores, top_results[top_k_indices]\n",
    "    \n",
    "    \n",
    "    \n",
    "def format_topk_sql(topk_table_columns: dict[str, list[tuple[str, str, str]]], shuffle: bool) -> str:\n",
    "    if len(topk_table_columns) == 0:\n",
    "        return \"\"\n",
    "\n",
    "    md_str = \"\\n\"\n",
    "    # shuffle the keys in topk_table_columns\n",
    "    table_names = list(topk_table_columns.keys())\n",
    "    if shuffle:\n",
    "        np.random.seed(0)\n",
    "        np.random.shuffle(table_names)\n",
    "    for table_name in table_names:\n",
    "        columns_str = \"\"\n",
    "        columns = topk_table_columns[table_name]\n",
    "        if shuffle:\n",
    "            np.random.seed(0)\n",
    "            np.random.shuffle(columns)\n",
    "        for column_tuple in columns:\n",
    "            if len(column_tuple) > 2:\n",
    "                columns_str += (\n",
    "                    f\"\\n  {column_tuple[0]} {column_tuple[1]}, --{column_tuple[2]}\"\n",
    "                )\n",
    "            else:\n",
    "                columns_str += f\"\\n  {column_tuple[0]} {column_tuple[1]}, \"\n",
    "        md_str += f\"CREATE TABLE {table_name} ({columns_str}\\n);\\n\"\n",
    "    md_str += \"\\n\"\n",
    "    return md_str"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "question = 'Fetch the count of male and female in the data.'\n",
    "EMBEDDING_PATH = 'embs'\n",
    "TOP_K_LIMIT = 25 # number of columns to include in the prompt\n",
    "PRUNE_LIMIT = 5 # minimum number of columns above which a given DDL will be pruned\n",
    "num_cols = 0\n",
    "TAB_DETAILS = []\n",
    "\n",
    "\n",
    "for col in sqlglot.parse_one(schema, dialect='snowflake').find_all(sqlglot.exp.ColumnDef):\n",
    "    num_cols += 1\n",
    "    TAB_DETAILS.append([table_name, col.alias_or_name, col.find(sqlglot.exp.DataType).__str__(), col.find(sqlglot.exp.ColumnConstraint)])\n",
    "        \n",
    "# print(TAB_DETAILS)\n",
    "\n",
    "\n",
    "encoder = SentenceTransformer(\"mixedbread-ai/mxbai-embed-large-v1\", device='cpu')\n",
    "\n",
    "column_descriptions = []\n",
    "column_descriptions_typed = []\n",
    "\n",
    "for row in TAB_DETAILS:\n",
    "    tab_name, col_name, col_dtype, col_desc = row\n",
    "\n",
    "    col_str = f\"{tab_name}.{col_name}:{col_desc}\"\n",
    "    col_str_typed = f\"{tab_name}.{col_name},{col_dtype},{col_desc}\"\n",
    "\n",
    "    column_descriptions.append(col_str)\n",
    "    column_descriptions_typed.append(col_str_typed)\n",
    "column_embs = encoder.encode(column_descriptions, convert_to_tensor=True, device='cpu')\n",
    "\n",
    "\n",
    "# 1a) get top k columns\n",
    "top_k_scores, top_k_indices = knn_(question, column_embs, top_k=5, threshold=0.0)\n",
    "topk_table_columns = {}\n",
    "table_column_names = set()\n",
    "\n",
    "for score, index in zip(top_k_scores, top_k_indices):\n",
    "    table_name, column_info = column_descriptions_typed[index].split(\".\", 1)\n",
    "    column_tuple = re.split(r',\\s*(?![^()]*\\))', column_info, maxsplit=2) #split only on commas outside parantheses\n",
    "    if table_name not in topk_table_columns:\n",
    "        topk_table_columns[table_name] = []\n",
    "    topk_table_columns[table_name].append(column_tuple)\n",
    "    table_column_names.add(f\"{table_name}.{column_tuple[0]}\")\n",
    "    # print(\"INCLUDED by embs: \", column_tuple)\n",
    "\n",
    "# 1b) get columns which match terms in question\n",
    "nlp = spacy.load(\"en_core_web_trf\")\n",
    "question_doc = nlp(question)\n",
    "q_filtered_tokens = [token.lemma_.lower() for token in question_doc if not token.is_stop]\n",
    "q_alpha_tokens = [i for i in q_filtered_tokens if (len(i)>1 and i.isalpha())]\n",
    "\n",
    "\n",
    "TIME_TERMS = ['when', 'time', 'hour', 'minute', 'second', \n",
    "            'day', 'yesterday', 'today', 'tomorrow', \n",
    "            'week', 'month', 'year', \n",
    "            'duration', 'date']\n",
    "\n",
    "time_in_q = False\n",
    "\n",
    "nlp_ner = spacy.load(\"en_core_web_md\")\n",
    "q_ner_doc = nlp_ner(question)\n",
    "ent_types = [w.label_ for w in q_ner_doc.ents]\n",
    "\n",
    "if 'DATE' in ent_types or 'TIME' in ent_types:\n",
    "    time_in_q = True\n",
    "elif any([term in question.lower() for term in TIME_TERMS]):\n",
    "    time_in_q = True\n",
    "elif set(q_alpha_tokens).intersection(set(TIME_TERMS)):\n",
    "    time_in_q = True\n",
    "\n",
    "for col_details in column_descriptions_typed:\n",
    "    table_name, column_info = col_details.split(\".\", 1)\n",
    "    column_tuple = re.split(r',\\s*(?![^()]*\\))', column_info, maxsplit=2) #split only on commas outside parantheses\n",
    "    col_name = column_tuple[0]\n",
    "\n",
    "    if column_tuple in topk_table_columns[table_name]:\n",
    "        # print(\"SKIPPING: \", column_tuple)\n",
    "        continue\n",
    "\n",
    "    # if question concerns time, add time-related columns\n",
    "    if time_in_q and any([timetype in column_tuple[1] for timetype in ['DATE', 'TIMESTAMP']]):\n",
    "        if table_name not in topk_table_columns:\n",
    "            topk_table_columns[table_name] = []\n",
    "        if column_tuple not in topk_table_columns[table_name]:\n",
    "            topk_table_columns[table_name].append(column_tuple)\n",
    "        table_column_names.add(f\"{table_name}.{column_tuple[0]}\")\n",
    "        continue\n",
    "\n",
    "    # if question-token-lemmas overlap with column-token-lemmas, add the column\n",
    "    column_doc = nlp(col_name.replace('_', ' '))\n",
    "    col_tokens = [token.lemma_.lower() for token in column_doc if not token.is_stop]\n",
    "    col_alpha_tokens = [i for i in col_tokens if (len(i)>1 and i.isalpha())]\n",
    "    if set(col_alpha_tokens).intersection(set(q_alpha_tokens)):\n",
    "        if table_name not in topk_table_columns:\n",
    "            topk_table_columns[table_name] = []\n",
    "        if column_tuple not in topk_table_columns[table_name]:\n",
    "            topk_table_columns[table_name].append(column_tuple)\n",
    "        table_column_names.add(f\"{table_name}.{column_tuple[0]}\")\n",
    "\n",
    "# 4) format metadata string\n",
    "pruned_schema = format_topk_sql(topk_table_columns, shuffle=False)\n",
    "print(pruned_schema)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# pruned_schema = '''CREATE TABLE young-people-survey (\n",
    "#   Gender TEXT, --None\n",
    "#   Number of siblings FLOAT, --None\n",
    "#   Height FLOAT, --None\n",
    "#   Age FLOAT, --None\n",
    "#   Alternative FLOAT, --None\n",
    "# );'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "### Task\n",
      "Generate a SQL query to answer [QUESTION]Fetch the count of male and female in the data.[/QUESTION] \n",
      "\n",
      "### Instructions\n",
      "- If you cannot answer the question with the available database schema, return 'I do not know'\n",
      "\n",
      "## Database Schema \n",
      "This query will run on a database whose schema is represented in this string: CREATE TABLE young-people-survey (\n",
      "  Gender TEXT, --None\n",
      "  Number of siblings FLOAT, --None\n",
      "  Height FLOAT, --None\n",
      "  Age FLOAT, --None\n",
      "  Alternative FLOAT, --None\n",
      "); \n",
      "\n",
      "### Answer \n",
      "Given the database schema, here is the SQL query that answers [QUESTION]Fetch the count of male and female in the data.[/QUESTION] [SQL]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "prompt_template = \"\"\"### Task\n",
    "Generate a SQL query to answer [QUESTION]{question}[/QUESTION] \n",
    "\n",
    "### Instructions\n",
    "- If you cannot answer the question with the available database schema, return 'I do not know'\n",
    "\n",
    "## Database Schema \n",
    "This query will run on a database whose schema is represented in this string: {db_schema} \n",
    "\n",
    "### Answer \n",
    "Given the database schema, here is the SQL query that answers [QUESTION]{question}[/QUESTION] [SQL]\n",
    "\"\"\"\n",
    "\n",
    "prompt = prompt_template.format(question=question, db_schema = pruned_schema)\n",
    "print(prompt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prompt the LLM \n",
    "\n",
    "We will use Ollama (running locally or over the network) to prompt our LLM of choice (in this case, we are using `sqlglot-7b-2`)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = \"defog/sqlcoder-7b-2\"\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "\n",
    "# else, load in 8 bits – this is a bit slower\n",
    "model = AutoModelForCausalLM.from_pretrained(\n",
    "    model_name,\n",
    "    trust_remote_code=True,\n",
    "    # torch_dtype=torch.float16,\n",
    "    load_in_8bit=True,\n",
    "    device_map=\"auto\",\n",
    "    use_cache=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_query(quest):\n",
    "    updated_prompt = prompt.format(question=quest)\n",
    "    print('Updated prompt: \\n', updated_prompt)\n",
    "    inputs = tokenizer(updated_prompt, return_tensors=\"pt\").to(\"cuda\")\n",
    "    generated_ids = model.generate(\n",
    "        **inputs,\n",
    "        num_return_sequences=1,\n",
    "        eos_token_id=tokenizer.eos_token_id,\n",
    "        pad_token_id=tokenizer.eos_token_id,\n",
    "        max_new_tokens=400,\n",
    "        do_sample=False,\n",
    "        num_beams=1,\n",
    "    )\n",
    "    outputs = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)\n",
    "\n",
    "    torch.cuda.empty_cache()\n",
    "    torch.cuda.synchronize()\n",
    "    # empty cache so that you do generate more results w/o memory crashing\n",
    "    # particularly important on Colab – memory management is much more straightforward\n",
    "    # when running on an inference service\n",
    "    return outputs[0].split(\"[SQL]\")[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "generated_sql = generate_query(question)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(generated_sql)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fixing LLM hallucinations\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sqlglot import parse_one, exp, parse, table, column, to_identifier\n",
    "from Levenshtein import distance\n",
    "from sentence_transformers import SentenceTransformer\n",
    "import numpy as np\n",
    "from sklearn.metrics.pairwise import cosine_similarity\n",
    "import time\n",
    "\n",
    "class queryPostprocessing:\n",
    "\n",
    "    def __init__(self, query, table_metadata, embedding_model_name):\n",
    "        self.query = query.upper().split(';')[0]\n",
    "        self.table_metadata = table_metadata\n",
    "        self.col_mapping = {}\n",
    "        self.table_mapping = {}\n",
    "        self.embedding_model = SentenceTransformer(embedding_model_name, device='cpu')\n",
    "    \n",
    "    def getIncorrectColumns(self):\n",
    "        column_names_query = [col.name for col in parse_one(self.query).find_all(exp.Column)]\n",
    "        column_names_query = np.unique(column_names_query).tolist()\n",
    "        column_names_table = self.table_metadata['columns']\n",
    " \n",
    "        matching_cols = list(set(column_names_query).intersection(set(column_names_table)))\n",
    "        invalid_cols = list(set(column_names_query).difference(set(matching_cols)))\n",
    "        available_set = list(set(column_names_table).difference(set(matching_cols)))\n",
    "\n",
    "        # check if enough columns are available to substitute (3 columns in query but only 2 in table)\n",
    "        for col in matching_cols:\n",
    "            self.col_mapping[col] = col\n",
    "        \n",
    "        for col in invalid_cols:\n",
    "            closest_col = self.embedding_distance(col, available_set)\n",
    "            self.col_mapping[col] = closest_col\n",
    "            available_set.remove(closest_col)\n",
    "        \n",
    "        return self.col_mapping\n",
    "    \n",
    "    def getIncorrectTables(self):\n",
    "        ast = parse_one(self.query)\n",
    "        table_list = []\n",
    "        for tbl in ast.find_all(exp.Table):\n",
    "            tbl.set('alias',None)\n",
    "            table_name = tbl.sql()\n",
    "            table_list.append(table_name)\n",
    "        \n",
    "        unique_tables = np.unique(table_list).tolist()\n",
    "        if len(unique_tables):\n",
    "            table_name = unique_tables[0]\n",
    "            self.table_mapping[table_name] = self.table_metadata['table_name']\n",
    "        \n",
    "        return self.table_mapping\n",
    "    \n",
    "    def lv_distance(self, col, available_set):\n",
    "        distances = []\n",
    "        for column in available_set:\n",
    "            d = distance(col, column)\n",
    "            distances.append(d)\n",
    "        most_similar_column = available_set[np.argmin(distances)]\n",
    "        return most_similar_column\n",
    "    \n",
    "    def embedding_distance(self, col, available_set):\n",
    "        col_embedding = self.embedding_model.encode([col])\n",
    "        available_set_embedding = self.embedding_model.encode(available_set)\n",
    "        similarity_list = cosine_similarity(col_embedding, available_set_embedding)\n",
    "        most_similar_column = available_set[np.argmax(similarity_list)]\n",
    "        return most_similar_column\n",
    "    \n",
    "    def formatQuery(self):\n",
    "        _ = self.getIncorrectColumns()\n",
    "        _ = self.getIncorrectTables()\n",
    "        updated_query = self.query\n",
    "        \n",
    "        for tbl, updated_tbl in self.table_mapping.items():\n",
    "            updated_query = updated_query.replace(tbl, updated_tbl)\n",
    "        \n",
    "        for col, updated_col in self.col_mapping.items():\n",
    "            updated_query = updated_query.replace(col, updated_col)\n",
    "        \n",
    "        return updated_query\n",
    "    \n",
    "    def formatQuerySQLglot(self):\n",
    "        _ = self.getIncorrectColumns()\n",
    "        _ = self.getIncorrectTables()\n",
    "        query_ast = parse_one(self.query)\n",
    "        alias_cols = [al.alias for al in query_ast.find_all(exp.Alias)]\n",
    "\n",
    "        for col in alias_cols:\n",
    "            self.col_mapping.pop(col, None)\n",
    "        \n",
    "        for tbl in query_ast.find_all(exp.Table):\n",
    "            table_alias = None\n",
    "            if 'alias' in tbl.args:\n",
    "                table_alias = tbl.alias\n",
    "            tbl.set('alias',None)\n",
    "            table_name = tbl.sql()\n",
    "            if table_name in self.table_mapping:\n",
    "                new_table = table(table=self.table_mapping[table_name], quoted=False, alias=table_alias)\n",
    "                tbl.replace(new_table)\n",
    "        \n",
    "        for col in query_ast.find_all(exp.Column):\n",
    "            column_name = col.this.this\n",
    "            if column_name in self.col_mapping:\n",
    "                col.this.set('this',self.col_mapping[column_name])\n",
    "        \n",
    "        return query_ast.sql(dialect='snowflake')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "generated_sql = 'SELECT * from table'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running post-processing...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/cloud-user/data/miniconda3/envs/myenv/lib/python3.12/site-packages/huggingface_hub/file_download.py:1132: 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"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Done! Processed query: SELECT * FROM young-people-survey\n"
     ]
    }
   ],
   "source": [
    "print('Running post-processing...')\n",
    "try:\n",
    "    qp = queryPostprocessing(generated_sql.upper(), \n",
    "                             {'table_name':table_name, 'columns':table_desc_df[table_name]['name'].tolist()}, \n",
    "                             embedding_model_name='mixedbread-ai/mxbai-embed-large-v1')\n",
    "    processed_query = qp.formatQuerySQLglot()\n",
    "except Exception as e:\n",
    "    print(e)\n",
    "    processed_query = generated_sql\n",
    "\n",
    "processed_query = processed_query.replace('\"\"', '\"')\n",
    "print(f\"Done! Processed query: {processed_query}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "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>Music</th>\n",
       "      <th>Slow songs or fast songs</th>\n",
       "      <th>Dance</th>\n",
       "      <th>Folk</th>\n",
       "      <th>Country</th>\n",
       "      <th>Classical music</th>\n",
       "      <th>Musical</th>\n",
       "      <th>Pop</th>\n",
       "      <th>Rock</th>\n",
       "      <th>Metal or Hardrock</th>\n",
       "      <th>...</th>\n",
       "      <th>Age</th>\n",
       "      <th>Height</th>\n",
       "      <th>Weight</th>\n",
       "      <th>Number of siblings</th>\n",
       "      <th>Gender</th>\n",
       "      <th>Left - right handed</th>\n",
       "      <th>Education</th>\n",
       "      <th>Only child</th>\n",
       "      <th>Village - town</th>\n",
       "      <th>House - block of flats</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>5.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>5.0</td>\n",
       "      <td>5.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>...</td>\n",
       "      <td>20.0</td>\n",
       "      <td>163.0</td>\n",
       "      <td>48.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>female</td>\n",
       "      <td>right handed</td>\n",
       "      <td>college/bachelor degree</td>\n",
       "      <td>no</td>\n",
       "      <td>village</td>\n",
       "      <td>block of flats</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>4.0</td>\n",
       "      <td>4.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>5.0</td>\n",
       "      <td>4.0</td>\n",
       "      <td>...</td>\n",
       "      <td>19.0</td>\n",
       "      <td>163.0</td>\n",
       "      <td>58.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>female</td>\n",
       "      <td>right handed</td>\n",
       "      <td>college/bachelor degree</td>\n",
       "      <td>no</td>\n",
       "      <td>city</td>\n",
       "      <td>block of flats</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>5.0</td>\n",
       "      <td>5.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>4.0</td>\n",
       "      <td>5.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>5.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>...</td>\n",
       "      <td>20.0</td>\n",
       "      <td>176.0</td>\n",
       "      <td>67.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>female</td>\n",
       "      <td>right handed</td>\n",
       "      <td>secondary school</td>\n",
       "      <td>no</td>\n",
       "      <td>city</td>\n",
       "      <td>block of flats</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>5.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>...</td>\n",
       "      <td>22.0</td>\n",
       "      <td>172.0</td>\n",
       "      <td>59.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>female</td>\n",
       "      <td>right handed</td>\n",
       "      <td>college/bachelor degree</td>\n",
       "      <td>yes</td>\n",
       "      <td>city</td>\n",
       "      <td>house/bungalow</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>5.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>4.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>4.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>5.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>...</td>\n",
       "      <td>20.0</td>\n",
       "      <td>170.0</td>\n",
       "      <td>59.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>female</td>\n",
       "      <td>right handed</td>\n",
       "      <td>secondary school</td>\n",
       "      <td>no</td>\n",
       "      <td>village</td>\n",
       "      <td>house/bungalow</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1005</th>\n",
       "      <td>5.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>5.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>5.0</td>\n",
       "      <td>4.0</td>\n",
       "      <td>4.0</td>\n",
       "      <td>4.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>...</td>\n",
       "      <td>20.0</td>\n",
       "      <td>164.0</td>\n",
       "      <td>57.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>female</td>\n",
       "      <td>right handed</td>\n",
       "      <td>secondary school</td>\n",
       "      <td>no</td>\n",
       "      <td>city</td>\n",
       "      <td>house/bungalow</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1006</th>\n",
       "      <td>4.0</td>\n",
       "      <td>4.0</td>\n",
       "      <td>5.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>4.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>4.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>...</td>\n",
       "      <td>27.0</td>\n",
       "      <td>183.0</td>\n",
       "      <td>80.0</td>\n",
       "      <td>5.0</td>\n",
       "      <td>male</td>\n",
       "      <td>left handed</td>\n",
       "      <td>masters degree</td>\n",
       "      <td>no</td>\n",
       "      <td>village</td>\n",
       "      <td>house/bungalow</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1007</th>\n",
       "      <td>4.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>4.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>...</td>\n",
       "      <td>18.0</td>\n",
       "      <td>173.0</td>\n",
       "      <td>75.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>female</td>\n",
       "      <td>right handed</td>\n",
       "      <td>secondary school</td>\n",
       "      <td>yes</td>\n",
       "      <td>city</td>\n",
       "      <td>block of flats</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1008</th>\n",
       "      <td>5.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>4.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>...</td>\n",
       "      <td>25.0</td>\n",
       "      <td>173.0</td>\n",
       "      <td>58.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>female</td>\n",
       "      <td>right handed</td>\n",
       "      <td>college/bachelor degree</td>\n",
       "      <td>no</td>\n",
       "      <td>city</td>\n",
       "      <td>block of flats</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1009</th>\n",
       "      <td>5.0</td>\n",
       "      <td>5.0</td>\n",
       "      <td>4.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>4.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>...</td>\n",
       "      <td>21.0</td>\n",
       "      <td>185.0</td>\n",
       "      <td>72.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>male</td>\n",
       "      <td>right handed</td>\n",
       "      <td>secondary school</td>\n",
       "      <td>no</td>\n",
       "      <td>village</td>\n",
       "      <td>house/bungalow</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>1010 rows × 150 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "      Music  Slow songs or fast songs  Dance  Folk  Country  Classical music  \\\n",
       "0       5.0                       3.0    2.0   1.0      2.0              2.0   \n",
       "1       4.0                       4.0    2.0   1.0      1.0              1.0   \n",
       "2       5.0                       5.0    2.0   2.0      3.0              4.0   \n",
       "3       5.0                       3.0    2.0   1.0      1.0              1.0   \n",
       "4       5.0                       3.0    4.0   3.0      2.0              4.0   \n",
       "...     ...                       ...    ...   ...      ...              ...   \n",
       "1005    5.0                       2.0    5.0   2.0      2.0              5.0   \n",
       "1006    4.0                       4.0    5.0   1.0      3.0              4.0   \n",
       "1007    4.0                       3.0    1.0   1.0      2.0              2.0   \n",
       "1008    5.0                       3.0    3.0   3.0      1.0              3.0   \n",
       "1009    5.0                       5.0    4.0   3.0      2.0              3.0   \n",
       "\n",
       "      Musical  Pop  Rock  Metal or Hardrock  ...   Age  Height  Weight  \\\n",
       "0         1.0  5.0   5.0                1.0  ...  20.0   163.0    48.0   \n",
       "1         2.0  3.0   5.0                4.0  ...  19.0   163.0    58.0   \n",
       "2         5.0  3.0   5.0                3.0  ...  20.0   176.0    67.0   \n",
       "3         1.0  2.0   2.0                1.0  ...  22.0   172.0    59.0   \n",
       "4         3.0  5.0   3.0                1.0  ...  20.0   170.0    59.0   \n",
       "...       ...  ...   ...                ...  ...   ...     ...     ...   \n",
       "1005      4.0  4.0   4.0                3.0  ...  20.0   164.0    57.0   \n",
       "1006      1.0  4.0   1.0                1.0  ...  27.0   183.0    80.0   \n",
       "1007      2.0  3.0   4.0                1.0  ...  18.0   173.0    75.0   \n",
       "1008      1.0  3.0   4.0                1.0  ...  25.0   173.0    58.0   \n",
       "1009      3.0  4.0   1.0                1.0  ...  21.0   185.0    72.0   \n",
       "\n",
       "      Number of siblings  Gender  Left - right handed  \\\n",
       "0                    1.0  female         right handed   \n",
       "1                    2.0  female         right handed   \n",
       "2                    2.0  female         right handed   \n",
       "3                    1.0  female         right handed   \n",
       "4                    1.0  female         right handed   \n",
       "...                  ...     ...                  ...   \n",
       "1005                 1.0  female         right handed   \n",
       "1006                 5.0    male          left handed   \n",
       "1007                 0.0  female         right handed   \n",
       "1008                 1.0  female         right handed   \n",
       "1009                 1.0    male         right handed   \n",
       "\n",
       "                    Education  Only child  Village - town  \\\n",
       "0     college/bachelor degree          no         village   \n",
       "1     college/bachelor degree          no            city   \n",
       "2            secondary school          no            city   \n",
       "3     college/bachelor degree         yes            city   \n",
       "4            secondary school          no         village   \n",
       "...                       ...         ...             ...   \n",
       "1005         secondary school          no            city   \n",
       "1006           masters degree          no         village   \n",
       "1007         secondary school         yes            city   \n",
       "1008  college/bachelor degree          no            city   \n",
       "1009         secondary school          no         village   \n",
       "\n",
       "      House - block of flats  \n",
       "0             block of flats  \n",
       "1             block of flats  \n",
       "2             block of flats  \n",
       "3             house/bungalow  \n",
       "4             house/bungalow  \n",
       "...                      ...  \n",
       "1005          house/bungalow  \n",
       "1006          house/bungalow  \n",
       "1007          block of flats  \n",
       "1008          block of flats  \n",
       "1009          house/bungalow  \n",
       "\n",
       "[1010 rows x 150 columns]"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "engine = create_engine(f\"sqlite:///mysqlitedb.db\")\n",
    "with engine.connect() as conn, conn.begin():\n",
    "    query_result = pd.read_sql_query(processed_query, conn)\n",
    "\n",
    "query_result"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}