prompt_tuning.ipynb

{
    "cells": [
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "This project is an experiment to see how one can incorporate real world knowledge into classical ML models. \n",
                "\n",
                "Deep Learning models are not competitive in tabular datasets. Models like XGBoost are on par with the best DL approaches have to offer but require considerably less training time and can be tuned easily. On the other hand, DL models can levarage large scale language models like GPT-2. These language models have been trained on a trove of information and contain knowledge that can improve downstream performance on many tasks. A typical way DL models incorporate pre-trained language models is by incorporating them as a submodule and training on the whole dataset. Such approaches cannot be easily extended onto classical ML models that dont use gradient descent as the way to optimize them.\n",
                "\n",
                "The \"knowledge\" from language models are usually contained in the embeddings they generate. Since these embeddings are just vectors, we can just use them as the inuts of classical ML models. So how do we get these embeddings? A simple way would be to just convert a row in a table as a sentence and feed it to the language model. But the way we construct our sentence has a big impact on the embeddings generated by the language model. Language models are trained to missing or the next word words given a sentence. If we naively convert a table row to a sentence, the model might not get the task that we are expecting it to do (eg. generating an embedding that can help a decision tree regressor better predict a certain value. Ideally this embedding will contain some external information that can be used to infer our target variable).\n",
                "\n",
                "So how can we get the model to give us the \"right\" knowledge? This is where prompting language models come in. Prompting involves crafting our input in such a way that the it helps the language model get the larger context of what it expects us to do. Prompting has been succesfull in making large language models like GPT-3 perform zero-shot tasks on a variety of tasks that it was never trained for.\n",
                "\n",
                "There are a few ways of crafting such prompts. Here we will be trying out something called **prompt tuning**. Before we dive into prompting, let us first get a dataset for which we can easily incorporate some real world knowledge and build a baseline XGBoost model. "
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "### Imports"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 1,
            "metadata": {},
            "outputs": [
                {
                    "name": "stderr",
                    "output_type": "stream",
                    "text": [
                        "c:\\Users\\rahul\\miniconda3\\envs\\tabular\\lib\\site-packages\\xgboost\\compat.py:36: FutureWarning: pandas.Int64Index is deprecated and will be removed from pandas in a future version. Use pandas.Index with the appropriate dtype instead.\n",
                        "  from pandas import MultiIndex, Int64Index\n"
                    ]
                }
            ],
            "source": [
                "from copy import deepcopy\n",
                "import itertools\n",
                "from functools import partial\n",
                "\n",
                "import polars as pl\n",
                "import numpy as np\n",
                "\n",
                "from sklearn.preprocessing import LabelEncoder\n",
                "from sklearn.model_selection import train_test_split\n",
                "from sklearn.metrics import r2_score, mean_squared_error, mean_absolute_error\n",
                "from sklearn.metrics import mean_squared_error\n",
                "\n",
                "import torch\n",
                "from torch.optim import Adam\n",
                "from torch.nn import CosineEmbeddingLoss\n",
                "from torch.utils.data import DataLoader\n",
                "\n",
                "from transformers import AutoTokenizer, AutoModel\n",
                "from sentence_transformers import SentenceTransformer, util\n",
                "import xgboost\n",
                "\n",
                "# custom modules\n",
                "from data_utils import load_data\n",
                "from model import PrefrozenEmbeddings, embed_sentences\n",
                "from prompt_utils import *\n",
                "\n",
                "from tqdm.notebook import tqdm\n"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "### Load the Dataset"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "> We will be using a library called Polars for loading our datasets into dataframes. Polars is simiar to Pandas but provides a more elegant API"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 2,
            "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>\n",
                            "<th>\n",
                            "Year\n",
                            "</th>\n",
                            "<th>\n",
                            "Month\n",
                            "</th>\n",
                            "<th>\n",
                            "DayofMonth\n",
                            "</th>\n",
                            "<th>\n",
                            "DayOfWeek\n",
                            "</th>\n",
                            "<th>\n",
                            "Origin\n",
                            "</th>\n",
                            "<th>\n",
                            "counts\n",
                            "</th>\n",
                            "<th>\n",
                            "airport\n",
                            "</th>\n",
                            "</tr>\n",
                            "<tr>\n",
                            "<td>\n",
                            "i64\n",
                            "</td>\n",
                            "<td>\n",
                            "i64\n",
                            "</td>\n",
                            "<td>\n",
                            "i64\n",
                            "</td>\n",
                            "<td>\n",
                            "i64\n",
                            "</td>\n",
                            "<td>\n",
                            "str\n",
                            "</td>\n",
                            "<td>\n",
                            "u32\n",
                            "</td>\n",
                            "<td>\n",
                            "str\n",
                            "</td>\n",
                            "</tr>\n",
                            "</thead>\n",
                            "<tbody>\n",
                            "<tr>\n",
                            "<td>\n",
                            "2002\n",
                            "</td>\n",
                            "<td>\n",
                            "1\n",
                            "</td>\n",
                            "<td>\n",
                            "5\n",
                            "</td>\n",
                            "<td>\n",
                            "6\n",
                            "</td>\n",
                            "<td>\n",
                            "\"BPT\"\n",
                            "</td>\n",
                            "<td>\n",
                            "4\n",
                            "</td>\n",
                            "<td>\n",
                            "\"Southeast Texas Regional\"\n",
                            "</td>\n",
                            "</tr>\n",
                            "<tr>\n",
                            "<td>\n",
                            "2002\n",
                            "</td>\n",
                            "<td>\n",
                            "7\n",
                            "</td>\n",
                            "<td>\n",
                            "10\n",
                            "</td>\n",
                            "<td>\n",
                            "3\n",
                            "</td>\n",
                            "<td>\n",
                            "\"BIL\"\n",
                            "</td>\n",
                            "<td>\n",
                            "8\n",
                            "</td>\n",
                            "<td>\n",
                            "\"Billings Logan Intl\"\n",
                            "</td>\n",
                            "</tr>\n",
                            "<tr>\n",
                            "<td>\n",
                            "2002\n",
                            "</td>\n",
                            "<td>\n",
                            "11\n",
                            "</td>\n",
                            "<td>\n",
                            "29\n",
                            "</td>\n",
                            "<td>\n",
                            "5\n",
                            "</td>\n",
                            "<td>\n",
                            "\"BET\"\n",
                            "</td>\n",
                            "<td>\n",
                            "4\n",
                            "</td>\n",
                            "<td>\n",
                            "\"Bethel\"\n",
                            "</td>\n",
                            "</tr>\n",
                            "<tr>\n",
                            "<td>\n",
                            "2003\n",
                            "</td>\n",
                            "<td>\n",
                            "1\n",
                            "</td>\n",
                            "<td>\n",
                            "16\n",
                            "</td>\n",
                            "<td>\n",
                            "4\n",
                            "</td>\n",
                            "<td>\n",
                            "\"SGU\"\n",
                            "</td>\n",
                            "<td>\n",
                            "10\n",
                            "</td>\n",
                            "<td>\n",
                            "\"St George Muni\"\n",
                            "</td>\n",
                            "</tr>\n",
                            "<tr>\n",
                            "<td>\n",
                            "2003\n",
                            "</td>\n",
                            "<td>\n",
                            "6\n",
                            "</td>\n",
                            "<td>\n",
                            "23\n",
                            "</td>\n",
                            "<td>\n",
                            "1\n",
                            "</td>\n",
                            "<td>\n",
                            "\"SMX\"\n",
                            "</td>\n",
                            "<td>\n",
                            "6\n",
                            "</td>\n",
                            "<td>\n",
                            "\"Santa Maria Pub/Capt G Allan Hancock \"\n",
                            "</td>\n",
                            "</tr>\n",
                            "</tbody>\n",
                            "</table>\n",
                            "</div>"
                        ],
                        "text/plain": [
                            "shape: (5, 7)\n",
                            "┌──────┬───────┬────────────┬───────────┬────────┬────────┬─────────────────────────────────────┐\n",
                            "│ Year ┆ Month ┆ DayofMonth ┆ DayOfWeek ┆ Origin ┆ counts ┆ airport                             │\n",
                            "│ ---  ┆ ---   ┆ ---        ┆ ---       ┆ ---    ┆ ---    ┆ ---                                 │\n",
                            "│ i64  ┆ i64   ┆ i64        ┆ i64       ┆ str    ┆ u32    ┆ str                                 │\n",
                            "╞══════╪═══════╪════════════╪═══════════╪════════╪════════╪═════════════════════════════════════╡\n",
                            "│ 2002 ┆ 1     ┆ 5          ┆ 6         ┆ BPT    ┆ 4      ┆ Southeast Texas Regional            │\n",
                            "├╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤\n",
                            "│ 2002 ┆ 7     ┆ 10         ┆ 3         ┆ BIL    ┆ 8      ┆ Billings Logan Intl                 │\n",
                            "├╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤\n",
                            "│ 2002 ┆ 11    ┆ 29         ┆ 5         ┆ BET    ┆ 4      ┆ Bethel                              │\n",
                            "├╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤\n",
                            "│ 2003 ┆ 1     ┆ 16         ┆ 4         ┆ SGU    ┆ 10     ┆ St George Muni                      │\n",
                            "├╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤\n",
                            "│ 2003 ┆ 6     ┆ 23         ┆ 1         ┆ SMX    ┆ 6      ┆ Santa Maria Pub/Capt G Allan Han... │\n",
                            "└──────┴───────┴────────────┴───────────┴────────┴────────┴─────────────────────────────────────┘"
                        ]
                    },
                    "execution_count": 2,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "df = load_data(2002, 2005)\n",
                "df.head()"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "The dataset above contains the number of flights from a specific airport in the United States during a particular date. The number of flights is given by the counts column. Our task is to predict this number, given the date and airport.\n",
                "\n",
                "Origin represents the unique iata indentifier for the airport in question. The actual name of the airport is given in the 'airport' column\n",
                "\n",
                "Right away we can see how external knowledge can help us here. Take an airport like Chicago O'Hare International. It is known to be pretty busy with an average of around 1000 flights coming and going out of the airport per day. Compare that to another airport like Dawson Community Airport, one of the quietest airports in the United States, we can easily guess the average number of flights per day.\n",
                "\n",
                "Let's take an XGBoost Regressor as the baseline and see how well it fares on this dataset. As usual with any ML task, we need to first split the dataset and prep it for the XGBoost model. The following code does this"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Split the dataset into train and test sets"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 3,
            "metadata": {},
            "outputs": [],
            "source": [
                "# Take data from 2002 to 2004 inclusive\n",
                "df_train = df.filter((pl.col('Year') >= 2002) & (pl.col('Year') <= 2004))\n",
                "# Shuffle the dataframe using train_test_split. 0.0 as test_size ensures that all the data is forwarded to the test_split\n",
                "df_train, _ = train_test_split(df_train, test_size=0.1)\n",
                "\n",
                "# Take data from the year 2005 for test split\n",
                "df_test = df.filter((pl.col('Year') <= 2005) & (pl.col('Year') >= 2005))\n",
                "# Shuffle the data and take a random sample of 1/2 of the total data\n",
                "_, df_test = train_test_split(df_test, test_size=0.5)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 4,
            "metadata": {},
            "outputs": [],
            "source": [
                "# Need to make sure that there are no airports present in test split that are not in trian split\n",
                "train_airports = set(df_train.select('Origin').distinct()['Origin'].to_list())\n",
                "test_airports = set(df_test.select('Origin').distinct()['Origin'].to_list())\n",
                "\n",
                "df_test = df_test.filter(\n",
                "    pl.col('Origin').is_in(list(test_airports - train_airports)).is_not()\n",
                ")"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "---"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "### Train a baseline XGBoost model\n",
                "\n",
                "With our train and test sets defined, lets train an XGBoost regression model. For the sake of simplicity, lets just use the default model parameters instead of doing a hyperparameter search"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Featurize the train and test sets"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 5,
            "metadata": {},
            "outputs": [],
            "source": [
                "# We need to convert each airport id from a string to an integer. Using the Label encoder from scikit-learn for this purpose\n",
                "# The label encoder is fit for the whole dataset to prevent OOV when trying to transform the test set\n",
                "all_origins = df.select('Origin').distinct().Origin.to_numpy()\n",
                "origin_encoder = LabelEncoder()\n",
                "origin_encoder.fit(all_origins)\n",
                "\n",
                "# Applying the fitted label encoder to get the featurize the train set\n",
                "X_train = df_train.with_column(\n",
                "    pl.Series('origin_encoded', origin_encoder.transform(df_train.Origin.to_numpy()))\n",
                ").select([\n",
                "    pl.all().exclude(['Year', 'Origin', 'counts', 'airport', 'Month_name'])\n",
                "]).to_numpy()\n",
                "\n",
                "y_train = df_train.counts.to_numpy()\n",
                "\n",
                "\n",
                "# Applying the fitted label encoder to get the featurize the test set\n",
                "X_test = df_test.with_column(\n",
                "    pl.Series('origin_encoded', origin_encoder.transform(df_test.Origin.to_numpy()))\n",
                ").select([\n",
                "    pl.all().exclude(['Year', 'Origin', 'counts', 'airport', 'Month_name'])\n",
                "]).to_numpy()\n",
                "\n",
                "y_test = df_test.counts.to_numpy()"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Train the baseline XGBoost model"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 6,
            "metadata": {},
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "Mean Absolute Error (MSE): \t12.551184803597232 \n",
                        "Mean Squared Error (MAE): \t1082.0375655038551 \n",
                        "R2 Score (R2):\t\t\t0.9519790981716346\n"
                    ]
                }
            ],
            "source": [
                "xgb = xgboost.XGBRegressor(\n",
                "    objective='reg:squarederror',\n",
                "    n_jobs = -1,\n",
                ")\n",
                "\n",
                "model = xgb.fit(X_train, y_train)\n",
                "y_pred = model.predict(X_test)\n",
                "\n",
                "r2 = r2_score(y_test, y_pred)\n",
                "mse = mean_squared_error(y_test, y_pred)\n",
                "mae = mean_absolute_error(y_test, y_pred)\n",
                "\n",
                "print(f'Mean Absolute Error (MSE): \\t{mae} \\n\\\n",
                "Mean Squared Error (MAE): \\t{mse} \\n\\\n",
                "R2 Score (R2):\\t\\t\\t{r2}')"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "The above metrics gives the performance of our baseline model. We want MAE and MSE to decrase while R2 to increase if we have a better model than the above baseline one."
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "---"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "### Prompting\n",
                "\n",
                "Before going into prompt tuning, lets try out the approach of converting a table row into a sentence. Ideally we want to use only columns that makes sense to a language model"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 7,
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "text/plain": [
                            "['Year', 'Month', 'DayofMonth', 'DayOfWeek', 'Origin', 'counts', 'airport']"
                        ]
                    },
                    "execution_count": 7,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "# Print all the columns in the dataframe\n",
                "df.columns"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Right away we can see that the **airport** column is a good candidate for feeding into the language model. Take an airport like *Chicago O'Hare International*. It is known to be pretty busy with an average of around 1000 flights coming and going out of the airport per day. Compare that to another airport like Dawson Community Airport, one of the quietest airports in the United States, we can easily guess the average number of flights per day. \n",
                "\n",
                "A large language model should have come across these associations. If we can somehow craft a sentence that to tease this knowledfe out, we can feed it into an XGBoost model for it to learn."
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "One idea is to craft a of query and a set of value sentences as follows:\n",
                "\n",
                ">**Query**:\n",
                ">\n",
                ">How crowded is the Chicago O'Hare International airport?\n",
                "\n",
                ">**Values**: \n",
                ">- The Chicago O'Hare International airport is very crowded. There are more than 800 flights every day\n",
                ">- The Chicago O'Hare International airport moderately crowded. There around 400 flights every day\n",
                ">- The Chicago O'Hare International airport is slightly crowded. There are around 100 flights every day\n",
                ">- The Chicago O'Hare International airport is not crowded. There are less than 50 flights every day\n",
                "\n",
                "Here we have a set of value sentences that tells the level of *crowdedness* of the airport. We can ask a language model for the value sentence that makes the most sense given the query sentence. Chicago O'Hare International airport is pretty crowded and it has more than 800 flights everyday (on average). This is sort of a Question Answer task. Here the right answer is the 1st sentence in the values list. Hopefully the language model has learnt enough to pick the right one. Also note that the value sentence has an estimate of the value that we want to predict (the number of flights from the airport). This information should help the downstream model that we use for the regression task.\n",
                "\n",
                "So how do we use this in practice and how can we feed the information in the correct value sentence to an XGBoost model?"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "A straightforward (maybe a bit simplistic) approach is to convert each sentence into an embedding. Thus we have a query embedding and a set of value embeddings. Finding the right value sentence can be seen as selecting the value embedding that has the highest cosine similarity score with the query embedding.\n",
                "\n",
                "> This approach of using cosine similarity doesnt always work, hence why I mentioned it is a bit simplistic. We will look at a better approach to prompting later.\n",
                "\n",
                "After getting the right value embedding, we can feed it as part of the input to an XGBoost model. We can use the *sentence_transformers* library for the purpose of embedding query and value sentences. "
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 8,
            "metadata": {},
            "outputs": [],
            "source": [
                "# The main task of the following function is to return a set of embeddings that capture the level of `crowdedness` of airports.\n",
                "# The df variable is a dataframe that can be either the train or test split of the original dataset. This function returns a matrix\n",
                "# having the same number of rows as df. \n",
                "# Each row in the matrix is an embedding that captures the level of `crowdedness` of the airport\n",
                "# present in the same row of the df variable.\n",
                "def generate_embeddings(df, language_model):\n",
                "    # Remove duplicate entries in the df. This is to prevent duplicated effort when embedding sentences\n",
                "    origins = df.select(['Origin', 'airport']).distinct()\n",
                "\n",
                "    embeddings = []\n",
                "    for airport in origins['airport']:\n",
                "        # For every airport, craft the possible query and value sentences. Then embed them using the language model\n",
                "        query = language_model.encode(f'How crowded is the {airport} airport?')\n",
                "        values = language_model.encode([\n",
                "            f'The {airport} airport is very crowded. There are more than 800 flights every day',\n",
                "            f'The {airport} airport is moderately crowded. There around 400 flights every day',\n",
                "            f'The {airport} airport is slightly crowded. There are around 100 flights every day',\n",
                "            f'The {airport} airport is not crowded. There are less than 50 flights every day'\n",
                "        ])\n",
                "\n",
                "        # Find the most similar value embedding to a query and append it to the list\n",
                "        sims = util.cos_sim(query, values)\n",
                "        most_sim_idx = np.argmax(sims)\n",
                "        embeddings.append(values[most_sim_idx])\n",
                "\n",
                "    embeddings = np.vstack(embeddings)\n",
                "    origins = origins.with_column(pl.Series('prompt_embeddings', embeddings))\n",
                "\n",
                "    # Since we had removed duplicate entries before, we need to repopulate them when we eventually merge\n",
                "    # the embeddings data with other dataset features.\n",
                "    origin_embedding_map = {k:v for k,v in origins.select(['Origin', 'prompt_embeddings']).rows()}\n",
                "    embeddings = df.select(\n",
                "        pl.col('Origin').apply(lambda x: origin_embedding_map[x]).alias('embeddings')\n",
                "    )\n",
                "    embeddings = embeddings['embeddings'].to_list()\n",
                "    embeddings = np.vstack(embeddings)\n",
                "    return embeddings"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Generate the airport embeddings for both train and test set. Instead of a large scale model like BERT as our language model, lets use a much smaller one to keep our computational overhead low"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 9,
            "metadata": {},
            "outputs": [],
            "source": [
                "language_model = SentenceTransformer('all-MiniLM-L6-v2', device='cuda', cache_folder='model_cache')\n",
                "X_embs_train = generate_embeddings(df_train, language_model)\n",
                "X_embs_test = generate_embeddings(df_test, language_model)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Featurize the train and test sets and attach the airport embeddings that we got earlier"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 10,
            "metadata": {},
            "outputs": [],
            "source": [
                "all_origins = df.select('Origin').distinct().Origin.to_numpy()\n",
                "origin_encoder = LabelEncoder()\n",
                "origin_encoder.fit(all_origins)\n",
                "\n",
                "X_train = df_train.with_column(\n",
                "    pl.Series('origin_encoded', origin_encoder.transform(df_train.Origin.to_numpy()))\n",
                ").select([\n",
                "    pl.all().exclude(['Year', 'Origin', 'counts', 'airport', 'Month_name'])\n",
                "]).to_numpy()\n",
                "\n",
                "X_train = np.hstack([X_train, X_embs_train])\n",
                "y_train = df_train.counts.to_numpy()\n",
                "\n",
                "\n",
                "X_test = df_test.with_column(\n",
                "    pl.Series('origin_encoded', origin_encoder.transform(df_test.Origin.to_numpy()))\n",
                ").select([\n",
                "    pl.all().exclude(['Year', 'Origin', 'counts', 'airport', 'Month_name'])\n",
                "]).to_numpy()\n",
                "\n",
                "X_test = np.hstack([X_test, X_embs_test])\n",
                "y_test = df_test.counts.to_numpy()"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Train the knowledge enchanged XGBoost model"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 11,
            "metadata": {},
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "Mean Absolute Error (MSE): \t11.725260123551378 \n",
                        "Mean Squared Error (MAE): \t1007.6142761469155 \n",
                        "R2 Score (R2):\t\t\t0.9552820088892394\n"
                    ]
                }
            ],
            "source": [
                "xgb = xgboost.XGBRegressor(\n",
                "    objective='reg:squarederror',\n",
                "    n_jobs=-1)\n",
                "\n",
                "model = xgb.fit(X_train, y_train)\n",
                "y_pred = model.predict(X_test)\n",
                "\n",
                "r2 = r2_score(y_test, y_pred)\n",
                "mse = mean_squared_error(y_test, y_pred)\n",
                "mae = mean_absolute_error(y_test, y_pred)\n",
                "\n",
                "print(f'Mean Absolute Error (MSE): \\t{mae} \\n\\\n",
                "Mean Squared Error (MAE): \\t{mse} \\n\\\n",
                "R2 Score (R2):\\t\\t\\t{r2}')"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "We can already see an improvement in performance by incorporating those embeddings. But this approach comes with a lot of downsides:\n",
                "- How do we make sure that the format of the query and value sentences are okay?\n",
                "- We have defined 4 levels of *crowdedness* for our value sentences. What if we used only 3?\n",
                "- In each level, we have mentioned a threshold. Is that the right one to use?\n",
                "- What if we want to incorporate another attribute like month name?\n",
                "\n",
                "I arrived at the template you saw before after a lot of trial and error. This approach is not feasible in the long run"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "---"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "### Prompt Tuning\n",
                "\n",
                "Till now we have been trying to craft the model to get it to do what we want. Prompt tuning turns the idea around; make the model decide the kind of input it needs to see to get the desired output. This idea was first specified in this [paper](https://arxiv.org/abs/2104.08691) where it was used in the context of NLP tasks. Here we extend this idea onto tabular datasets.\n",
                "\n",
                "<br>\n",
                "\n",
                "So how does it work? Lets start with the same query value approach as last time but we here we use only 1 value statement\n",
                "\n",
                ">Query -  _ _ _ _ _ ; Airport Chicago O'Hare International, Number of flights: _\n",
                "\n",
                ">Value - Airport: Chicago O'Hare International, Number of flights: 800\n",
                "\n",
                "Notice the _ present in the query statement. There are a bunch of them at the start of the sentence and another one at the end. We are going to let the model figure out the way to fill in those blanks in such a way that the cosine similarity of the filled query embedding and the value embedding is maximized. \n",
                "\n",
                "\n",
                "So how do we make the model fill them? The idea is to introduce special trainable embeddings. During the training phase the language model decides how these embeddings should be modified. To make things easier for the language model, each airport has its own set of trainable embeddings. This of these as specific notes that the language model creates for each airport that helps it decide the number of flights. Our query sentence, in practice looks like the one below:\n",
                "\n",
                ">Query: [ORD_1] [ORD_2] [ORD_3] [ORD_4] [ORD_5]; Airport Chicago O'Hare International, Number of flights: [VAL]\n",
                "\n",
                "ORD is the iata code for Chicago O'Hare International airport (the *Origin* column in the dataframe has this information). [ORD_1] is a special token that represents the first trainable embedding for the Chicago O'Hare International airport. [VAL], on the other hand, is a special token to represent the value of the number of flights."
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Generate generate query and value pairs using the prompt tuning template"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 12,
            "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>\n",
                            "<th>\n",
                            "query\n",
                            "</th>\n",
                            "<th>\n",
                            "value\n",
                            "</th>\n",
                            "</tr>\n",
                            "<tr>\n",
                            "<td>\n",
                            "str\n",
                            "</td>\n",
                            "<td>\n",
                            "str\n",
                            "</td>\n",
                            "</tr>\n",
                            "</thead>\n",
                            "<tbody>\n",
                            "<tr>\n",
                            "<td>\n",
                            "\"[MEI_0] [MEI_1] [MEI_2] [MEI_3] [MEI_4] [MEI_5] [MEI_6] [MEI_7]; Airport: Key , Number of flights: [VAL]\"\n",
                            "</td>\n",
                            "<td>\n",
                            "\"Airport: Key , Number of flights: 0\"\n",
                            "</td>\n",
                            "</tr>\n",
                            "<tr>\n",
                            "<td>\n",
                            "\"[LYH_0] [LYH_1] [LYH_2] [LYH_3] [LYH_4] [LYH_5] [LYH_6] [LYH_7]; Airport: Lynchburg Municipal-Preston Glenn , Number of flights: [VAL]\"\n",
                            "</td>\n",
                            "<td>\n",
                            "\"Airport: Lynchburg Municipal-Preston Glenn , Number of flights: 0\"\n",
                            "</td>\n",
                            "</tr>\n",
                            "<tr>\n",
                            "<td>\n",
                            "\"[SMX_0] [SMX_1] [SMX_2] [SMX_3] [SMX_4] [SMX_5] [SMX_6] [SMX_7]; Airport: Santa Maria Pub/Capt G Allan Hancock , Number of flights: [VAL]\"\n",
                            "</td>\n",
                            "<td>\n",
                            "\"Airport: Santa Maria Pub/Capt G Allan Hancock , Number of flights: 0\"\n",
                            "</td>\n",
                            "</tr>\n",
                            "<tr>\n",
                            "<td>\n",
                            "\"[MAF_0] [MAF_1] [MAF_2] [MAF_3] [MAF_4] [MAF_5] [MAF_6] [MAF_7]; Airport: Midland International, Number of flights: [VAL]\"\n",
                            "</td>\n",
                            "<td>\n",
                            "\"Airport: Midland International, Number of flights: 0\"\n",
                            "</td>\n",
                            "</tr>\n",
                            "<tr>\n",
                            "<td>\n",
                            "\"[XNA_0] [XNA_1] [XNA_2] [XNA_3] [XNA_4] [XNA_5] [XNA_6] [XNA_7]; Airport: Northwest Arkansas Regional, Number of flights: [VAL]\"\n",
                            "</td>\n",
                            "<td>\n",
                            "\"Airport: Northwest Arkansas Regional, Number of flights: 0\"\n",
                            "</td>\n",
                            "</tr>\n",
                            "</tbody>\n",
                            "</table>\n",
                            "</div>"
                        ],
                        "text/plain": [
                            "shape: (5, 2)\n",
                            "┌─────────────────────────────────────┬─────────────────────────────────────┐\n",
                            "│ query                               ┆ value                               │\n",
                            "│ ---                                 ┆ ---                                 │\n",
                            "│ str                                 ┆ str                                 │\n",
                            "╞═════════════════════════════════════╪═════════════════════════════════════╡\n",
                            "│ [MEI_0] [MEI_1] [MEI_2] [MEI_3] ... ┆ Airport: Key , Number of flights... │\n",
                            "├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤\n",
                            "│ [LYH_0] [LYH_1] [LYH_2] [LYH_3] ... ┆ Airport: Lynchburg Municipal-Pre... │\n",
                            "├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤\n",
                            "│ [SMX_0] [SMX_1] [SMX_2] [SMX_3] ... ┆ Airport: Santa Maria Pub/Capt G ... │\n",
                            "├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤\n",
                            "│ [MAF_0] [MAF_1] [MAF_2] [MAF_3] ... ┆ Airport: Midland International, ... │\n",
                            "├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤\n",
                            "│ [XNA_0] [XNA_1] [XNA_2] [XNA_3] ... ┆ Airport: Northwest Arkansas Regi... │\n",
                            "└─────────────────────────────────────┴─────────────────────────────────────┘"
                        ]
                    },
                    "execution_count": 12,
                    "metadata": {},
                    "output_type": "execute_result"
                }
            ],
            "source": [
                "query_prompt_format = pl.format(\n",
                "    '{}; Airport: {}, Number of flights: {}', \n",
                "    pl.col('airport_tokens'), pl.col('airport'), pl.lit('[VAL]')\n",
                ")\n",
                "\n",
                "value_prompt_format = pl.format(\n",
                "    'Airport: {}, Number of flights: {}', \n",
                "    pl.col('airport'), pl.col('counts')\n",
                ")\n",
                "\n",
                "num_prompts = 8\n",
                "query_value_prompts_train = generate_query_value_prompts(df_train, query_prompt_format, value_prompt_format, num_prompts)\n",
                "query_value_prompts_test = generate_query_value_prompts(df_test, query_prompt_format, value_prompt_format, num_prompts)\n",
                "\n",
                "query_value_prompts_train.head()"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "The above gives a sample of how the query and value statements look like for some other airports. Here we have set 8 as the number of special trainable embeddings per airport. This can be modified by setting the **num_prompts** variable"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Now we are in a position to download the necessary language model and modify it to accomodate these special embeddings. As before, we will be using a small language model instead of BERT. The way we going to train this model is by freezing everythin in the language model except the special token embeddings that we are going to add into it. I found this method, instead of finetuning the whole model, to be more reliable."
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Download and set up the language model"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 13,
            "metadata": {},
            "outputs": [],
            "source": [
                "airport_tokens = df.select('Origin').distinct()\n",
                "airport_tokens = airport_tokens['Origin']\n",
                "airport_tokens = airport_tokens.apply(partial(airport_token_sequencer, num_prompts=num_prompts)).to_list()\n",
                "airport_tokens = list(itertools.chain.from_iterable(airport_tokens))\n",
                "\n",
                "# Get the language model and tokenizer from huggingface\n",
                "tokenizer = AutoTokenizer.from_pretrained('nreimers/MiniLM-L6-H384-uncased')\n",
                "language_model = AutoModel.from_pretrained('nreimers/MiniLM-L6-H384-uncased')\n",
                "\n",
                "# Add the special tokens into the tokenizer\n",
                "num_added_tokens = tokenizer.add_tokens(airport_tokens, special_tokens=True)    # Add airport tokens\n",
                "num_added_tokens += tokenizer.add_tokens(['[VAL]'], special_tokens=True)        # Add [VAL] token\n",
                "assert num_added_tokens == len(airport_tokens) + 1\n",
                "\n",
                "# create a copy of the original word token embeddings\n",
                "pretrained_word_embeddings = deepcopy(language_model.embeddings.word_embeddings)\n",
                "# call resize_token_embeddings to accomodate the new special tokens\n",
                "language_model.resize_token_embeddings(len(tokenizer))\n",
                "\n",
                "# freeze all layers in the model\n",
                "for param in language_model.parameters():\n",
                "    param.requires_grad = False\n",
                "\n",
                "# create an instance of PrefrozenEmbeddings. This lets us have an embedding layer where only a portion of it is not frozen\n",
                "prefrozen_word_embeddings = PrefrozenEmbeddings(pretrained_word_embeddings, num_added_tokens)\n",
                "# update the instance of word_embeddings to prefrozen ones\n",
                "language_model.embeddings.word_embeddings = prefrozen_word_embeddings"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Set up dataloaders for training"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 14,
            "metadata": {},
            "outputs": [],
            "source": [
                "train_dataloader = DataLoader(\n",
                "    query_value_prompts_train.distinct().to_numpy(),        # type: ignore\n",
                "    shuffle=True, batch_size=128, collate_fn=lambda x: x\n",
                ")\n",
                "      \n",
                "test_dataloader = DataLoader(\n",
                "    query_value_prompts_test.distinct().to_numpy(),         # type: ignore\n",
                "    shuffle=True, batch_size=128, collate_fn=lambda x: x\n",
                ")        "
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 15,
            "metadata": {},
            "outputs": [],
            "source": [
                "def train_step(model, optimizer, criterion, tokenizer, device):\n",
                "    epoch_train_loss = 0.0\n",
                "    model.train()\n",
                "    for batch in train_dataloader:\n",
                "            batch = np.vstack(batch)\n",
                "            query_embeddings = embed_sentences(batch[:, 0].tolist(), model, tokenizer, device)\n",
                "            value_embeddings = embed_sentences(batch[:, 1].tolist(), model, tokenizer, device)\n",
                "            labels = torch.ones(batch.shape[0], dtype=torch.int64).to(device)\n",
                "            loss = criterion(query_embeddings, value_embeddings, labels)\n",
                "\n",
                "            loss.backward()\n",
                "            optimizer.step()\n",
                "            optimizer.zero_grad()\n",
                "            epoch_train_loss += loss.item()\n",
                "    \n",
                "    return epoch_train_loss\n",
                "\n",
                "\n",
                "def valid_step(model, criterion, tokenizer, device):\n",
                "    epoch_valid_loss = 0.0\n",
                "    model.eval()\n",
                "    with torch.no_grad():\n",
                "            for batch in test_dataloader:\n",
                "                batch = np.vstack(batch)\n",
                "                query_embeddings = embed_sentences(batch[:, 0].tolist(), model, tokenizer, device)\n",
                "                value_embeddings = embed_sentences(batch[:, 1].tolist(), model, tokenizer, device)\n",
                "                labels = torch.ones(batch.shape[0], dtype=torch.int64).to(device)\n",
                "                loss = criterion(query_embeddings, value_embeddings, labels)\n",
                "\n",
                "                epoch_valid_loss += loss.item()\n",
                "\n",
                "    return epoch_valid_loss\n",
                "\n",
                "\n",
                "def train(model, optimizer, criterion, device):\n",
                "    train_losses = []\n",
                "    valid_losses = []\n",
                "\n",
                "    num_epocs = 100\n",
                "    for e in tqdm(range(num_epocs)):\n",
                "\n",
                "        epoch_train_loss = train_step(model, optimizer, criterion, tokenizer, device)\n",
                "        epoch_valid_loss = valid_step(model, criterion, tokenizer, device)\n",
                "\n",
                "        train_losses.append(epoch_train_loss / len(train_dataloader))\n",
                "        valid_losses.append(epoch_valid_loss / len(test_dataloader))\n",
                "        \n",
                "        if e%10 == 0 or e == (num_epocs-1) or e == 0:\n",
                "            print(e, train_losses[-1], valid_losses[-1]) "
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Train the language model with the Adam optimizer"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 16,
            "metadata": {},
            "outputs": [
                {
                    "data": {
                        "application/vnd.jupyter.widget-view+json": {
                            "model_id": "c0da8d6674b44808824cccd88020864d",
                            "version_major": 2,
                            "version_minor": 0
                        },
                        "text/plain": [
                            "  0%|          | 0/100 [00:00<?, ?it/s]"
                        ]
                    },
                    "metadata": {},
                    "output_type": "display_data"
                },
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "0 0.1013887698451678 0.0847081442674001\n",
                        "10 0.07361186544100444 0.05333315829435984\n",
                        "20 0.0615144781768322 0.038893599063158035\n",
                        "30 0.05279811720053355 0.030487289031346638\n",
                        "40 0.0488720511396726 0.02507865180571874\n",
                        "50 0.04606983934839567 0.02186124709745248\n",
                        "60 0.04395903026064237 0.01973813648025195\n",
                        "70 0.04311550781130791 0.018111898253361385\n",
                        "80 0.04106645534435908 0.01684059699376424\n",
                        "90 0.04134113217393557 0.015901586040854454\n",
                        "99 0.039788744101921715 0.015044853401680788\n"
                    ]
                }
            ],
            "source": [
                "device = 'cuda:0' if torch.cuda.is_available() else 'cpu'\n",
                "optimizer = Adam(language_model.parameters(), lr=1e-2)\n",
                "criterion = CosineEmbeddingLoss()\n",
                "language_model = language_model.to(device)\n",
                "\n",
                "train(language_model, optimizer, criterion, device)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 17,
            "metadata": {},
            "outputs": [],
            "source": [
                "# The main task of the following function is to return a set of embeddings from the language model that we trained just now.\n",
                "# The df variable is a dataframe that can be either the train or test split of the original dataset. This function returns a matrix\n",
                "# having the same number of rows as df. \n",
                "# Each row in the matrix is an embedding that captures the level of `crowdedness` of the airport present in the same row of the dataframe\n",
                "def generate_embeddings(df, model, tokenizer, device):\n",
                "    \n",
                "    # For every distinct airport in the dataframe, get the appropriate query and value formats for it.\n",
                "    # The query and value entences will be in the format that was described just before.\n",
                "    query_prompts = pl.concat([\n",
                "        df.select('Origin'),\n",
                "        generate_query_value_prompts(\n",
                "            df, query_prompt_format, value_prompt_format, num_prompts\n",
                "        ).select('query')\n",
                "    ], how='horizontal').distinct()\n",
                "    \n",
                "    with torch.no_grad():\n",
                "        # Use the trained model to convert the query sentences into embeddings\n",
                "        prompt_embs_train = embed_sentences(\n",
                "            query_prompts.query.to_list(), \n",
                "            model, tokenizer, device\n",
                "        ).cpu()\n",
                "\n",
                "    # Since we removed duplicate entries to speed up calculations, we need to repopulate them\n",
                "    origin_prompt_emb_map = dict()\n",
                "    for i in range(query_prompts.shape[0]):\n",
                "        origin = query_prompts['Origin'][i]\n",
                "        emb = prompt_embs_train[i, :].numpy()\n",
                "\n",
                "        origin_prompt_emb_map[origin] = emb\n",
                "\n",
                "    return np.vstack(\n",
                "        df.select([\n",
                "            pl.col('Origin').apply(lambda x: origin_prompt_emb_map[x]).alias('emb')\n",
                "        ]).emb.to_list()\n",
                "    )"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Get the query embeddings for both train and test split from the trained model. We will be using them to train the XGBoost model"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 18,
            "metadata": {},
            "outputs": [],
            "source": [
                "X_embs_train = generate_embeddings(df_train, language_model, tokenizer, device)\n",
                "X_embs_test = generate_embeddings(df_test, language_model, tokenizer, device)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Prepare the input features and labels for the XGBoost model"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 19,
            "metadata": {},
            "outputs": [],
            "source": [
                "all_origins = df.select('Origin').distinct().Origin.to_numpy()\n",
                "origin_encoder = LabelEncoder()\n",
                "origin_encoder.fit(all_origins)\n",
                "\n",
                "X_train = df_train.with_column(\n",
                "    pl.Series('origin_encoded', origin_encoder.transform(df_train.Origin.to_numpy()))\n",
                ").select([\n",
                "    pl.all().exclude(['Year', 'Origin', 'counts', 'airport', 'Month_name'])\n",
                "]).to_numpy()\n",
                "\n",
                "X_train = np.hstack([X_train, X_embs_train])\n",
                "y_train = df_train.counts.to_numpy()\n",
                "\n",
                "X_test = df_test.with_column(\n",
                "    pl.Series('origin_encoded', origin_encoder.transform(df_test.Origin.to_numpy()))\n",
                ").select([\n",
                "    pl.all().exclude(['Year', 'Origin', 'counts', 'airport', 'Month_name'])\n",
                "]).to_numpy()\n",
                "\n",
                "X_test = np.hstack([X_test, X_embs_test])\n",
                "y_test = df_test.counts.to_numpy()"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Train the XGBoost model with the extra embeddings"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 20,
            "metadata": {},
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "Mean Absolute Error (MSE): \t11.630789953473045 \n",
                        "Mean Squared Error (MAE): \t1016.9502722153172 \n",
                        "R2 Score (R2):\t\t\t0.9548676767394476\n"
                    ]
                }
            ],
            "source": [
                "xgb = xgboost.XGBRegressor(\n",
                "    objective='reg:squarederror', \n",
                "    n_jobs=-1)\n",
                "\n",
                "model = xgb.fit(X_train, y_train)\n",
                "y_pred = model.predict(X_test)\n",
                "\n",
                "r2 = r2_score(y_test, y_pred)\n",
                "mse = mean_squared_error(y_test, y_pred)\n",
                "mae = mean_absolute_error(y_test, y_pred)\n",
                "\n",
                "print(f'Mean Absolute Error (MSE): \\t{mae} \\n\\\n",
                "Mean Squared Error (MAE): \\t{mse} \\n\\\n",
                "R2 Score (R2):\\t\\t\\t{r2}')"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Looking at the above metrics, prompt tuned XGBoost is pretty close to the case where we hardcoded prompts based XGBoost (it might be slightly better or worse depending on runtime randomness). And we didnt have to go through the hassle of crafting prompts. Just plain old gradient descent to the rescue. So success ?"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "---"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "### Caveats\n",
                "\n",
                "So here is the part where I point out the shortcomings of this approach as well as the tests that we did till now. Here are a few on the top of my head\n",
                "\n",
                "- This whole exercise is done on just a single dataset. It could very much be the case that this approach of using language model embeddings does not generalize\n",
                "- One thing you may have noticed is that I have not performed any sort of data analysis in the begining. These should improve the overall performance of the baseline XGBoost model\n",
                "- For now we have only looked at 1 column for encoding real world knowledge. Its not clear what happens if we try to add more. It could be the case that presenting the language model with more columns might overwhelm it\n",
                "- The embeddings generated from the language model is large compared to the other inputs to the XGBoost model. We could try to its dimensions down using a simple linear layer at the end"
            ]
        }
    ],
    "metadata": {
        "interpreter": {
            "hash": "976d43950396151e7cab9926e73b5abe20542c9ce0b57c070f05867d690adfc3"
        },
        "kernelspec": {
            "display_name": "Python 3.8.13 ('tabular')",
            "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.8.13"
        },
        "orig_nbformat": 4
    },
    "nbformat": 4,
    "nbformat_minor": 2
}