examples/11_sasrec_dataframes_comparison.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Comparison of performance for different parts of pipeline depends on input data format (PySpark, Pandas, Polars)\n",
    "\n",
    "Note that this example designed only for comparison goals, to get more detailed information about pipeline steps see `09_sasrec_example.ipynb` and `10_bert4rec_example.ipynb`\n",
    "\n",
    "Example was executed on Intel(R) Xeon(R) Gold 6248R: 12 cores, HT:on, Turbo:on. OS: Ubuntu 20.04 LTS, total memory of 96 GB.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import lightning as L\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "from replay.metrics import OfflineMetrics, Recall, Precision, MAP, NDCG, HitRate, MRR\n",
    "from replay.splitters import LastNSplitter\n",
    "from replay.utils import get_spark_session, DataFrameLike\n",
    "from replay.data import (\n",
    "    FeatureHint,\n",
    "    FeatureInfo,\n",
    "    FeatureSchema,\n",
    "    FeatureSource,\n",
    "    FeatureType,\n",
    "    Dataset,\n",
    ")\n",
    "from replay.models.nn.optimizer_utils import FatOptimizerFactory\n",
    "from replay.models.nn.sequential.callbacks import (\n",
    "    ValidationMetricsCallback,\n",
    "    SparkPredictionCallback,\n",
    "    PandasPredictionCallback, \n",
    "    PolarsPredictionCallback,\n",
    ")\n",
    "from replay.models.nn.sequential.postprocessors import RemoveSeenItems\n",
    "from replay.data.nn import (\n",
    "    SequenceTokenizer,\n",
    "    SequentialDataset,\n",
    "    TensorFeatureSource,\n",
    "    TensorSchema,\n",
    "    TensorFeatureInfo\n",
    ")\n",
    "from replay.models.nn.sequential import SasRec\n",
    "from replay.models.nn.sequential.sasrec import (\n",
    "    SasRecPredictionDataset,\n",
    "    SasRecTrainingDataset,\n",
    "    SasRecValidationDataset,\n",
    ")\n",
    "\n",
    "import pandas as pd\n",
    "import polars as pl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "spark_session = get_spark_session()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Observe and load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install rs-datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<style type='text/css'>\n",
       ".datatable table.frame { margin-bottom: 0; }\n",
       ".datatable table.frame thead { border-bottom: none; }\n",
       ".datatable table.frame tr.coltypes td {  color: #FFFFFF;  line-height: 6px;  padding: 0 0.5em;}\n",
       ".datatable .bool    { background: #DDDD99; }\n",
       ".datatable .object  { background: #565656; }\n",
       ".datatable .int     { background: #5D9E5D; }\n",
       ".datatable .float   { background: #4040CC; }\n",
       ".datatable .str     { background: #CC4040; }\n",
       ".datatable .time    { background: #40CC40; }\n",
       ".datatable .row_index {  background: var(--jp-border-color3);  border-right: 1px solid var(--jp-border-color0);  color: var(--jp-ui-font-color3);  font-size: 9px;}\n",
       ".datatable .frame tbody td { text-align: left; }\n",
       ".datatable .frame tr.coltypes .row_index {  background: var(--jp-border-color0);}\n",
       ".datatable th:nth-child(2) { padding-left: 12px; }\n",
       ".datatable .hellipsis {  color: var(--jp-cell-editor-border-color);}\n",
       ".datatable .vellipsis {  background: var(--jp-layout-color0);  color: var(--jp-cell-editor-border-color);}\n",
       ".datatable .na {  color: var(--jp-cell-editor-border-color);  font-size: 80%;}\n",
       ".datatable .sp {  opacity: 0.25;}\n",
       ".datatable .footer { font-size: 9px; }\n",
       ".datatable .frame_dimensions {  background: var(--jp-border-color3);  border-top: 1px solid var(--jp-border-color0);  color: var(--jp-ui-font-color3);  display: inline-block;  opacity: 0.6;  padding: 1px 10px 1px 5px;}\n",
       "</style>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from rs_datasets import MovieLens\n",
    "\n",
    "movielens = MovieLens(\"20m\")\n",
    "interactions_pandas = movielens.ratings\n",
    "item_features_pandas = movielens.items"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>user_id</th>\n",
       "      <th>item_id</th>\n",
       "      <th>rating</th>\n",
       "      <th>timestamp</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>2</td>\n",
       "      <td>3.5</td>\n",
       "      <td>1112486027</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>29</td>\n",
       "      <td>3.5</td>\n",
       "      <td>1112484676</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1</td>\n",
       "      <td>32</td>\n",
       "      <td>3.5</td>\n",
       "      <td>1112484819</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1</td>\n",
       "      <td>47</td>\n",
       "      <td>3.5</td>\n",
       "      <td>1112484727</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1</td>\n",
       "      <td>50</td>\n",
       "      <td>3.5</td>\n",
       "      <td>1112484580</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   user_id  item_id  rating   timestamp\n",
       "0        1        2     3.5  1112486027\n",
       "1        1       29     3.5  1112484676\n",
       "2        1       32     3.5  1112484819\n",
       "3        1       47     3.5  1112484727\n",
       "4        1       50     3.5  1112484580"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "interactions_pandas.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>item_id</th>\n",
       "      <th>title</th>\n",
       "      <th>genres</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>Toy Story (1995)</td>\n",
       "      <td>Adventure|Animation|Children|Comedy|Fantasy</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2</td>\n",
       "      <td>Jumanji (1995)</td>\n",
       "      <td>Adventure|Children|Fantasy</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3</td>\n",
       "      <td>Grumpier Old Men (1995)</td>\n",
       "      <td>Comedy|Romance</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>4</td>\n",
       "      <td>Waiting to Exhale (1995)</td>\n",
       "      <td>Comedy|Drama|Romance</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>5</td>\n",
       "      <td>Father of the Bride Part II (1995)</td>\n",
       "      <td>Comedy</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   item_id                               title  \\\n",
       "0        1                    Toy Story (1995)   \n",
       "1        2                      Jumanji (1995)   \n",
       "2        3             Grumpier Old Men (1995)   \n",
       "3        4            Waiting to Exhale (1995)   \n",
       "4        5  Father of the Bride Part II (1995)   \n",
       "\n",
       "                                        genres  \n",
       "0  Adventure|Animation|Children|Comedy|Fantasy  \n",
       "1                   Adventure|Children|Fantasy  \n",
       "2                               Comedy|Romance  \n",
       "3                         Comedy|Drama|Romance  \n",
       "4                                       Comedy  "
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "item_features_pandas.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>user_id</th>\n",
       "      <th>item_id</th>\n",
       "      <th>rating</th>\n",
       "      <th>timestamp</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>4182421</th>\n",
       "      <td>28507</td>\n",
       "      <td>1176</td>\n",
       "      <td>4.0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>18950979</th>\n",
       "      <td>131160</td>\n",
       "      <td>1079</td>\n",
       "      <td>3.0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>18950936</th>\n",
       "      <td>131160</td>\n",
       "      <td>47</td>\n",
       "      <td>5.0</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>18950930</th>\n",
       "      <td>131160</td>\n",
       "      <td>21</td>\n",
       "      <td>3.0</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12341178</th>\n",
       "      <td>85252</td>\n",
       "      <td>45</td>\n",
       "      <td>3.0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7819902</th>\n",
       "      <td>53930</td>\n",
       "      <td>118706</td>\n",
       "      <td>3.5</td>\n",
       "      <td>1429</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2508834</th>\n",
       "      <td>16978</td>\n",
       "      <td>2093</td>\n",
       "      <td>3.5</td>\n",
       "      <td>707</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12898546</th>\n",
       "      <td>89081</td>\n",
       "      <td>55232</td>\n",
       "      <td>3.5</td>\n",
       "      <td>1607</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12898527</th>\n",
       "      <td>89081</td>\n",
       "      <td>52458</td>\n",
       "      <td>4.0</td>\n",
       "      <td>1608</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12675921</th>\n",
       "      <td>87586</td>\n",
       "      <td>7151</td>\n",
       "      <td>3.5</td>\n",
       "      <td>669</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>20000263 rows × 4 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "          user_id  item_id  rating  timestamp\n",
       "4182421     28507     1176     4.0          0\n",
       "18950979   131160     1079     3.0          0\n",
       "18950936   131160       47     5.0          1\n",
       "18950930   131160       21     3.0          2\n",
       "12341178    85252       45     3.0          0\n",
       "...           ...      ...     ...        ...\n",
       "7819902     53930   118706     3.5       1429\n",
       "2508834     16978     2093     3.5        707\n",
       "12898546    89081    55232     3.5       1607\n",
       "12898527    89081    52458     4.0       1608\n",
       "12675921    87586     7151     3.5        669\n",
       "\n",
       "[20000263 rows x 4 columns]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "interactions_pandas[\"timestamp\"] = interactions_pandas[\"timestamp\"].astype(\"int64\")\n",
    "interactions_pandas = interactions_pandas.sort_values(by=\"timestamp\")\n",
    "interactions_pandas[\"timestamp\"] = interactions_pandas.groupby(\"user_id\").cumcount()\n",
    "interactions_pandas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "interactions_spark = spark_session.createDataFrame(interactions_pandas)\n",
    "item_features_spark = spark_session.createDataFrame(item_features_pandas)\n",
    "\n",
    "interactions_polars = pl.from_pandas(interactions_pandas)\n",
    "item_features_polars = pl.from_pandas(item_features_pandas)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Splitting comparison"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "splitter = LastNSplitter(\n",
    "    N=1,\n",
    "    divide_column=\"user_id\",\n",
    "    query_column=\"user_id\",\n",
    "    strategy=\"interactions\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n 2 -r 3\n",
    "\n",
    "splitter.split(interactions_pandas)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "18.3 s ± 29.7 ms per loop (mean ± std. dev. of 3 runs, 2 loops each)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For a correct comparison, it is necessary to trigger the graph of calculations in the pyspark"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n 2 -r 3\n",
    "\n",
    "train, test = splitter.split(interactions_spark)\n",
    "train.count(), test.count()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "7.12 s ± 1.82 s per loop (mean ± std. dev. of 3 runs, 2 loops each)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n 2 -r 3\n",
    "\n",
    "splitter.split(interactions_polars)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "1.66 s ± 6.05 ms per loop (mean ± std. dev. of 3 runs, 2 loops each)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# pandas datasets\n",
    "raw_test_events_pandas, raw_test_gt_pandas = splitter.split(interactions_pandas)\n",
    "raw_validation_events_pandas, raw_validation_gt_pandas = splitter.split(raw_test_events_pandas)\n",
    "raw_train_events_pandas = raw_validation_events_pandas\n",
    "\n",
    "\n",
    "# pyspark datasets\n",
    "raw_test_events_spark, raw_test_gt_spark = splitter.split(interactions_spark)\n",
    "raw_validation_events_spark, raw_validation_gt_spark = splitter.split(raw_test_events_spark)\n",
    "raw_train_events_spark = raw_validation_events_spark\n",
    "\n",
    "# trigger for correct performance measurements\n",
    "raw_train_events_spark.cache(), raw_train_events_spark.count()\n",
    "raw_validation_events_spark.cache(), raw_validation_events_spark.count()\n",
    "raw_validation_gt_spark.cache(), raw_validation_gt_spark.count()\n",
    "raw_test_events_spark.cache(), raw_test_events_spark.count()\n",
    "raw_test_gt_spark.cache(), raw_test_gt_spark.count()\n",
    "\n",
    "# polars datasets\n",
    "raw_test_events_polars, raw_test_gt_polars = splitter.split(interactions_polars)\n",
    "raw_validation_events_polars, raw_validation_gt_polars = splitter.split(raw_test_events_polars)\n",
    "raw_train_events_polars = raw_validation_events_polars"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating Dataset objects"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def prepare_feature_schema(is_ground_truth: bool) -> FeatureSchema:\n",
    "    base_features = FeatureSchema(\n",
    "        [\n",
    "            FeatureInfo(\n",
    "                column=\"user_id\",\n",
    "                feature_hint=FeatureHint.QUERY_ID,\n",
    "                feature_type=FeatureType.CATEGORICAL,\n",
    "            ),\n",
    "            FeatureInfo(\n",
    "                column=\"item_id\",\n",
    "                feature_hint=FeatureHint.ITEM_ID,\n",
    "                feature_type=FeatureType.CATEGORICAL,\n",
    "            ),\n",
    "        ]\n",
    "    )\n",
    "    if is_ground_truth:\n",
    "        return base_features\n",
    "\n",
    "    all_features = base_features + FeatureSchema(\n",
    "        [\n",
    "            FeatureInfo(\n",
    "                column=\"timestamp\",\n",
    "                feature_type=FeatureType.NUMERICAL,\n",
    "                feature_hint=FeatureHint.TIMESTAMP,\n",
    "            ),\n",
    "        ]\n",
    "    )\n",
    "    return all_features\n",
    "\n",
    "\n",
    "def create_dataset_object(\n",
    "    schema: FeatureSchema,\n",
    "    interactions: DataFrameLike,\n",
    "    item_features: DataFrameLike = None,\n",
    "):\n",
    "    return Dataset(\n",
    "        feature_schema=schema,\n",
    "        interactions=interactions,\n",
    "        item_features=item_features,\n",
    "        check_consistency=True,\n",
    "        categorical_encoded=False,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "schema = prepare_feature_schema(is_ground_truth=False)\n",
    "gt_schema = prepare_feature_schema(is_ground_truth=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset_pandas = create_dataset_object(\n",
    "    schema, raw_train_events_pandas, item_features_pandas\n",
    ")\n",
    "validation_dataset_pandas = create_dataset_object(\n",
    "    schema, raw_validation_events_pandas, item_features_pandas\n",
    ")\n",
    "validation_gt_pandas = create_dataset_object(\n",
    "    gt_schema, raw_validation_gt_pandas\n",
    ")\n",
    "test_dataset_pandas = create_dataset_object(\n",
    "    schema, raw_test_events_pandas, item_features_pandas\n",
    ")\n",
    "test_gt_pandas = create_dataset_object(\n",
    "    gt_schema, raw_test_gt_pandas\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "train_dataset_spark = create_dataset_object(\n",
    "    schema, raw_train_events_spark, item_features_spark\n",
    ")\n",
    "validation_dataset_spark = create_dataset_object(\n",
    "    schema, raw_validation_events_spark, item_features_spark\n",
    ")\n",
    "validation_gt_spark = create_dataset_object(\n",
    "    gt_schema, raw_validation_gt_spark\n",
    ")\n",
    "test_dataset_spark = create_dataset_object(\n",
    "    schema, raw_test_events_spark, item_features_spark\n",
    ")\n",
    "test_gt_spark = create_dataset_object(\n",
    "    gt_schema, raw_test_gt_spark\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset_polars = create_dataset_object(\n",
    "    schema, raw_train_events_polars, item_features_polars\n",
    ")\n",
    "validation_dataset_polars = create_dataset_object(\n",
    "    schema, raw_validation_events_polars, item_features_polars\n",
    ")\n",
    "validation_gt_polars = create_dataset_object(\n",
    "    gt_schema, raw_validation_gt_polars\n",
    ")\n",
    "test_dataset_polars = create_dataset_object(\n",
    "    schema, raw_test_events_polars, item_features_polars\n",
    ")\n",
    "test_gt_polars = create_dataset_object(\n",
    "    gt_schema, raw_test_gt_polars\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "ITEM_FEATURE_NAME = \"item_id_seq\"\n",
    "\n",
    "tensor_schema = TensorSchema(\n",
    "    TensorFeatureInfo(\n",
    "        name=ITEM_FEATURE_NAME,\n",
    "        is_seq=True,\n",
    "        feature_type=FeatureType.CATEGORICAL,\n",
    "        feature_sources=[TensorFeatureSource(FeatureSource.INTERACTIONS, train_dataset_pandas.feature_schema.item_id_column)],\n",
    "        feature_hint=FeatureHint.ITEM_ID,\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tokenizing comparison"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Fit tokenizer comparison"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n 2 -r 3\n",
    "\n",
    "tokenizer = SequenceTokenizer(tensor_schema, allow_collect_to_master=True)\n",
    "tokenizer.fit(train_dataset_pandas)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "269 ms ± 18.8 ms per loop (mean ± std. dev. of 3 runs, 2 loops each)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n 2 -r 3\n",
    "\n",
    "tokenizer = SequenceTokenizer(tensor_schema, allow_collect_to_master=True)\n",
    "tokenizer.fit(train_dataset_spark)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "5.39 s ± 425 ms per loop (mean ± std. dev. of 3 runs, 2 loops each)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n 2 -r 3\n",
    "\n",
    "tokenizer = SequenceTokenizer(tensor_schema, allow_collect_to_master=True)\n",
    "tokenizer.fit(train_dataset_polars)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "180 ms ± 40.7 ms per loop (mean ± std. dev. of 3 runs, 2 loops each)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer_pandas = SequenceTokenizer(tensor_schema, allow_collect_to_master=True)\n",
    "tokenizer_spark = SequenceTokenizer(tensor_schema, allow_collect_to_master=True)\n",
    "tokenizer_polars = SequenceTokenizer(tensor_schema, allow_collect_to_master=True)\n",
    "\n",
    "tokenizer_pandas.fit(train_dataset_pandas)\n",
    "tokenizer_spark.fit(train_dataset_spark)\n",
    "tokenizer_polars.fit(train_dataset_polars)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Transform dataset with fitted tokenizer comparison"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n 2 -r 3\n",
    "\n",
    "tokenizer_pandas.transform(train_dataset_pandas)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "24 s ± 87.1 ms per loop (mean ± std. dev. of 3 runs, 2 loops each)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n 2 -r 3\n",
    "\n",
    "tokenizer_spark.transform(train_dataset_spark)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "41.2 s ± 2.56 s per loop (mean ± std. dev. of 3 runs, 2 loops each)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n 2 -r 3\n",
    "\n",
    "tokenizer_polars.transform(train_dataset_polars)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2.29 s ± 20.2 ms per loop (mean ± std. dev. of 3 runs, 2 loops each)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Due to there are only `PandasSequentialDataset` and `PolarsSequentialDataset`, we can get rid of spark frames now"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "sequential_train_dataset = tokenizer_pandas.transform(train_dataset_pandas)\n",
    "sequential_validation_dataset = tokenizer_pandas.transform(validation_dataset_pandas)\n",
    "sequential_validation_gt = tokenizer_pandas.transform(validation_gt_pandas, [tensor_schema.item_id_feature_name])\n",
    "\n",
    "sequential_validation_dataset, sequential_validation_gt = SequentialDataset.keep_common_query_ids(\n",
    "    sequential_validation_dataset, sequential_validation_gt\n",
    ")\n",
    "\n",
    "test_query_ids = test_gt_pandas.query_ids\n",
    "test_query_ids_np = tokenizer_pandas.query_id_encoder.transform(test_query_ids)[\"user_id\"].values\n",
    "sequential_test_dataset = tokenizer_pandas.transform(test_dataset_pandas).filter_by_query_id(test_query_ids_np)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "MAX_SEQ_LEN = 200\n",
    "BATCH_SIZE = 512\n",
    "NUM_WORKERS = 4"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will not waste time training the model, so we will get the training time for only 1 epoch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "GPU available: True (cuda), used: True\n",
      "TPU available: False, using: 0 TPU cores\n",
      "IPU available: False, using: 0 IPUs\n",
      "HPU available: False, using: 0 HPUs\n"
     ]
    }
   ],
   "source": [
    "model = SasRec(\n",
    "    tensor_schema,\n",
    "    block_count=2,\n",
    "    head_count=2,\n",
    "    max_seq_len=MAX_SEQ_LEN,\n",
    "    hidden_size=300,\n",
    "    dropout_rate=0.5,\n",
    "    optimizer_factory=FatOptimizerFactory(learning_rate=0.001),\n",
    ")\n",
    "\n",
    "validation_metrics_callback = ValidationMetricsCallback(\n",
    "    metrics=[\"map\", \"ndcg\", \"recall\"],\n",
    "    ks=[1, 5, 10, 20],\n",
    "    item_count=train_dataset_pandas.item_count,\n",
    "    postprocessors=[RemoveSeenItems(sequential_validation_dataset)]\n",
    ")\n",
    "\n",
    "trainer = L.Trainer(\n",
    "    max_epochs=1,\n",
    "    callbacks=[validation_metrics_callback],\n",
    "    logger=False,\n",
    "    log_every_n_steps=1000,\n",
    ")\n",
    "\n",
    "train_dataloader = DataLoader(\n",
    "    dataset=SasRecTrainingDataset(\n",
    "        sequential_train_dataset,\n",
    "        max_sequence_length=MAX_SEQ_LEN,\n",
    "    ),\n",
    "    batch_size=BATCH_SIZE,\n",
    "    shuffle=True,\n",
    "    num_workers=4,\n",
    "    pin_memory=True,\n",
    ")\n",
    "\n",
    "validation_dataloader = DataLoader(\n",
    "    dataset=SasRecValidationDataset(\n",
    "        sequential_validation_dataset,\n",
    "        sequential_validation_gt,\n",
    "        sequential_train_dataset,\n",
    "        max_sequence_length=MAX_SEQ_LEN,\n",
    "    ),\n",
    "    batch_size=BATCH_SIZE,\n",
    "    num_workers=NUM_WORKERS,\n",
    "    pin_memory=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "You are using a CUDA device ('NVIDIA A100 80GB PCIe') that has Tensor Cores. To properly utilize them, you should set `torch.set_float32_matmul_precision('medium' | 'high')` which will trade-off precision for performance. For more details, read https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html#torch.set_float32_matmul_precision\n",
      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n",
      "\n",
      "  | Name   | Type             | Params\n",
      "--------------------------------------------\n",
      "0 | _model | SasRecModel      | 9.3 M \n",
      "1 | _loss  | CrossEntropyLoss | 0     \n",
      "--------------------------------------------\n",
      "9.3 M     Trainable params\n",
      "0         Non-trainable params\n",
      "9.3 M     Total params\n",
      "37.321    Total estimated model params size (MB)\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4ffa34ac3ef04e1a9f84191f40d718aa",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Sanity Checking: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d8617379ba6a43d491b8796e7af183aa",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Training: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "b4de40496a5540b1be02e1fced25a367",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Validation: |          | 0/? [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "k              1        10        20         5\n",
      "map     0.032406  0.063699  0.068970  0.056406\n",
      "ndcg    0.032406  0.085727  0.105130  0.067868\n",
      "recall  0.032406  0.158535  0.235694  0.102901\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "`Trainer.fit` stopped: `max_epochs=1` reached.\n"
     ]
    }
   ],
   "source": [
    "trainer.fit(\n",
    "    model,\n",
    "    train_dataloaders=train_dataloader,\n",
    "    val_dataloaders=validation_dataloader,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inference callbacks comparison\n",
    "\n",
    "Let's launch the inference of the model. \n",
    "\n",
    "Please note that no time measurement is performed at this stage, because the result will be obtained inside the callbacks in the torch.Tensor format.\n",
    "\n",
    "Conversion to dataframes will begin only when we call the `get_result` function for callbacks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "TOPK = [1, 10, 20, 100]\n",
    "postprocessors = [RemoveSeenItems(sequential_test_dataset)]\n",
    "\n",
    "prediction_dataloader = DataLoader(\n",
    "    dataset=SasRecPredictionDataset(\n",
    "        sequential_test_dataset,\n",
    "        max_sequence_length=MAX_SEQ_LEN,\n",
    "    ),\n",
    "    batch_size=BATCH_SIZE,\n",
    "    num_workers=NUM_WORKERS,\n",
    "    pin_memory=True,\n",
    ")\n",
    "\n",
    "spark_prediction_callback = SparkPredictionCallback(\n",
    "    spark_session=spark_session,\n",
    "    top_k=max(TOPK),\n",
    "    query_column=\"user_id\",\n",
    "    item_column=\"item_id\",\n",
    "    rating_column=\"score\",\n",
    "    postprocessors=postprocessors,\n",
    ")\n",
    "\n",
    "pandas_prediction_callback = PandasPredictionCallback(\n",
    "    top_k=max(TOPK),\n",
    "    query_column=\"user_id\",\n",
    "    item_column=\"item_id\",\n",
    "    rating_column=\"score\",\n",
    "    postprocessors=postprocessors,\n",
    ")\n",
    "\n",
    "polars_prediction_callback = PolarsPredictionCallback(\n",
    "    top_k=max(TOPK),\n",
    "    query_column=\"user_id\",\n",
    "    item_column=\"item_id\",\n",
    "    rating_column=\"score\",\n",
    "    postprocessors=postprocessors,\n",
    ")\n",
    "\n",
    "trainer = L.Trainer(\n",
    "    callbacks=[\n",
    "        spark_prediction_callback,\n",
    "        pandas_prediction_callback,\n",
    "        polars_prediction_callback\n",
    "    ], \n",
    "    inference_mode=True\n",
    ")\n",
    "\n",
    "trainer.predict(model, dataloaders=prediction_dataloader, return_predictions=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n 2 -r 3\n",
    "\n",
    "pandas_prediction_callback.get_result()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "3.82 s ± 15.5 ms per loop (mean ± std. dev. of 3 runs, 2 loops each)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n 2 -r 3\n",
    "\n",
    "spark_prediction_callback.get_result()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "13.4 s ± 312 ms per loop (mean ± std. dev. of 3 runs, 2 loops each)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n 2 -r 3\n",
    "\n",
    "polars_prediction_callback.get_result()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "777 ms ± 27.2 ms per loop (mean ± std. dev. of 3 runs, 2 loops each)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Metrics comparison"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "result_spark = spark_prediction_callback.get_result()\n",
    "result_pandas = pandas_prediction_callback.get_result()\n",
    "result_polars = polars_prediction_callback.get_result()\n",
    "\n",
    "recommendations_spark = tokenizer_spark.query_and_item_id_encoder.inverse_transform(result_spark)\n",
    "recommendations_pandas = tokenizer_pandas.query_and_item_id_encoder.inverse_transform(result_pandas)\n",
    "recommendations_polars = tokenizer_spark.query_and_item_id_encoder.inverse_transform(result_polars)\n",
    "# Polars is sensitive to different dtypes, so we need to match them in user column\n",
    "raw_test_gt_polars = raw_test_gt_polars.with_columns(\n",
    "    pl.col(\"user_id\").cast(recommendations_polars.get_column(\"user_id\").dtype)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "init_args = {\"query_column\": \"user_id\", \"item_column\": \"item_id\", \"rating_column\": \"score\"}\n",
    "result_metrics = OfflineMetrics(\n",
    "    [Recall(TOPK), Precision(TOPK), MAP(TOPK), NDCG(TOPK), MRR(TOPK), HitRate(TOPK)], **init_args\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n 2 -r 3\n",
    "\n",
    "result_metrics(recommendations_pandas, raw_test_gt_pandas)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "3min 28s ± 15.7 s per loop (mean ± std. dev. of 3 runs, 2 loops each)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n 2 -r 3\n",
    "\n",
    "result_metrics(recommendations_spark, raw_test_gt_spark)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "4min 32s ± 3.86 s per loop (mean ± std. dev. of 3 runs, 2 loops each)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%timeit -n 2 -r 3\n",
    "\n",
    "result_metrics(recommendations_polars, raw_test_gt_polars)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "14.5 s ± 35.3 ms per loop (mean ± std. dev. of 3 runs, 2 loops each)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "replay_venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.19"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}