partitioned-models/m5-forecasting-lightgbm-with-timeseries-splits.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Effortless models deployment with MLflow\r\n",
        "\r\n",
        "## Working with partitioned models (many models) in MLflow\r\n",
        "\r\n",
        "This example demonstrates how to package multiple models with MLflow to work all together to produce inference over partitioned data. In this scenario, the data can be partitioned based on a given attribute and then individual models are trained for each slice of the data. This is a typical approach in time series data when training individual models for each partition outperforms training a larger model over the entire dataset. "
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Working with the M5 forecasting dataset\r\n",
        "\r\n",
        "To demonstrate this scenario, we are going to solve the [M5 Forecasting competition](https://www.kaggle.com/competitions/m5-forecasting-accuracy/overview) problem which tries to predict the demand of diffent products on different retail stores across the US. Instead of training 1 model over the entire dataset, we are going to partition the data by store and then train one model per each store. In this dataset, there are 10 different stores so this will produce 10 different models.\r\n",
        "\r\n",
        "At the end, we will package all those 10 models into a single entity that can be deploy using MLflow."
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import mlflow\r\n",
        "import pandas as pd\r\n",
        "import numpy as np\r\n",
        "import gc\r\n",
        "from typing import List, Dict, Any, Union, Tuple"
      ],
      "outputs": [],
      "execution_count": 1,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1688746688284
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "We start by reading the input data from the 3 CSV files:"
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "sales = pd.read_csv('data/sales_train_validation.csv')\r\n",
        "calendar = pd.read_csv('data/calendar.csv')\r\n",
        "prices = pd.read_csv('data/sell_prices.csv')"
      ],
      "outputs": [],
      "execution_count": 2,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1688746697638
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Combining multiple datasets into one"
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's define a function that can take the three datasets and build a unified one:"
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def build_dataset(sales, calendar, prices, remove_old: bool = True):\r\n",
        "    data = sales\r\n",
        "    data.drop(columns=['id'], inplace=True)\r\n",
        "    calendar.drop(columns=['weekday', 'year'], inplace=True)\r\n",
        "\r\n",
        "    data = pd.melt(data, id_vars = ['item_id', 'dept_id', 'cat_id', 'store_id', 'state_id'], var_name = 'day', value_name = 'demand')\r\n",
        "\r\n",
        "    data = pd.merge(data, calendar, how='left', left_on=['day'], right_on=['d'])\r\n",
        "    data.drop(columns=['day', 'd'], inplace=True)\r\n",
        "\r\n",
        "    data = data.merge(prices, on = ['store_id', 'item_id', 'wm_yr_wk'], how = 'left')\r\n",
        "    data.drop(columns=['wm_yr_wk'], inplace=True)\r\n",
        "\r\n",
        "    if remove_old:\r\n",
        "        del sales\r\n",
        "        del calendar\r\n",
        "        del prices\r\n",
        "        gc.collect()\r\n",
        "\r\n",
        "    return data"
      ],
      "outputs": [],
      "execution_count": 3,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1688746697827
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "The function `build_dataset` will construct our final dataset. We then we delete the old datasets to save memory in the compute we are using as this data can grow large."
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "data = build_dataset(sales, calendar, prices)"
      ],
      "outputs": [],
      "execution_count": 4,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1688746798510
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Data would look as follows:"
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "data.sample(10)"
      ],
      "outputs": [
        {
          "output_type": "execute_result",
          "execution_count": 5,
          "data": {
            "text/plain": "                  item_id      dept_id     cat_id store_id state_id  demand  \\\n15034829    HOBBIES_1_217    HOBBIES_1    HOBBIES     CA_2       CA       0   \n28898415      FOODS_3_821      FOODS_3      FOODS     WI_1       WI       0   \n516347    HOUSEHOLD_1_511  HOUSEHOLD_1  HOUSEHOLD     WI_3       WI       0   \n58221494  HOUSEHOLD_1_281  HOUSEHOLD_1  HOUSEHOLD     TX_2       TX       0   \n35581043      FOODS_3_037      FOODS_3      FOODS     WI_3       WI       2   \n54824175    HOBBIES_1_111    HOBBIES_1    HOBBIES     CA_2       CA       0   \n8432284       FOODS_1_191      FOODS_1      FOODS     TX_2       TX       1   \n48242199  HOUSEHOLD_1_364  HOUSEHOLD_1  HOUSEHOLD     CA_3       CA       0   \n53302449      FOODS_3_656      FOODS_3      FOODS     CA_2       CA       0   \n27719728  HOUSEHOLD_2_173  HOUSEHOLD_2  HOUSEHOLD     CA_2       CA       1   \n\n                date  wday  month   event_name_1 event_type_1 event_name_2  \\\n15034829  2012-06-05     4      6            NaN          NaN          NaN   \n28898415  2013-09-02     3      9       LaborDay     National          NaN   \n516347    2011-02-14     3      2  ValentinesDay     Cultural          NaN   \n58221494  2016-04-21     6      4            NaN          NaN          NaN   \n35581043  2014-04-09     5      4            NaN          NaN          NaN   \n54824175  2016-01-01     7      1        NewYear     National          NaN   \n8432284   2011-11-01     4     11            NaN          NaN          NaN   \n48242199  2015-05-30     1      5            NaN          NaN          NaN   \n53302449  2015-11-12     6     11            NaN          NaN          NaN   \n27719728  2013-07-26     7      7            NaN          NaN          NaN   \n\n         event_type_2  snap_CA  snap_TX  snap_WI  sell_price  \n15034829          NaN        1        1        1        3.47  \n28898415          NaN        1        0        1        4.98  \n516347            NaN        0        0        1         NaN  \n58221494          NaN        0        0        0        4.94  \n35581043          NaN        1        1        1        2.50  \n54824175          NaN        1        1        0        2.92  \n8432284           NaN        1        1        0        5.72  \n48242199          NaN        0        0        0        7.48  \n53302449          NaN        0        1        1        3.48  \n27719728          NaN        0        0        0        6.97  ",
            "text/html": "<div>\n<style scoped>\n    .dataframe tbody tr th:only-of-type {\n        vertical-align: middle;\n    }\n\n    .dataframe tbody tr th {\n        vertical-align: top;\n    }\n\n    .dataframe thead th {\n        text-align: right;\n    }\n</style>\n<table border=\"1\" class=\"dataframe\">\n  <thead>\n    <tr style=\"text-align: right;\">\n      <th></th>\n      <th>item_id</th>\n      <th>dept_id</th>\n      <th>cat_id</th>\n      <th>store_id</th>\n      <th>state_id</th>\n      <th>demand</th>\n      <th>date</th>\n      <th>wday</th>\n      <th>month</th>\n      <th>event_name_1</th>\n      <th>event_type_1</th>\n      <th>event_name_2</th>\n      <th>event_type_2</th>\n      <th>snap_CA</th>\n      <th>snap_TX</th>\n      <th>snap_WI</th>\n      <th>sell_price</th>\n    </tr>\n  </thead>\n  <tbody>\n    <tr>\n      <th>15034829</th>\n      <td>HOBBIES_1_217</td>\n      <td>HOBBIES_1</td>\n      <td>HOBBIES</td>\n      <td>CA_2</td>\n      <td>CA</td>\n      <td>0</td>\n      <td>2012-06-05</td>\n      <td>4</td>\n      <td>6</td>\n      <td>NaN</td>\n      <td>NaN</td>\n      <td>NaN</td>\n      <td>NaN</td>\n      <td>1</td>\n      <td>1</td>\n      <td>1</td>\n      <td>3.47</td>\n    </tr>\n    <tr>\n      <th>28898415</th>\n      <td>FOODS_3_821</td>\n      <td>FOODS_3</td>\n      <td>FOODS</td>\n      <td>WI_1</td>\n      <td>WI</td>\n      <td>0</td>\n      <td>2013-09-02</td>\n      <td>3</td>\n      <td>9</td>\n      <td>LaborDay</td>\n      <td>National</td>\n      <td>NaN</td>\n      <td>NaN</td>\n      <td>1</td>\n      <td>0</td>\n      <td>1</td>\n      <td>4.98</td>\n    </tr>\n    <tr>\n      <th>516347</th>\n      <td>HOUSEHOLD_1_511</td>\n      <td>HOUSEHOLD_1</td>\n      <td>HOUSEHOLD</td>\n      <td>WI_3</td>\n      <td>WI</td>\n      <td>0</td>\n      <td>2011-02-14</td>\n      <td>3</td>\n      <td>2</td>\n      <td>ValentinesDay</td>\n      <td>Cultural</td>\n      <td>NaN</td>\n      <td>NaN</td>\n      <td>0</td>\n      <td>0</td>\n      <td>1</td>\n      <td>NaN</td>\n    </tr>\n    <tr>\n      <th>58221494</th>\n      <td>HOUSEHOLD_1_281</td>\n      <td>HOUSEHOLD_1</td>\n      <td>HOUSEHOLD</td>\n      <td>TX_2</td>\n      <td>TX</td>\n      <td>0</td>\n      <td>2016-04-21</td>\n      <td>6</td>\n      <td>4</td>\n      <td>NaN</td>\n      <td>NaN</td>\n      <td>NaN</td>\n      <td>NaN</td>\n      <td>0</td>\n      <td>0</td>\n      <td>0</td>\n      <td>4.94</td>\n    </tr>\n    <tr>\n      <th>35581043</th>\n      <td>FOODS_3_037</td>\n      <td>FOODS_3</td>\n      <td>FOODS</td>\n      <td>WI_3</td>\n      <td>WI</td>\n      <td>2</td>\n      <td>2014-04-09</td>\n      <td>5</td>\n      <td>4</td>\n      <td>NaN</td>\n      <td>NaN</td>\n      <td>NaN</td>\n      <td>NaN</td>\n      <td>1</td>\n      <td>1</td>\n      <td>1</td>\n      <td>2.50</td>\n    </tr>\n    <tr>\n      <th>54824175</th>\n      <td>HOBBIES_1_111</td>\n      <td>HOBBIES_1</td>\n      <td>HOBBIES</td>\n      <td>CA_2</td>\n      <td>CA</td>\n      <td>0</td>\n      <td>2016-01-01</td>\n      <td>7</td>\n      <td>1</td>\n      <td>NewYear</td>\n      <td>National</td>\n      <td>NaN</td>\n      <td>NaN</td>\n      <td>1</td>\n      <td>1</td>\n      <td>0</td>\n      <td>2.92</td>\n    </tr>\n    <tr>\n      <th>8432284</th>\n      <td>FOODS_1_191</td>\n      <td>FOODS_1</td>\n      <td>FOODS</td>\n      <td>TX_2</td>\n      <td>TX</td>\n      <td>1</td>\n      <td>2011-11-01</td>\n      <td>4</td>\n      <td>11</td>\n      <td>NaN</td>\n      <td>NaN</td>\n      <td>NaN</td>\n      <td>NaN</td>\n      <td>1</td>\n      <td>1</td>\n      <td>0</td>\n      <td>5.72</td>\n    </tr>\n    <tr>\n      <th>48242199</th>\n      <td>HOUSEHOLD_1_364</td>\n      <td>HOUSEHOLD_1</td>\n      <td>HOUSEHOLD</td>\n      <td>CA_3</td>\n      <td>CA</td>\n      <td>0</td>\n      <td>2015-05-30</td>\n      <td>1</td>\n      <td>5</td>\n      <td>NaN</td>\n      <td>NaN</td>\n      <td>NaN</td>\n      <td>NaN</td>\n      <td>0</td>\n      <td>0</td>\n      <td>0</td>\n      <td>7.48</td>\n    </tr>\n    <tr>\n      <th>53302449</th>\n      <td>FOODS_3_656</td>\n      <td>FOODS_3</td>\n      <td>FOODS</td>\n      <td>CA_2</td>\n      <td>CA</td>\n      <td>0</td>\n      <td>2015-11-12</td>\n      <td>6</td>\n      <td>11</td>\n      <td>NaN</td>\n      <td>NaN</td>\n      <td>NaN</td>\n      <td>NaN</td>\n      <td>0</td>\n      <td>1</td>\n      <td>1</td>\n      <td>3.48</td>\n    </tr>\n    <tr>\n      <th>27719728</th>\n      <td>HOUSEHOLD_2_173</td>\n      <td>HOUSEHOLD_2</td>\n      <td>HOUSEHOLD</td>\n      <td>CA_2</td>\n      <td>CA</td>\n      <td>1</td>\n      <td>2013-07-26</td>\n      <td>7</td>\n      <td>7</td>\n      <td>NaN</td>\n      <td>NaN</td>\n      <td>NaN</td>\n      <td>NaN</td>\n      <td>0</td>\n      <td>0</td>\n      <td>0</td>\n      <td>6.97</td>\n    </tr>\n  </tbody>\n</table>\n</div>"
          },
          "metadata": {}
        }
      ],
      "execution_count": 5,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1688746822725
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Train and validation split"
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Next, we will separate our data in train and validation. Notice that this function doesn't return the splits of the data but the row indices selected for each data set. Since the data is big, you will notice across all this notebook that we will index the data instead of making copies of it."
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def split_train_test_idxs(data, cutoff_date):\r\n",
        "    data.sort_values('date', inplace=True)\r\n",
        "\r\n",
        "    train_idxs = data['date'] <= cutoff_date\r\n",
        "    valid_idxs = data['date'] > cutoff_date\r\n",
        "\r\n",
        "    data.drop(columns=['date'], inplace=True)\r\n",
        "\r\n",
        "    gc.collect()\r\n",
        "\r\n",
        "    return train_idxs, valid_idxs"
      ],
      "outputs": [],
      "execution_count": 5,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1688600693753
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's apply the function to generate the splits. We are considering anything before 2014-04-24 as part of the training data:"
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "train_idxs, valid_idxs = split_train_test_idxs(data, cutoff_date='2014-04-24')"
      ],
      "outputs": [],
      "execution_count": 6,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1688600760982
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Feature engineering\r\n",
        "\r\n",
        "We are going to separate our target variable from the features:"
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "features, target = data.loc[:, data.columns != \"demand\"], data[[\"demand\"]]\r\n",
        "\r\n",
        "del data\r\n",
        "gc.collect()"
      ],
      "outputs": [
        {
          "output_type": "execute_result",
          "execution_count": 7,
          "data": {
            "text/plain": "0"
          },
          "metadata": {}
        }
      ],
      "execution_count": 7,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1688600784473
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now, let's build now a function that can pre-process the input features. To keep it simple, we are only going to focus on doing categorical encoding of the variables and handling missing values. The function will also return the transformation that can be use to transform data later. We will put this transformation inside of our final model:"
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from sklearn.preprocessing import OrdinalEncoder\r\n",
        "from sklearn.compose import ColumnTransformer\r\n",
        "from tqdm import tqdm\r\n",
        "\r\n",
        "def preprocess(data: pd.DataFrame, categorical_features: List[str]) -> Tuple[pd.DataFrame, ColumnTransformer]:\r\n",
        "    \"\"\"\r\n",
        "    Preprocess the input data features.\r\n",
        "\r\n",
        "    Parameters\r\n",
        "    ----------\r\n",
        "    data: pd.DataFrame\r\n",
        "        The input data features\r\n",
        "    categorical_features: List[str]\r\n",
        "        The features that you want to treat as categorical\r\n",
        "\r\n",
        "    Returns\r\n",
        "    -------\r\n",
        "    Tuple[pd.DataFrame, ColumnTransformer]:\r\n",
        "        The transformed data and the associated transformations.\r\n",
        "    \"\"\"\r\n",
        "    transformations = ColumnTransformer(\r\n",
        "            [\r\n",
        "                ('encoding', OrdinalEncoder(handle_unknown='use_encoded_value', unknown_value=np.nan, encoded_missing_value=-1), categorical_features),\r\n",
        "            ],\r\n",
        "            remainder='passthrough'\r\n",
        "        )\r\n",
        "\r\n",
        "    transformed = transformations.fit_transform(data)\r\n",
        "    columns = [name.split('__')[1] for name in transformations.get_feature_names_out()]\r\n",
        "\r\n",
        "    return pd.DataFrame(data=transformed, columns=columns, index=data.index).infer_objects(), transformations"
      ],
      "outputs": [],
      "execution_count": 8,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1688600785068
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's run it over the features:"
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "categorical_features = ['item_id', 'dept_id', 'cat_id', 'store_id', 'state_id', 'event_name_1', 'event_type_1', 'event_name_2', 'event_type_2']\r\n",
        "features, transforms = preprocess(features, categorical_features)"
      ],
      "outputs": [],
      "execution_count": 9,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1688603032605
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Preparating the training framework"
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now it's time to generate a training function. The following training function trans a LightGBM model for the given datasets. It receives the input data (both containing training and validation), and the train-tests indices within this dataset. That means that indices in `train_idxs` will be used for training while `valid_idxs` will be used for validation. The argument `params` will be used to pass the hyperparameters of the model."
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from sklearn import metrics\r\n",
        "from typing import List, Dict, Any\r\n",
        "from mlflow.models.signature import infer_signature\r\n",
        "import lightgbm as lgb\r\n",
        "\r\n",
        "def train_and_evaluate(features: pd.DataFrame, target: pd.DataFrame, train_idxs: pd.Index, valid_idxs: pd.Index,\r\n",
        "                       params: Dict[str, Any], model_name: str = None) -> Tuple[lgb.Booster, float]:\r\n",
        "    \"\"\"\r\n",
        "    Trains a model using a LightGBM.\r\n",
        "\r\n",
        "    Parameters\r\n",
        "    ----------\r\n",
        "    features: pd.DataFrame\r\n",
        "        The input features, already preprocessed.\r\n",
        "    target: pd.DataFrame\r\n",
        "        The column to predict.\r\n",
        "    train_idxs: pd.Index\r\n",
        "        The indices of the data correspoding to training.\r\n",
        "    valid_idxs: pd.Index\r\n",
        "        The indices of the data corresponding to validation.\r\n",
        "    params: Dict[str, Any]:\r\n",
        "        The training parameters for the model.\r\n",
        "    metrics_prefix: Union[str, None]\r\n",
        "        A prefix to be used for logging metrics.\r\n",
        "    \"\"\"\r\n",
        "    X_train, X_valid = features[train_idxs], features[valid_idxs]\r\n",
        "    y_train, y_valid = target[train_idxs], target[valid_idxs]\r\n",
        "\r\n",
        "    dtrain = lgb.Dataset(X_train, label=y_train)\r\n",
        "    dvalid = lgb.Dataset(X_valid, label=y_valid)\r\n",
        "\r\n",
        "    clf = lgb.train(params, dtrain, 2500, valid_sets = [dtrain, dvalid], callbacks=[lgb.early_stopping(stopping_rounds=10)])\r\n",
        "\r\n",
        "    y_pred_valid = clf.predict(X_valid, num_iteration=clf.best_iteration)\r\n",
        "\r\n",
        "    val_error = np.sqrt(metrics.mean_squared_error(y_valid, y_pred_valid))\r\n",
        "    val_score = metrics.r2_score(y_valid, y_pred_valid)\r\n",
        "\r\n",
        "    mlflow.log_metric(\"valid_rmse\", val_error)\r\n",
        "    mlflow.log_metric(\"valid_r2\", val_score)\r\n",
        "\r\n",
        "    if model_name:\r\n",
        "        mlflow.lightgbm.log_model(clf, model_name, signature=infer_signature(features, target))\r\n",
        "    \r\n",
        "    return (clf, val_error)"
      ],
      "outputs": [],
      "execution_count": 10,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1688603033858
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "For sake of simplicity, the hyper-parameters we will use will be the same:"
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "params = {\r\n",
        "    'num_leaves': 500,\r\n",
        "    'min_child_weight': 0.034,\r\n",
        "    'feature_fraction': 0.379,\r\n",
        "    'bagging_fraction': 0.418,\r\n",
        "    'min_data_in_leaf': 106,\r\n",
        "    'objective': 'regression',\r\n",
        "    'max_depth': -1,\r\n",
        "    'learning_rate': 0.005,\r\n",
        "    \"boosting_type\": \"gbdt\",\r\n",
        "    \"bagging_seed\": 42,\r\n",
        "    \"metric\": 'rmse',\r\n",
        "    \"verbosity\": -1,\r\n",
        "    'reg_alpha': 0.3899,\r\n",
        "    'reg_lambda': 0.648,\r\n",
        "    'random_state': 222,\r\n",
        "}"
      ],
      "outputs": [],
      "execution_count": 11,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1688603034582
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "We are going to need another helper function, which will return the `store_id` associated with a given \"categorical encoding\" of this column. Since our feature dataset is preprocess, the store \"CA_1\" is encoded as a numeric value. The function then will return for an input `0` the associated label, for instance `CA_1`."
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def get_encoded_value(transformations: ColumnTransformer, column: str, value: float) -> str:\r\n",
        "    \"\"\"\r\n",
        "    Returns the label associated with a given encoded value for a given column.\r\n",
        "\r\n",
        "    Parameters\r\n",
        "    ----------\r\n",
        "    transformations: ColumnTransformer\r\n",
        "        The transformations used to encode the columns.\r\n",
        "    column: str\r\n",
        "        The name of the column\r\n",
        "    value: float\r\n",
        "        The encoded value of the column indicated.\r\n",
        "    \"\"\"\r\n",
        "    column_encoder_idxs = [i for i in range(len(transformations.transformers_)) if column in transformations.transformers_[i][2]]\r\n",
        "    if column_encoder_idxs:\r\n",
        "        column_encoder_idx = column_encoder_idxs[0]\r\n",
        "        column_index = transformations.transformers_[column_encoder_idx][2].index(column)\r\n",
        "        return transformations.transformers_[column_encoder_idx][1].categories_[column_index][int(value)]\r\n",
        "    else:\r\n",
        "        raise ValueError(f'There is no transformation applied to column ``{column}``')"
      ],
      "outputs": [],
      "execution_count": 12,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": true
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1688603035294
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now, let's try to design our model. We want to create a partitioned multi-model, meaning, the data will be partitioned by a given key (in this case \"store_id\") and we will train one model per each of those partitions. When running inference, we want to apply the correct model based on the partition the data belongs to.\r\n",
        "\r\n",
        "To solve this problem we will create a custom model in MLflow. The `PartitionedModelEnsemble` will be a meta model that will be able to apply multiple models to a given input data set based on a partition key. Those models can be trained independently or all together, like in this example.\r\n",
        "\r\n",
        "> Tip: Notice how models are loaded using `mlflow.pyfunc.load_model`."
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from mlflow.pyfunc import PythonModel, PythonModelContext\r\n",
        "from sklearn.compose import ColumnTransformer\r\n",
        "import pandas as pd\r\n",
        "import mlflow\r\n",
        "\r\n",
        "class PartitionedModelEnsemble(PythonModel):\r\n",
        "    \"\"\"\r\n",
        "    A custom model that implements a paritioned inferrencing strategy over an ensamble of models.\r\n",
        "    \"\"\"\r\n",
        "    def __init__(self, key: str, prediction_col: str, transformations: ColumnTransformer):\r\n",
        "        \"\"\"\r\n",
        "        Creates a new instance of the mode.\r\n",
        "\r\n",
        "        Parameters\r\n",
        "        ----------\r\n",
        "        key: str\r\n",
        "            The name of the column the data will be partitioned on.\r\n",
        "        prediction_col: str\r\n",
        "            The name of the column the model generates predictions on.\r\n",
        "        transformations: ColumnTransformer\r\n",
        "            Any given transformation that needs to be applied to the data before sending to the model.\r\n",
        "        \"\"\"\r\n",
        "        self.pred_col = prediction_col\r\n",
        "        self.key = key\r\n",
        "        self.transformations = transformations\r\n",
        "\r\n",
        "    def load_context(self, context: PythonModelContext):\r\n",
        "        \"\"\"\r\n",
        "        Loads all the models for the given partitions. This method assumes the models were logged with the\r\n",
        "        different of the column `key` as artifact key.\r\n",
        "        \"\"\"\r\n",
        "        self.models = { key: mlflow.pyfunc.load_model(model_path) for key, model_path in context.artifacts.items() }\r\n",
        "        \r\n",
        "    def predict(self, context: PythonModelContext, data: pd.DataFrame) -> pd.DataFrame:\r\n",
        "        if isinstance(data, pd.DataFrame):\r\n",
        "            predictions = pd.DataFrame(0, index=data.index, columns=[self.pred_col])\r\n",
        "            \r\n",
        "            # Get all the unique partition's value in the input data\r\n",
        "            key_ids = data[self.key].unique()\r\n",
        "\r\n",
        "            # We will run 1 predict call per each partition\r\n",
        "            for key_id in key_ids:\r\n",
        "                input_data_idx = data[self.key] == key_id\r\n",
        "\r\n",
        "                if self.transformations:\r\n",
        "                    columns = [name.split('__')[1] for name in self.transformations.get_feature_names_out()]\r\n",
        "                    transformed_data = pd.DataFrame(self.transformations.transform(data[input_data_idx]), columns=columns)\r\n",
        "                else:\r\n",
        "                    transformed_data = data[input_data_idx]\r\n",
        "                predictions[input_data_idx] = self.models[key_id].predict(transformed_data.loc[:, transformed_data.columns != self.key]).reshape(-1,1)\r\n",
        "\r\n",
        "            return predictions\r\n",
        "        \r\n",
        "        raise TypeError(\"This implementation can only take pandas DataFrame as inputs\")"
      ],
      "outputs": [],
      "execution_count": 15,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1688748127260
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Model ensemble training"
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "It's time now to train our model. Let's start by configuring the run with MLflow:"
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "mlflow.set_experiment(experiment_name=\"m5-forecasting-mlflow\")"
      ],
      "outputs": [
        {
          "output_type": "execute_result",
          "execution_count": 14,
          "data": {
            "text/plain": "<Experiment: artifact_location='', creation_time=1687962235375, experiment_id='a0ca7d8e-43e1-4d8f-aa9e-7eb636b58fc7', last_update_time=None, lifecycle_stage='active', name='m5-forecasting-mlflow', tags={}>"
          },
          "metadata": {}
        }
      ],
      "execution_count": 14,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1688603037578
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "We will use the following training routine which will train all the models, one per each partition. It does the following:\r\n",
        "\r\n",
        "1. It gets all the different values for the stores in the dataset. There are 10 different stores.\r\n",
        "2. For each of them:\r\n",
        "\r\n",
        "    a. We start a new child run in MLflow. We will have one child run per each store. This makes comparation and evaluation much easier. \r\n",
        "    b. Selects the indices of the data that belong to this particular store. This is done for training and validation datasets.\r\n",
        "    c. Calls the `train_and_evaluate()` function passing as arguments the features (all of them but the partition key column), the training indices and the test indices for the given store.\r\n",
        "    d. The function returns the trained model and the associated performance of it. Autolog is taking care of all the logging so we don't do much here.\r\n",
        "    e. In a dictionary called `artifacts`, we will record which is the artifact associated with the given store. The artifact is indicated in the form `runs:/<RUN_ID>/<ARTIFACT_PATH`. The reason for that is that this makes very clear from where the associated model came from, which we can use for lineage purposes later."
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from tqdm import tqdm\r\n",
        "\r\n",
        "with mlflow.start_run() as run:\r\n",
        "    models = {}\r\n",
        "    model_error = {}\r\n",
        "    artifacts = {}\r\n",
        "\r\n",
        "    partition_key = 'store_id'\r\n",
        "    model_features = features.columns.difference([partition_key]).values\r\n",
        "    stores = features[partition_key].unique()\r\n",
        "\r\n",
        "    mlflow.lightgbm.autolog(log_datasets=False, log_models=False)\r\n",
        "\r\n",
        "    for store in tqdm(stores):\r\n",
        "        store_id = get_encoded_value(transforms, partition_key, store)\r\n",
        "        with mlflow.start_run(run_name=f\"train_store_{store_id}\", nested=True) as store_run:\r\n",
        "            store_train_idx = (train_idxs) & (features[partition_key] == store)\r\n",
        "            store_valid_idx = (valid_idxs) & (features[partition_key] == store)\r\n",
        "            models[store_id], model_error[store_id] = train_and_evaluate(features=features.loc[:, features.columns != partition_key],\r\n",
        "                                                                         target=target, \r\n",
        "                                                                         train_idxs=store_train_idx, \r\n",
        "                                                                         valid_idxs=store_valid_idx,\r\n",
        "                                                                         params=params,\r\n",
        "                                                                         model_name=store_id)\r\n",
        "\r\n",
        "            artifacts[store_id] = f\"runs:/{store_run.info.run_id}/{store_id}\"\r\n",
        "\r\n",
        "    mlflow.pyfunc.log_model(\"model\", \r\n",
        "                            python_model=PartitionedModelEnsemble(partition_key, target, transforms),\r\n",
        "                            artifacts=artifacts,\r\n",
        "                            registered_model_name=\"m5-forcasting-partitioned\")"
      ],
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": "  0%|          | 0/10 [00:00<?, ?it/s]2023/07/05 23:51:34 WARNING mlflow.utils.autologging_utils: MLflow autologging encountered a warning: \"/anaconda/envs/azureml_py310_sdkv2/lib/python3.10/site-packages/_distutils_hack/__init__.py:33: UserWarning: Setuptools is replacing distutils.\"\n100%|██████████| 10/10 [56:54<00:00, 341.50s/it]\n2023/07/06 00:45:42 INFO mlflow.types.utils: Unsupported type hint: <class 'pandas.core.frame.DataFrame'>, skipping schema inference\n2023/07/06 00:45:42 INFO mlflow.types.utils: Unsupported type hint: <class 'pandas.core.frame.DataFrame'>, skipping schema inference\nRegistered model 'm5-forcasting-partitioned' already exists. Creating a new version of this model...\n2023/07/06 00:46:43 INFO mlflow.tracking._model_registry.client: Waiting up to 300 seconds for model version to finish creation. Model name: m5-forcasting-partitioned, version 2\nCreated version '2' of model 'm5-forcasting-partitioned'.\n"
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": "Training until validation scores don't improve for 10 rounds\nEarly stopping, best iteration is:\n[762]\ttraining's rmse: 3.34108\tvalid_1's rmse: 3.28104\nTraining until validation scores don't improve for 10 rounds\nEarly stopping, best iteration is:\n[1063]\ttraining's rmse: 2.72586\tvalid_1's rmse: 3.22566\nTraining until validation scores don't improve for 10 rounds\nEarly stopping, best iteration is:\n[1119]\ttraining's rmse: 1.76246\tvalid_1's rmse: 2.21233\nTraining until validation scores don't improve for 10 rounds\nEarly stopping, best iteration is:\n[762]\ttraining's rmse: 3.49055\tvalid_1's rmse: 3.46766\nTraining until validation scores don't improve for 10 rounds\nEarly stopping, best iteration is:\n[415]\ttraining's rmse: 3.63284\tvalid_1's rmse: 2.89144\nTraining until validation scores don't improve for 10 rounds\nEarly stopping, best iteration is:\n[1245]\ttraining's rmse: 2.80333\tvalid_1's rmse: 3.46991\nTraining until validation scores don't improve for 10 rounds\nEarly stopping, best iteration is:\n[661]\ttraining's rmse: 2.69411\tvalid_1's rmse: 2.66093\nTraining until validation scores don't improve for 10 rounds\nEarly stopping, best iteration is:\n[550]\ttraining's rmse: 2.26246\tvalid_1's rmse: 2.50658\nTraining until validation scores don't improve for 10 rounds\nEarly stopping, best iteration is:\n[762]\ttraining's rmse: 5.06703\tvalid_1's rmse: 4.67845\nTraining until validation scores don't improve for 10 rounds\nEarly stopping, best iteration is:\n[676]\ttraining's rmse: 1.63436\tvalid_1's rmse: 1.73468\n"
        }
      ],
      "execution_count": 15,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1688604403625
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Notice that at the end of the routine, we are logging a new model in a very particular way. We are logging an instance of our custom model, `PartitionedModelEnsemble`. Pay closer attention about how artifacts is indicated. Artifacts are used to indicate all the models this aggregator model ensembles together. When artifacts are indicated in the form of `runs:/`, MLflow will automatically pull those artifacts from the MLflow server and log them inside of the model as a single unit. So you still have a single unit and reproducible asset that you can move along. However, the advantage of doing it this way is that you can easily see from where these models came from."
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Running the model\r\n",
        "\r\n",
        "Let's see now how this model works in practice when running inference. Let's start by pulling the model from the registry:"
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "model = mlflow.pyfunc.load_model(f\"models:/m5-forcasting-partitioned/latest\")"
      ],
      "outputs": [],
      "execution_count": 6,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1688746855822
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's build a sample dataset we can use to test this out:"
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "sales = pd.read_csv('data/sales_train_validation.csv')\r\n",
        "calendar = pd.read_csv('data/calendar.csv')\r\n",
        "prices = pd.read_csv('data/sell_prices.csv')\r\n",
        "data = build_dataset(sales, calendar, prices)"
      ],
      "outputs": [],
      "execution_count": 4,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1688600480806
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's remove the target column:"
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "data.drop(columns=[\"demand\"], inplace=True)"
      ],
      "outputs": [],
      "execution_count": 12,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1688748032993
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Let's create the input data. We will create a copy of it:"
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "input_data = data.sample(n=1000).copy().reset_index()"
      ],
      "outputs": [],
      "execution_count": 13,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1688748035892
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "model.predict(input_data)"
      ],
      "outputs": [
        {
          "output_type": "execute_result",
          "execution_count": 18,
          "data": {
            "text/plain": "       demand\n0   -0.146778\n1    0.381161\n2    0.071890\n3   -0.024440\n4    0.070622\n..        ...\n995  0.186330\n996  0.059465\n997  0.116552\n998  1.126394\n999  0.103094\n\n[1000 rows x 1 columns]",
            "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 tr th {\n        text-align: left;\n    }\n</style>\n<table border=\"1\" class=\"dataframe\">\n  <thead>\n    <tr>\n      <th></th>\n      <th>demand</th>\n    </tr>\n  </thead>\n  <tbody>\n    <tr>\n      <th>0</th>\n      <td>-0.146778</td>\n    </tr>\n    <tr>\n      <th>1</th>\n      <td>0.381161</td>\n    </tr>\n    <tr>\n      <th>2</th>\n      <td>0.071890</td>\n    </tr>\n    <tr>\n      <th>3</th>\n      <td>-0.024440</td>\n    </tr>\n    <tr>\n      <th>4</th>\n      <td>0.070622</td>\n    </tr>\n    <tr>\n      <th>...</th>\n      <td>...</td>\n    </tr>\n    <tr>\n      <th>995</th>\n      <td>0.186330</td>\n    </tr>\n    <tr>\n      <th>996</th>\n      <td>0.059465</td>\n    </tr>\n    <tr>\n      <th>997</th>\n      <td>0.116552</td>\n    </tr>\n    <tr>\n      <th>998</th>\n      <td>1.126394</td>\n    </tr>\n    <tr>\n      <th>999</th>\n      <td>0.103094</td>\n    </tr>\n  </tbody>\n</table>\n<p>1000 rows × 1 columns</p>\n</div>"
          },
          "metadata": {}
        }
      ],
      "execution_count": 18,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1688748146882
        }
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "Notice how the signature of the \"downstream models\" is different from the one in the aggregator. Particularly, the downstream models 1) take transformed data as inputs and 2) the column `store_id` is not required."
      ],
      "metadata": {
        "nteract": {
          "transient": {
            "deleting": false
          }
        }
      }
    },
    {
      "cell_type": "code",
      "source": [
        "model._model_impl.python_model.models[\"CA_1\"].metadata.signature"
      ],
      "outputs": [
        {
          "output_type": "execute_result",
          "execution_count": 25,
          "data": {
            "text/plain": "inputs: \n  ['item_id': double, 'dept_id': double, 'cat_id': double, 'store_id': double, 'state_id': double, 'event_name_1': double, 'event_type_1': double, 'event_name_2': double, 'event_type_2': double, 'wday': double, 'month': double, 'snap_CA': double, 'snap_TX': double, 'snap_WI': double, 'sell_price': double]\noutputs: \n  [Tensor('float64', (-1,))]"
          },
          "metadata": {}
        }
      ],
      "execution_count": 25,
      "metadata": {
        "jupyter": {
          "source_hidden": false,
          "outputs_hidden": false
        },
        "nteract": {
          "transient": {
            "deleting": false
          }
        },
        "gather": {
          "logged": 1688447808463
        }
      }
    }
  ],
  "metadata": {
    "kernelspec": {
      "name": "python310-sdkv2",
      "language": "python",
      "display_name": "Python 3.10 - SDK v2"
    },
    "language_info": {
      "name": "python",
      "version": "3.10.9",
      "mimetype": "text/x-python",
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "pygments_lexer": "ipython3",
      "nbconvert_exporter": "python",
      "file_extension": ".py"
    },
    "microsoft": {
      "ms_spell_check": {
        "ms_spell_check_language": "en",
        "ms_ignore_dictionary": [
          "huggingface"
        ]
      },
      "host": {
        "AzureML": {
          "notebookHasBeenCompleted": true
        }
      }
    },
    "kernel_info": {
      "name": "python310-sdkv2"
    },
    "nteract": {
      "version": "nteract-front-end@1.0.0"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 4
}