004_module_ranking.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50a19ce1-bb8a-40c1-8ef3-ff78f60b71d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "#| default_exp ranking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "885c918b-5721-4a0d-947e-62451a6d4c33",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ec96dcb-56ee-412f-bc63-253c5b1a19b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "from nbdev import nbdev_export\n",
    "from fastcore.test import test_eq, test_fail"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14f00408-3ef3-475c-8ab4-eea43088ecae",
   "metadata": {},
   "source": [
    "# ranking\n",
    "> Reference API related to the ranking framework"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e79fb1c-ff24-4887-842d-6230f7c599f7",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-10-23 07:30:31.510166: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n"
     ]
    }
   ],
   "source": [
    "#|export\n",
    "import json\n",
    "import os\n",
    "import os.path\n",
    "from typing import Optional\n",
    "import pandas as pd\n",
    "import tensorflow as tf\n",
    "import tensorflow_ranking as tfr\n",
    "import keras_tuner as kt\n",
    "from tensorflow.keras.layers import Normalization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "71d02afd-0f6c-426e-b9e7-5803c5582095",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def keras_linear_model(\n",
    "    number_documents_per_query,  # Number of documents per query to reshape the listwise prediction.\n",
    "    number_features,  # Number of features used per document.\n",
    ") -> tf.keras.Sequential:  # The uncompiled Keras model.\n",
    "    \"linear model with a lasso constrain on the kernel weights.\"\n",
    "    \n",
    "    model = tf.keras.Sequential()\n",
    "    model.add(\n",
    "        tf.keras.layers.Input(shape=(number_documents_per_query, number_features))\n",
    "    )\n",
    "    model.add(\n",
    "        tf.keras.layers.Dense(\n",
    "            1,\n",
    "            use_bias=False,\n",
    "            activation=None,\n",
    "        )\n",
    "    )\n",
    "    model.add(tf.keras.layers.Reshape((number_documents_per_query,)))\n",
    "    return model\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0cea677-f77e-4d40-b08f-a35c0a4560b8",
   "metadata": {},
   "source": [
    "Usage:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "413cc7d4-c01f-4919-941a-9834319d31bd",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2022-10-23 07:30:36.573173: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n"
     ]
    }
   ],
   "source": [
    "#|warning: false\n",
    "klm = keras_linear_model(\n",
    "    number_documents_per_query=10, \n",
    "    number_features=5\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "027ab379-8972-4a94-a95d-1d87c6950070",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def keras_lasso_linear_model(\n",
    "    number_documents_per_query,  # Number of documents per query to reshape the listwise prediction.\n",
    "    number_features,  # Number of features used per document.\n",
    "    l1_penalty,  # Controls the L1-norm penalty.\n",
    "    normalization_layer: Optional=None,  # Initialized normalization layers. Used when performing feature selection.\n",
    ") -> tf.keras.Sequential:  # The uncompiled Keras model.\n",
    "    \"linear model with a lasso constrain on the kernel weights.\"\n",
    "    \n",
    "    model = tf.keras.Sequential()\n",
    "    model.add(\n",
    "        tf.keras.layers.Input(shape=(number_documents_per_query, number_features))\n",
    "    )\n",
    "    if normalization_layer:\n",
    "        model.add(normalization_layer)\n",
    "    model.add(\n",
    "        tf.keras.layers.Dense(\n",
    "            1,\n",
    "            use_bias=False,\n",
    "            activation=None,\n",
    "            kernel_regularizer=tf.keras.regularizers.L1(l1_penalty),\n",
    "        )\n",
    "    )\n",
    "    model.add(tf.keras.layers.Reshape((number_documents_per_query,)))\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8111ecf3-d62e-408f-87d8-19df6f594af0",
   "metadata": {},
   "source": [
    "Usage:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7822bd6-9419-4014-a16c-f0a632217c0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "kllm = keras_lasso_linear_model(\n",
    "    number_documents_per_query=10, \n",
    "    number_features=5, \n",
    "    l1_penalty=0.01\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7309a497-02b1-427f-831d-767f9fd13e00",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "def keras_ndcg_compiled_model(\n",
    "    model,  # Uncompiled Keras model \n",
    "    learning_rate,  # Learning rate used in the Adagrad optim algo.\n",
    "    top_n  # Top n used when computing the NDCG metric\n",
    "):\n",
    "    \"Compile listwise Keras model with NDCG stateless metric and ApproxNDCGLoss\"\n",
    "    \n",
    "    ndcg = tfr.keras.metrics.NDCGMetric(topn=top_n)\n",
    "\n",
    "    def ndcg_stateless(y_true, y_pred):\n",
    "        ndcg.reset_states()\n",
    "        return ndcg(y_true, y_pred)\n",
    "\n",
    "    optimizer = tf.keras.optimizers.Adagrad(learning_rate)\n",
    "    model.compile(\n",
    "        optimizer=optimizer,\n",
    "        loss=tfr.keras.losses.ApproxNDCGLoss(),\n",
    "        metrics=ndcg_stateless,\n",
    "    )\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "022964ed-7176-4c78-ad3f-300c79982104",
   "metadata": {},
   "source": [
    "Usage:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b932c913-12ba-4955-ab50-946011e6b703",
   "metadata": {},
   "outputs": [],
   "source": [
    "compiled_klm = keras_ndcg_compiled_model(\n",
    "    model=klm, \n",
    "    learning_rate=0.1, \n",
    "    top_n=10\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e84ece11-96d5-4f1e-80ca-ec0f21f7c716",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class LinearHyperModel(kt.HyperModel):\n",
    "    \"\"\"\n",
    "    Define a KerasTuner search space for linear models\n",
    "    \"\"\"          \n",
    "    def __init__(\n",
    "        self,\n",
    "        number_documents_per_query,\n",
    "        number_features,\n",
    "        top_n=10,\n",
    "        learning_rate_range=None,\n",
    "    ):\n",
    "        self.number_documents_per_query = number_documents_per_query\n",
    "        self.number_features = number_features\n",
    "        self.top_n = top_n\n",
    "        if not learning_rate_range:\n",
    "            learning_rate_range = [1e-2, 1e2]\n",
    "        self.learning_rate_range = learning_rate_range\n",
    "        super().__init__()\n",
    "\n",
    "\n",
    "    def build(self, hp):\n",
    "        model = keras_linear_model(\n",
    "            number_documents_per_query=self.number_documents_per_query,\n",
    "            number_features=self.number_features,\n",
    "        )\n",
    "        compiled_model = keras_ndcg_compiled_model(\n",
    "            model=model,\n",
    "            learning_rate=hp.Float(\n",
    "                \"learning_rate\",\n",
    "                min_value=self.learning_rate_range[0],\n",
    "                max_value=self.learning_rate_range[1],\n",
    "                sampling=\"log\",\n",
    "            ),\n",
    "            top_n=self.top_n,\n",
    "        )\n",
    "        return compiled_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9156480d-b877-40e5-acb8-7c1c065c1dfe",
   "metadata": {},
   "outputs": [],
   "source": [
    "linear_hyper_model = LinearHyperModel(\n",
    "    number_documents_per_query=10, \n",
    "    number_features=10, \n",
    "    top_n=10, \n",
    "    learning_rate_range=[1e-2, 1e2]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ae80468-6aca-42b3-bcde-747806f7409f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class LassoHyperModel(kt.HyperModel):\n",
    "    \"\"\"\n",
    "    Define a KerasTuner search space for lasso models\n",
    "    \"\"\"              \n",
    "    def __init__(\n",
    "        self,\n",
    "        number_documents_per_query,\n",
    "        number_features,\n",
    "        trained_normalization_layer,\n",
    "        top_n=10,\n",
    "        l1_penalty_range=None,\n",
    "        learning_rate_range=None,\n",
    "    ):\n",
    "        self.number_documents_per_query = number_documents_per_query\n",
    "        self.number_features = number_features\n",
    "        self.trained_normalization_layer = trained_normalization_layer\n",
    "        self.top_n = top_n\n",
    "        if not l1_penalty_range:\n",
    "            l1_penalty_range = [1e-4, 1e-2]\n",
    "        self.l1_penalty_range = l1_penalty_range\n",
    "        if not learning_rate_range:\n",
    "            learning_rate_range = [1e-2, 1e2]\n",
    "        self.learning_rate_range = learning_rate_range\n",
    "        super().__init__()\n",
    "\n",
    "    def build(self, hp):\n",
    "        model = keras_lasso_linear_model(\n",
    "            number_documents_per_query=self.number_documents_per_query,\n",
    "            number_features=self.number_features,\n",
    "            l1_penalty=hp.Float(\n",
    "                \"lambda\",\n",
    "                min_value=self.l1_penalty_range[0],\n",
    "                max_value=self.l1_penalty_range[1],\n",
    "                sampling=\"log\",\n",
    "            ),\n",
    "            normalization_layer=self.trained_normalization_layer,\n",
    "        )\n",
    "        compiled_model = keras_ndcg_compiled_model(\n",
    "            model=model,\n",
    "            learning_rate=hp.Float(\n",
    "                \"learning_rate\",\n",
    "                min_value=self.learning_rate_range[0],\n",
    "                max_value=self.learning_rate_range[1],\n",
    "                sampling=\"log\",\n",
    "            ),\n",
    "            top_n=self.top_n,\n",
    "        )\n",
    "        return compiled_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "201e6305-9848-44fa-b71b-bbba3697f100",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class ListwiseRankingFramework:\n",
    "    def __init__(\n",
    "        self,\n",
    "        number_documents_per_query,\n",
    "        batch_size=32,\n",
    "        shuffle_buffer_size=1000,\n",
    "        tuner_max_trials=3,\n",
    "        tuner_executions_per_trial=1,\n",
    "        tuner_epochs=1,\n",
    "        tuner_early_stop_patience=None,\n",
    "        final_epochs=1,\n",
    "        top_n=10,\n",
    "        l1_penalty_range=None,\n",
    "        learning_rate_range=None,\n",
    "        folder_dir=os.getcwd(),\n",
    "    ):\n",
    "        \"Listwise ranking framework\"\n",
    "        self.number_documents_per_query = number_documents_per_query\n",
    "        self.batch_size = batch_size\n",
    "        self.shuffle_buffer_size = shuffle_buffer_size\n",
    "        self.tuner_max_trials = tuner_max_trials\n",
    "        self.tuner_executions_per_trial = tuner_executions_per_trial\n",
    "        self.tuner_epochs = tuner_epochs\n",
    "        self.tuner_early_stop_patience = tuner_early_stop_patience\n",
    "        self.final_epochs = final_epochs\n",
    "        self.top_n = top_n\n",
    "        self.l1_penalty_range = l1_penalty_range\n",
    "        self.learning_rate_range = learning_rate_range\n",
    "        self.folder_dir = folder_dir\n",
    "\n",
    "        self.query_id_name = \"query_id\"\n",
    "        self.target_name = \"label\"\n",
    "        self.distribute_strategy = tf.distribute.MirroredStrategy()\n",
    "        \n",
    "\n",
    "    def listwise_tf_dataset_from_df(\n",
    "        self, df, feature_names, shuffle_buffer_size, batch_size\n",
    "    ):\n",
    "        \"\"\"\n",
    "        Create TensorFlow dataframe suited for listwise loss function from pandas df.\n",
    "\n",
    "        :param df: Pandas df containing the data.\n",
    "        :param feature_names: Features to be used in the tensorflow model.\n",
    "        :param shuffle_buffer_size: The size of the buffer used to sample data from.\n",
    "        :param batch_size: The size of the batch for each sample from the dataset.\n",
    "        :return: TF dataset\n",
    "        \"\"\"\n",
    "        ds = tf.data.Dataset.from_tensor_slices(\n",
    "            {\n",
    "                \"features\": tf.cast(df[feature_names].values, tf.float32),\n",
    "                \"label\": tf.cast(df[self.target_name].values, tf.float32),\n",
    "                \"query_id\": tf.cast(df[self.query_id_name].values, tf.int64),\n",
    "            }\n",
    "        )\n",
    "\n",
    "        key_func = lambda x: x[self.query_id_name]\n",
    "        reduce_func = lambda key, dataset: dataset.batch(\n",
    "            self.number_documents_per_query, drop_remainder=True\n",
    "        )\n",
    "        listwise_ds = ds.group_by_window(\n",
    "            key_func=key_func,\n",
    "            reduce_func=reduce_func,\n",
    "            window_size=self.number_documents_per_query,\n",
    "        )\n",
    "        listwise_ds = listwise_ds.map(lambda x: (x[\"features\"], x[\"label\"]))\n",
    "        listwise_ds = listwise_ds.shuffle(buffer_size=shuffle_buffer_size).batch(\n",
    "            batch_size=batch_size\n",
    "        )\n",
    "        return listwise_ds\n",
    "\n",
    "    def listwise_tf_dataset_from_csv(\n",
    "        self, file_path, feature_names, shuffle_buffer_size, batch_size\n",
    "    ):\n",
    "        \"\"\"\n",
    "        Create TensorFlow dataframe suited for listwise loss function from a .csv file.\n",
    "\n",
    "        :param file_path: The path to the csv file.\n",
    "        :param feature_names: Features to be used in the tensorflow model.\n",
    "        :param shuffle_buffer_size: The size of the buffer used to sample data from.\n",
    "        :param batch_size: The size of the batch for each sample from the dataset.\n",
    "        :return: TF dataset\n",
    "        \"\"\"\n",
    "        ds = tf.data.experimental.make_csv_dataset(\n",
    "            file_path,\n",
    "            batch_size=1,\n",
    "            num_epochs=1,\n",
    "            shuffle_buffer_size=shuffle_buffer_size,\n",
    "        )\n",
    "\n",
    "        def create_dict_slices(x):\n",
    "            return {\n",
    "                \"query_id\": tf.reshape(tf.cast(x[\"query_id\"], tf.int64), []),\n",
    "                \"label\": tf.reshape(tf.cast(x[\"label\"], tf.float32), []),\n",
    "                \"features\": tf.cast(\n",
    "                    tf.reshape(\n",
    "                        [x[name] for name in feature_names], [len(feature_names)]\n",
    "                    ),\n",
    "                    tf.float32,\n",
    "                ),\n",
    "            }\n",
    "\n",
    "        ds_mapped = ds.map(lambda x: create_dict_slices(x))\n",
    "        key_func = lambda x: x[self.query_id_name]\n",
    "        reduce_func = lambda key, dataset: dataset.batch(\n",
    "            self.number_documents_per_query, drop_remainder=True\n",
    "        )\n",
    "        listwise_ds = ds_mapped.group_by_window(\n",
    "            key_func=key_func,\n",
    "            reduce_func=reduce_func,\n",
    "            window_size=self.number_documents_per_query,\n",
    "        )\n",
    "        listwise_ds = listwise_ds.map(lambda x: (x[\"features\"], x[\"label\"]))\n",
    "        listwise_ds = listwise_ds.batch(batch_size=batch_size)\n",
    "        return listwise_ds\n",
    "\n",
    "    def create_dataset(self, df_or_file, feature_names):\n",
    "        if isinstance(df_or_file, pd.DataFrame):\n",
    "            ds = self.listwise_tf_dataset_from_df(\n",
    "                df=df_or_file,\n",
    "                feature_names=feature_names,\n",
    "                shuffle_buffer_size=self.shuffle_buffer_size,\n",
    "                batch_size=self.batch_size,\n",
    "            )\n",
    "        else:\n",
    "            ds = self.listwise_tf_dataset_from_csv(\n",
    "                file_path=df_or_file,\n",
    "                feature_names=feature_names,\n",
    "                shuffle_buffer_size=self.shuffle_buffer_size,\n",
    "                batch_size=self.batch_size,\n",
    "            )\n",
    "        return ds\n",
    "\n",
    "    def create_and_train_normalization_layer(self, train_ds):\n",
    "        normalization_layer = tf.keras.layers.Normalization()\n",
    "        train_feature_ds = train_ds.map(lambda x, y: x)\n",
    "        normalization_layer.adapt(train_feature_ds)\n",
    "        return normalization_layer\n",
    "\n",
    "    def tune_model(self, model, train_ds, dev_ds):\n",
    "        tuner = kt.RandomSearch(\n",
    "            model,\n",
    "            objective=kt.Objective(\"val_ndcg_stateless\", direction=\"max\"),\n",
    "            directory=self.folder_dir,\n",
    "            project_name=\"keras_tuner\",\n",
    "            distribution_strategy=self.distribute_strategy,\n",
    "            overwrite=True,\n",
    "            max_trials=self.tuner_max_trials,\n",
    "            executions_per_trial=self.tuner_executions_per_trial,\n",
    "        )\n",
    "        callbacks = []\n",
    "        if self.tuner_early_stop_patience:\n",
    "            early_stopping_callback = tf.keras.callbacks.EarlyStopping(\n",
    "                monitor=\"val_ndcg_stateless\",\n",
    "                patience=self.tuner_early_stop_patience,\n",
    "                mode=\"max\",\n",
    "            )\n",
    "            callbacks.append(early_stopping_callback)\n",
    "        tuner.search(\n",
    "            train_ds,\n",
    "            validation_data=dev_ds,\n",
    "            epochs=self.tuner_epochs,\n",
    "            callbacks=callbacks,\n",
    "        )\n",
    "        return tuner.get_best_hyperparameters()[0]\n",
    "\n",
    "    def fit_linear_model(\n",
    "        self, train_data, dev_data, feature_names, hyperparameters=None\n",
    "    ):\n",
    "\n",
    "        number_features = len(feature_names)\n",
    "\n",
    "        train_ds = self.create_dataset(\n",
    "            df_or_file=train_data, feature_names=feature_names\n",
    "        )\n",
    "        dev_ds = self.create_dataset(df_or_file=dev_data, feature_names=feature_names)\n",
    "        with self.distribute_strategy.scope():\n",
    "            linear_hyper_model = LinearHyperModel(\n",
    "                number_documents_per_query=self.number_documents_per_query,\n",
    "                number_features=number_features,\n",
    "                top_n=self.top_n,\n",
    "                learning_rate_range=self.learning_rate_range,\n",
    "            )\n",
    "        if not hyperparameters:\n",
    "            best_hps = self.tune_model(\n",
    "                model=linear_hyper_model, train_ds=train_ds, dev_ds=dev_ds\n",
    "            )\n",
    "            best_hyperparams = best_hps.values\n",
    "        else:\n",
    "            best_hyperparams = hyperparameters\n",
    "            best_hps = kt.HyperParameters()\n",
    "            best_hps.values = hyperparameters\n",
    "        model = linear_hyper_model.build(best_hps)\n",
    "        model.fit(\n",
    "            train_ds,\n",
    "            validation_data=dev_ds,\n",
    "            epochs=self.final_epochs,\n",
    "        )\n",
    "        weights = model.get_weights()\n",
    "        weights = {\n",
    "            \"feature_names\": feature_names,\n",
    "            \"linear_model_weights\": [\n",
    "                float(weights[0][idx][0]) for idx in range(len(feature_names))\n",
    "            ],\n",
    "        }\n",
    "        eval_result_from_fit = model.history.history[\"val_ndcg_stateless\"][-1]\n",
    "\n",
    "        return weights, eval_result_from_fit, best_hyperparams\n",
    "\n",
    "    def fit_lasso_linear_model(\n",
    "        self, train_data, dev_data, feature_names, hyperparameters=None\n",
    "    ):\n",
    "\n",
    "        number_features = len(feature_names)\n",
    "        train_ds = self.create_dataset(\n",
    "            df_or_file=train_data, feature_names=feature_names\n",
    "        )\n",
    "        dev_ds = self.create_dataset(df_or_file=dev_data, feature_names=feature_names)\n",
    "        with self.distribute_strategy.scope():\n",
    "            trained_normalization_layer = self.create_and_train_normalization_layer(\n",
    "                train_ds=train_ds\n",
    "            )\n",
    "            lasso_hyper_model = LassoHyperModel(\n",
    "                number_documents_per_query=self.number_documents_per_query,\n",
    "                number_features=number_features,\n",
    "                trained_normalization_layer=trained_normalization_layer,\n",
    "                top_n=self.top_n,\n",
    "                l1_penalty_range=self.l1_penalty_range,\n",
    "                learning_rate_range=self.learning_rate_range,\n",
    "            )\n",
    "        if not hyperparameters:\n",
    "            best_hps = self.tune_model(\n",
    "                model=lasso_hyper_model, train_ds=train_ds, dev_ds=dev_ds\n",
    "            )\n",
    "            best_hyperparams = best_hps.values\n",
    "        else:\n",
    "            best_hyperparams = hyperparameters\n",
    "            best_hps = kt.HyperParameters()\n",
    "            best_hps.values = hyperparameters\n",
    "        model = lasso_hyper_model.build(best_hps)\n",
    "        model.fit(\n",
    "            train_ds,\n",
    "            validation_data=dev_ds,\n",
    "            epochs=self.final_epochs,\n",
    "        )\n",
    "        weights = model.get_weights()\n",
    "        weights = {\n",
    "            \"feature_names\": feature_names,\n",
    "            \"normalization_mean\": weights[0].tolist(),\n",
    "            \"normalization_sd\": weights[1].tolist(),\n",
    "            \"normalization_number_data\": int(weights[2]),\n",
    "            \"linear_model_weights\": [\n",
    "                float(weights[3][idx][0]) for idx in range(len(feature_names))\n",
    "            ],\n",
    "        }\n",
    "        eval_result_from_fit = model.history.history[\"val_ndcg_stateless\"][-1]\n",
    "\n",
    "        return weights, eval_result_from_fit, best_hyperparams\n",
    "\n",
    "    def lasso_model_search(\n",
    "        self,\n",
    "        train_data,\n",
    "        dev_data,\n",
    "        feature_names,\n",
    "        protected_features=None,\n",
    "        hyperparameter=None,\n",
    "        output_file=\"lasso_model_search.json\",\n",
    "    ):\n",
    "\n",
    "        output_file = os.path.join(self.folder_dir, output_file)\n",
    "        try:\n",
    "            with open(output_file, \"r\") as f:\n",
    "                results = json.load(f)\n",
    "                print(\"Lasso model search: Results from output file loaded.\")\n",
    "        except FileNotFoundError:\n",
    "            print(\"Lasso model search: File not found. Starting search from scratch.\")\n",
    "            results = []\n",
    "\n",
    "        if not protected_features:\n",
    "            protected_features = []\n",
    "        while (len(feature_names) >= len(protected_features)) and len(\n",
    "            feature_names\n",
    "        ) > 0:\n",
    "            (weights, evaluation, best_hyperparams) = self.fit_lasso_linear_model(\n",
    "                train_data=train_data,\n",
    "                dev_data=dev_data,\n",
    "                feature_names=feature_names,\n",
    "                hyperparameters=hyperparameter,\n",
    "            )\n",
    "            partial_result = {\n",
    "                \"evaluation\": evaluation,\n",
    "                \"weights\": weights,\n",
    "                \"best_hyperparams\": best_hyperparams,\n",
    "            }\n",
    "            results.append(partial_result)\n",
    "            with open(output_file, \"w\") as f:\n",
    "                json.dump(results, f)\n",
    "\n",
    "            weights = {\n",
    "                feature_name: float(model_weight)\n",
    "                for feature_name, model_weight in zip(\n",
    "                    weights[\"feature_names\"], weights[\"linear_model_weights\"]\n",
    "                )\n",
    "            }\n",
    "            print({k: round(weights[k], 2) for k in weights})\n",
    "            print(evaluation)\n",
    "\n",
    "            abs_weights = {k: abs(weights[k]) for k in weights}\n",
    "            if protected_features:\n",
    "                abs_weights = {\n",
    "                    k: abs_weights[k]\n",
    "                    for k in abs_weights\n",
    "                    if k not in protected_features\n",
    "                }\n",
    "            if len(abs_weights) > 0:\n",
    "                worst_feature = min(abs_weights, key=abs_weights.get)\n",
    "                feature_names = [x for x in feature_names if x != worst_feature]\n",
    "            else:\n",
    "                break\n",
    "\n",
    "        return results\n",
    "\n",
    "    def _forward_selection_iteration(\n",
    "        self, train_data, dev_data, feature_names, hyperparameter=None\n",
    "    ):\n",
    "        (weights, evaluation, best_hyperparams) = self.fit_lasso_linear_model(\n",
    "            train_data=train_data,\n",
    "            dev_data=dev_data,\n",
    "            feature_names=feature_names,\n",
    "            hyperparameters=hyperparameter,\n",
    "        )\n",
    "        partial_result = {\n",
    "            \"number_features\": len(feature_names),\n",
    "            \"evaluation\": evaluation,\n",
    "            \"weights\": weights,\n",
    "            \"best_hyperparams\": best_hyperparams,\n",
    "        }\n",
    "        weights = {\n",
    "            feature_name: float(model_weight)\n",
    "            for feature_name, model_weight in zip(\n",
    "                weights[\"feature_names\"], weights[\"linear_model_weights\"]\n",
    "            )\n",
    "        }\n",
    "        print({k: round(weights[k], 2) for k in weights})\n",
    "        print(evaluation)\n",
    "        return partial_result\n",
    "\n",
    "    def forward_selection_model_search(\n",
    "        self,\n",
    "        train_data,\n",
    "        dev_data,\n",
    "        feature_names,\n",
    "        maximum_number_of_features=None,\n",
    "        output_file=\"forward_selection_model_search.json\",\n",
    "        protected_features=None,\n",
    "        hyperparameter=None,\n",
    "    ):\n",
    "\n",
    "        output_file = os.path.join(self.folder_dir, output_file)\n",
    "        try:\n",
    "            with open(output_file, \"r\") as f:\n",
    "                results = json.load(f)\n",
    "                print(\n",
    "                    \"Forward selection model search: Results from output file loaded.\"\n",
    "                )\n",
    "        except FileNotFoundError:\n",
    "            print(\n",
    "                \"Forward selection model search: File not found. Starting search from scratch.\"\n",
    "            )\n",
    "            results = []\n",
    "\n",
    "        if not maximum_number_of_features:\n",
    "            maximum_number_of_features = len(feature_names)\n",
    "        maximum_number_of_features = min(maximum_number_of_features, len(feature_names))\n",
    "\n",
    "        if not protected_features:\n",
    "            protected_features = []\n",
    "        else:\n",
    "            partial_result = self._forward_selection_iteration(\n",
    "                train_data=train_data,\n",
    "                dev_data=dev_data,\n",
    "                feature_names=protected_features,\n",
    "                hyperparameter=hyperparameter,\n",
    "            )\n",
    "            results.append(partial_result)\n",
    "        while len(protected_features) < maximum_number_of_features:\n",
    "            best_eval = 0\n",
    "            best_features = None\n",
    "            feature_names = [x for x in feature_names if x not in protected_features]\n",
    "            for new_feature in feature_names:\n",
    "                experimental_features = protected_features + [new_feature]\n",
    "                partial_result = self._forward_selection_iteration(\n",
    "                    train_data=train_data,\n",
    "                    dev_data=dev_data,\n",
    "                    feature_names=experimental_features,\n",
    "                    hyperparameter=hyperparameter,\n",
    "                )\n",
    "                evaluation = partial_result[\"evaluation\"]\n",
    "                results.append(partial_result)\n",
    "                if evaluation > best_eval:\n",
    "                    best_eval = evaluation\n",
    "                    best_features = experimental_features\n",
    "                with open(output_file, \"w\") as f:\n",
    "                    json.dump(results, f)\n",
    "            protected_features = best_features\n",
    "        return results\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0d77781-8f0d-46d0-92f6-22ef8b25719d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#|hide\n",
    "nbdev_export()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "learntorank",
   "language": "python",
   "name": "learntorank"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}