examples/09_sasrec_example.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Example of the SASRec training and inference stages\n",
    "Note that all the given examples can be run without using PySpark, using only Pandas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import lightning as L\n",
    "from lightning.pytorch.loggers import CSVLogger\n",
    "from lightning.pytorch.callbacks import ModelCheckpoint\n",
    "from torch.utils.data import DataLoader\n",
    "import torch\n",
    "\n",
    "from replay.metrics import OfflineMetrics, Recall, Precision, MAP, NDCG, HitRate, MRR\n",
    "from replay.metrics.torch_metrics_builder import metrics_to_df\n",
    "from replay.splitters import LastNSplitter\n",
    "from replay.utils import get_spark_session\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",
    "    TorchPredictionCallback,\n",
    "    QueryEmbeddingsPredictionCallback\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",
    "    SasRecPredictionBatch,\n",
    "    SasRecModel\n",
    ")\n",
    "\n",
    "import pandas as pd"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Getting a spark session"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "spark_session = get_spark_session()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Prepare data\n",
    "### Load raw movielens-1M interactions, item features and user features.\n",
    "In the current implementation, the SASRec does not take into account the features of items or users. They are only used to get a complete list of users and items."
   ]
  },
  {
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "movielens = MovieLens(\"1m\")\n",
    "interactions = movielens.ratings\n",
    "user_features = movielens.users\n",
    "item_features = movielens.items"
   ]
  },
  {
   "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>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>1193</td>\n",
       "      <td>5</td>\n",
       "      <td>978300760</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>661</td>\n",
       "      <td>3</td>\n",
       "      <td>978302109</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1</td>\n",
       "      <td>914</td>\n",
       "      <td>3</td>\n",
       "      <td>978301968</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1</td>\n",
       "      <td>3408</td>\n",
       "      <td>4</td>\n",
       "      <td>978300275</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1</td>\n",
       "      <td>2355</td>\n",
       "      <td>5</td>\n",
       "      <td>978824291</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   user_id  item_id  rating  timestamp\n",
       "0        1     1193       5  978300760\n",
       "1        1      661       3  978302109\n",
       "2        1      914       3  978301968\n",
       "3        1     3408       4  978300275\n",
       "4        1     2355       5  978824291"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "interactions.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "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>gender</th>\n",
       "      <th>age</th>\n",
       "      <th>occupation</th>\n",
       "      <th>zip_code</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>F</td>\n",
       "      <td>1</td>\n",
       "      <td>10</td>\n",
       "      <td>48067</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2</td>\n",
       "      <td>M</td>\n",
       "      <td>56</td>\n",
       "      <td>16</td>\n",
       "      <td>70072</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3</td>\n",
       "      <td>M</td>\n",
       "      <td>25</td>\n",
       "      <td>15</td>\n",
       "      <td>55117</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>4</td>\n",
       "      <td>M</td>\n",
       "      <td>45</td>\n",
       "      <td>7</td>\n",
       "      <td>02460</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>5</td>\n",
       "      <td>M</td>\n",
       "      <td>25</td>\n",
       "      <td>20</td>\n",
       "      <td>55455</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   user_id gender  age  occupation zip_code\n",
       "0        1      F    1          10    48067\n",
       "1        2      M   56          16    70072\n",
       "2        3      M   25          15    55117\n",
       "3        4      M   45           7    02460\n",
       "4        5      M   25          20    55455"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "user_features.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "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>Animation|Children's|Comedy</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2</td>\n",
       "      <td>Jumanji (1995)</td>\n",
       "      <td>Adventure|Children's|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</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                        genres\n",
       "0        1                    Toy Story (1995)   Animation|Children's|Comedy\n",
       "1        2                      Jumanji (1995)  Adventure|Children's|Fantasy\n",
       "2        3             Grumpier Old Men (1995)                Comedy|Romance\n",
       "3        4            Waiting to Exhale (1995)                  Comedy|Drama\n",
       "4        5  Father of the Bride Part II (1995)                        Comedy"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "item_features.head()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Removing duplicates in the timestamp column without changing the original items order where timestamp is the same"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "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>1000138</th>\n",
       "      <td>6040</td>\n",
       "      <td>858</td>\n",
       "      <td>4</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1000153</th>\n",
       "      <td>6040</td>\n",
       "      <td>2384</td>\n",
       "      <td>4</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>999873</th>\n",
       "      <td>6040</td>\n",
       "      <td>593</td>\n",
       "      <td>5</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1000007</th>\n",
       "      <td>6040</td>\n",
       "      <td>1961</td>\n",
       "      <td>4</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1000192</th>\n",
       "      <td>6040</td>\n",
       "      <td>2019</td>\n",
       "      <td>5</td>\n",
       "      <td>4</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>825793</th>\n",
       "      <td>4958</td>\n",
       "      <td>2399</td>\n",
       "      <td>1</td>\n",
       "      <td>446</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>825438</th>\n",
       "      <td>4958</td>\n",
       "      <td>1407</td>\n",
       "      <td>5</td>\n",
       "      <td>447</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>825724</th>\n",
       "      <td>4958</td>\n",
       "      <td>3264</td>\n",
       "      <td>4</td>\n",
       "      <td>448</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>825731</th>\n",
       "      <td>4958</td>\n",
       "      <td>2634</td>\n",
       "      <td>3</td>\n",
       "      <td>449</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>825603</th>\n",
       "      <td>4958</td>\n",
       "      <td>1924</td>\n",
       "      <td>4</td>\n",
       "      <td>450</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>1000209 rows × 4 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "         user_id  item_id  rating  timestamp\n",
       "1000138     6040      858       4          0\n",
       "1000153     6040     2384       4          1\n",
       "999873      6040      593       5          2\n",
       "1000007     6040     1961       4          3\n",
       "1000192     6040     2019       5          4\n",
       "...          ...      ...     ...        ...\n",
       "825793      4958     2399       1        446\n",
       "825438      4958     1407       5        447\n",
       "825724      4958     3264       4        448\n",
       "825731      4958     2634       3        449\n",
       "825603      4958     1924       4        450\n",
       "\n",
       "[1000209 rows x 4 columns]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "interactions[\"timestamp\"] = interactions[\"timestamp\"].astype(\"int64\")\n",
    "interactions = interactions.sort_values(by=\"timestamp\")\n",
    "interactions[\"timestamp\"] = interactions.groupby(\"user_id\").cumcount()\n",
    "interactions"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Split interactions into the train, validation and test datasets using LastNSplitter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "splitter = LastNSplitter(\n",
    "    N=1,\n",
    "    divide_column=\"user_id\",\n",
    "    query_column=\"user_id\",\n",
    "    strategy=\"interactions\",\n",
    ")\n",
    "\n",
    "raw_test_events, raw_test_gt = splitter.split(interactions)\n",
    "raw_validation_events, raw_validation_gt = splitter.split(raw_test_events)\n",
    "raw_train_events = raw_validation_events"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prepare FeatureSchema required to create Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "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"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Dataset for the training stage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = Dataset(\n",
    "    feature_schema=prepare_feature_schema(is_ground_truth=False),\n",
    "    interactions=raw_train_events,\n",
    "    query_features=user_features,\n",
    "    item_features=item_features,\n",
    "    check_consistency=True,\n",
    "    categorical_encoded=False,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Datasets (events and ground_truth) for the validation stage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "validation_dataset = Dataset(\n",
    "    feature_schema=prepare_feature_schema(is_ground_truth=False),\n",
    "    interactions=raw_validation_events,\n",
    "    query_features=user_features,\n",
    "    item_features=item_features,\n",
    "    check_consistency=True,\n",
    "    categorical_encoded=False,\n",
    ")\n",
    "validation_gt = Dataset(\n",
    "    feature_schema=prepare_feature_schema(is_ground_truth=True),\n",
    "    interactions=raw_validation_gt,\n",
    "    check_consistency=True,\n",
    "    categorical_encoded=False,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Datasets (events and ground_truth) for the testing stage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_dataset = Dataset(\n",
    "    feature_schema=prepare_feature_schema(is_ground_truth=False),\n",
    "    interactions=raw_test_events,\n",
    "    query_features=user_features,\n",
    "    item_features=item_features,\n",
    "    check_consistency=True,\n",
    "    categorical_encoded=False,\n",
    ")\n",
    "test_gt = Dataset(\n",
    "    feature_schema=prepare_feature_schema(is_ground_truth=True),\n",
    "    interactions=raw_test_gt,\n",
    "    check_consistency=True,\n",
    "    categorical_encoded=False,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create the tensor schema\n",
    "A schema shows the correspondence of columns from the source dataset with the internal representation of tensors inside the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "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.feature_schema.item_id_column)],\n",
    "        feature_hint=FeatureHint.ITEM_ID,\n",
    "    )\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create sequential datasets using SequenceTokenizer\n",
    "The SequentialDataset internally store data in the form of sequences of items sorted by increasing interaction time (timestamp). A SequenceTokenizer is used to convert to this format. In addition, the SequenceTokenizer encodes all categorical columns from the source dataset and stores mapping inside itself.\n",
    "SequentialDataset.keep_common_query_ids is used to leave only sequences from the same users"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = SequenceTokenizer(tensor_schema, allow_collect_to_master=True)\n",
    "tokenizer.fit(train_dataset)\n",
    "\n",
    "sequential_train_dataset = tokenizer.transform(train_dataset)\n",
    "\n",
    "sequential_validation_dataset = tokenizer.transform(validation_dataset)\n",
    "sequential_validation_gt = tokenizer.transform(validation_gt, [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",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_query_ids = test_gt.query_ids\n",
    "test_query_ids_np = tokenizer.query_id_encoder.transform(test_query_ids)[\"user_id\"].values\n",
    "sequential_test_dataset = tokenizer.transform(test_dataset).filter_by_query_id(test_query_ids_np)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can get the user and item mapping and inverse mapping as follows"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(tokenizer.query_id_encoder.mapping, tokenizer.query_id_encoder.inverse_mapping)\n",
    "print(tokenizer.item_id_encoder.mapping, tokenizer.item_id_encoder.inverse_mapping)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train model\n",
    "### Create SASRec model instance and run the training stage using lightning\n",
    "After each epoch validation metrics are shown. You can change the list of validation metrics in ValidationMetricsCallback\n",
    "The model is determined to be the best and is saved if the metric updates its maximum during validation (see the ModelCheckpoint)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MAX_SEQ_LEN = 200\n",
    "BATCH_SIZE = 512\n",
    "NUM_WORKERS = 4\n",
    "\n",
    "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",
    "csv_logger = CSVLogger(save_dir=\".logs/train\", name=\"SASRec_example\")\n",
    "\n",
    "checkpoint_callback = ModelCheckpoint(\n",
    "    dirpath=\".checkpoints\",\n",
    "    save_top_k=1,\n",
    "    verbose=True,\n",
    "    # if you use multiple dataloaders, then add the serial number of the dataloader to the suffix of the metric name.\n",
    "    # For example,\"recall@10/dataloader_idx_0\"\n",
    "    monitor=\"recall@10\",\n",
    "    mode=\"max\",\n",
    ")\n",
    "\n",
    "validation_metrics_callback = ValidationMetricsCallback(\n",
    "    metrics=[\"map\", \"ndcg\", \"recall\"],\n",
    "    ks=[1, 5, 10, 20],\n",
    "    item_count=train_dataset.item_count,\n",
    "    postprocessors=[RemoveSeenItems(sequential_validation_dataset)]\n",
    ")\n",
    "\n",
    "trainer = L.Trainer(\n",
    "    max_epochs=100,\n",
    "    callbacks=[checkpoint_callback, validation_metrics_callback],\n",
    "    logger=csv_logger,\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",
    ")\n",
    "\n",
    "trainer.fit(\n",
    "    model,\n",
    "    train_dataloaders=train_dataloader,\n",
    "    val_dataloaders=validation_dataloader,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The path to the best model is saved inside checkpoint_callback"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "best_model = SasRec.load_from_checkpoint(checkpoint_callback.best_model_path)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Inference stage\n",
    "### Prepare Dataloader and logger"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "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",
    "csv_logger = CSVLogger(save_dir=\".logs/test\", name=\"SASRec_example\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run inference\n",
    "You can get the recommendations in three formats: PySpark DataFrame, Pandas DataFrame, PyTorch tensors. Each of the types corresponds a callback\n",
    "You can filter the results using postprocessors strategy. For example the RemoveSeenItems postprocessor is filtering out the items that already have been seen in test dataset\n",
    "You don't need to use all three callbacks. This is shown only for example\n",
    "\n",
    "Also, you can get user embeddings, that were used to perform predictions, using `get_query_embedding` method inside SasRecModel or `QueryEmbeddingsPredictionCallback` for lightning module."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "TOPK = [1, 10, 20, 100]\n",
    "\n",
    "postprocessors = [RemoveSeenItems(sequential_test_dataset)]\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",
    "torch_prediction_callback = TorchPredictionCallback(\n",
    "    top_k=max(TOPK),\n",
    "    postprocessors=postprocessors,\n",
    ")\n",
    "\n",
    "query_embeddings_callback = QueryEmbeddingsPredictionCallback()\n",
    "\n",
    "trainer = L.Trainer(\n",
    "    callbacks=[\n",
    "        spark_prediction_callback,\n",
    "        pandas_prediction_callback,\n",
    "        torch_prediction_callback,\n",
    "        query_embeddings_callback\n",
    "    ], \n",
    "    logger=csv_logger, \n",
    "    inference_mode=True\n",
    ")\n",
    "trainer.predict(best_model, dataloaders=prediction_dataloader, return_predictions=False)\n",
    "\n",
    "spark_res = spark_prediction_callback.get_result()\n",
    "pandas_res = pandas_prediction_callback.get_result()\n",
    "torch_user_ids, torch_item_ids, torch_scores = torch_prediction_callback.get_result()\n",
    "user_embeddings = query_embeddings_callback.get_result()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+-------+-------+------------------+\n",
      "|user_id|item_id|             score|\n",
      "+-------+-------+------------------+\n",
      "|      0|   2012| 8.138867378234863|\n",
      "|      0|   1642|  8.09422492980957|\n",
      "|      0|    360| 8.087743759155273|\n",
      "|      0|   2325| 7.976056098937988|\n",
      "|      0|     47|  7.48519229888916|\n",
      "|      0|   2020| 7.362086296081543|\n",
      "|      0|   2016| 7.105367660522461|\n",
      "|      0|   2009| 7.090343475341797|\n",
      "|      0|   1020|  7.08748197555542|\n",
      "|      0|   2014| 6.992861747741699|\n",
      "|      0|     12| 6.892394542694092|\n",
      "|      0|   1382| 6.864995956420898|\n",
      "|      0|    612| 6.860413074493408|\n",
      "|      0|   2027|  6.73327112197876|\n",
      "|      0|   3088| 6.721376419067383|\n",
      "|      0|   2018| 6.611262321472168|\n",
      "|      0|   1012| 6.602219581604004|\n",
      "|      0|   1019|6.5607428550720215|\n",
      "|      0|   3090| 6.448214530944824|\n",
      "|      0|   2965| 6.409302711486816|\n",
      "+-------+-------+------------------+\n",
      "only showing top 20 rows\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "spark_res.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "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>score</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>2012</td>\n",
       "      <td>8.138867</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>1642</td>\n",
       "      <td>8.094225</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>360</td>\n",
       "      <td>8.087744</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>2325</td>\n",
       "      <td>7.976056</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0</td>\n",
       "      <td>47</td>\n",
       "      <td>7.485192</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6039</th>\n",
       "      <td>6039</td>\n",
       "      <td>1251</td>\n",
       "      <td>3.320609</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6039</th>\n",
       "      <td>6039</td>\n",
       "      <td>2174</td>\n",
       "      <td>3.301445</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6039</th>\n",
       "      <td>6039</td>\n",
       "      <td>3197</td>\n",
       "      <td>3.286182</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6039</th>\n",
       "      <td>6039</td>\n",
       "      <td>3141</td>\n",
       "      <td>3.280689</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6039</th>\n",
       "      <td>6039</td>\n",
       "      <td>2849</td>\n",
       "      <td>3.2665</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>604000 rows × 3 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "      user_id item_id     score\n",
       "0           0    2012  8.138867\n",
       "0           0    1642  8.094225\n",
       "0           0     360  8.087744\n",
       "0           0    2325  7.976056\n",
       "0           0      47  7.485192\n",
       "...       ...     ...       ...\n",
       "6039     6039    1251  3.320609\n",
       "6039     6039    2174  3.301445\n",
       "6039     6039    3197  3.286182\n",
       "6039     6039    3141  3.280689\n",
       "6039     6039    2849    3.2665\n",
       "\n",
       "[604000 rows x 3 columns]"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pandas_res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(0) tensor([2012, 1642,  360, 2325,   47, 2020, 2016, 2009, 1020, 2014,   12, 1382,\n",
      "         612, 2027, 3088, 2018, 1012, 1019, 3090, 2965,  547, 2068, 2021, 1990,\n",
      "        2073, 2315, 2072,  700,  105, 1050,  582, 1346, 2011, 2918,  310, 1743,\n",
      "        1977,  314, 3546, 1058,  974, 1010, 1850, 2023,  358, 1979,  259, 2631,\n",
      "         667,   33, 2030, 2070, 2013, 1002, 2640, 3328, 2807,  592, 1964, 3144,\n",
      "        2830, 2985, 3327, 3329, 1262, 1017, 1547,  476,    1, 2285, 1459, 1664,\n",
      "        1018,   59, 2692, 1000, 1936, 1851, 3218, 3414, 1539,  826, 1011, 3106,\n",
      "         999,  791, 2092, 1004, 2047,  236, 2731, 3676, 2677, 3682,  156, 2071,\n",
      "        1008,  363, 1945, 2777]) tensor([8.1389, 8.0942, 8.0877, 7.9761, 7.4852, 7.3621, 7.1054, 7.0903, 7.0875,\n",
      "        6.9929, 6.8924, 6.8650, 6.8604, 6.7333, 6.7214, 6.6113, 6.6022, 6.5607,\n",
      "        6.4482, 6.4093, 6.3793, 6.3753, 6.2767, 6.2201, 6.2057, 6.1955, 6.1510,\n",
      "        6.1112, 6.0373, 6.0356, 5.9954, 5.9927, 5.9892, 5.9694, 5.9151, 5.8761,\n",
      "        5.8475, 5.8284, 5.7449, 5.7310, 5.7161, 5.6727, 5.6549, 5.6495, 5.6186,\n",
      "        5.6085, 5.6031, 5.5880, 5.5866, 5.5690, 5.5257, 5.4499, 5.3953, 5.3695,\n",
      "        5.3578, 5.3515, 5.3228, 5.3165, 5.3163, 5.3155, 5.2968, 5.2931, 5.2828,\n",
      "        5.2674, 5.2641, 5.2480, 5.2354, 5.2027, 5.1985, 5.1930, 5.1861, 5.1736,\n",
      "        5.1604, 5.1566, 5.1563, 5.1207, 5.0993, 5.0964, 5.0871, 5.0869, 5.0672,\n",
      "        5.0649, 5.0559, 5.0418, 5.0282, 5.0174, 4.9985, 4.9898, 4.9881, 4.8736,\n",
      "        4.8252, 4.7583, 4.7408, 4.7300, 4.6881, 4.6835, 4.6771, 4.6612, 4.6486,\n",
      "        4.5319])\n"
     ]
    }
   ],
   "source": [
    "print(torch_user_ids[0], torch_item_ids[0], torch_scores[0])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Suppose we want to get the recomendations in PySpark format. \n",
    "Let's get the inverse representation of labels using inverse_transform method.\n",
    "\n",
    "Note that the reverse representation can only be obtained for PySpark and Pandas formats. When working with PyTorch tensors, the reverse representation must be done manually"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "recommendations = tokenizer.query_and_item_id_encoder.inverse_transform(spark_res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+------------------+-------+-------+\n",
      "|             score|user_id|item_id|\n",
      "+------------------+-------+-------+\n",
      "| 8.138867378234863|      1|   2081|\n",
      "|  8.09422492980957|      1|   1688|\n",
      "| 8.087743759155273|      1|    364|\n",
      "| 7.976056098937988|      1|   2394|\n",
      "|  7.48519229888916|      1|     48|\n",
      "| 7.362086296081543|      1|   2089|\n",
      "| 7.105367660522461|      1|   2085|\n",
      "| 7.090343475341797|      1|   2078|\n",
      "|  7.08748197555542|      1|   1033|\n",
      "| 6.992861747741699|      1|   2083|\n",
      "| 6.892394542694092|      1|     13|\n",
      "| 6.864995956420898|      1|   1405|\n",
      "| 6.860413074493408|      1|    616|\n",
      "|  6.73327112197876|      1|   2096|\n",
      "| 6.721376419067383|      1|   3157|\n",
      "| 6.611262321472168|      1|   2087|\n",
      "| 6.602219581604004|      1|   1025|\n",
      "|6.5607428550720215|      1|   1032|\n",
      "| 6.448214530944824|      1|   3159|\n",
      "| 6.409302711486816|      1|   3034|\n",
      "+------------------+-------+-------+\n",
      "only showing top 20 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "recommendations.show()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Calculating metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "init_args = {\"query_column\": \"user_id\", \"rating_column\": \"score\"}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "result_metrics = OfflineMetrics(\n",
    "    [Recall(TOPK), Precision(TOPK), MAP(TOPK), NDCG(TOPK), MRR(TOPK), HitRate(TOPK)], **init_args\n",
    ")(recommendations.toPandas(), raw_test_gt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "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>k</th>\n",
       "      <th>1</th>\n",
       "      <th>10</th>\n",
       "      <th>20</th>\n",
       "      <th>100</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>HitRate</th>\n",
       "      <td>0.06904</td>\n",
       "      <td>0.296026</td>\n",
       "      <td>0.415397</td>\n",
       "      <td>0.680464</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>MAP</th>\n",
       "      <td>0.06904</td>\n",
       "      <td>0.126483</td>\n",
       "      <td>0.134644</td>\n",
       "      <td>0.141227</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>MRR</th>\n",
       "      <td>0.06904</td>\n",
       "      <td>0.126483</td>\n",
       "      <td>0.134644</td>\n",
       "      <td>0.141227</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>NDCG</th>\n",
       "      <td>0.06904</td>\n",
       "      <td>0.165998</td>\n",
       "      <td>0.196025</td>\n",
       "      <td>0.244784</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Precision</th>\n",
       "      <td>0.06904</td>\n",
       "      <td>0.029603</td>\n",
       "      <td>0.020770</td>\n",
       "      <td>0.006805</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>Recall</th>\n",
       "      <td>0.06904</td>\n",
       "      <td>0.296026</td>\n",
       "      <td>0.415397</td>\n",
       "      <td>0.680464</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "k              1         10        20        100\n",
       "HitRate    0.06904  0.296026  0.415397  0.680464\n",
       "MAP        0.06904  0.126483  0.134644  0.141227\n",
       "MRR        0.06904  0.126483  0.134644  0.141227\n",
       "NDCG       0.06904  0.165998  0.196025  0.244784\n",
       "Precision  0.06904  0.029603  0.020770  0.006805\n",
       "Recall     0.06904  0.296026  0.415397  0.680464"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "metrics_to_df(result_metrics)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### User embeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Got 6040 x 300 user embeddings, because among all 12 batches: \n",
    "\n",
    "11 batches contains 512 samples\n",
    "\n",
    "1 batch contains 408 left samples\n",
    "\n",
    "11 * 512 + 408 == 6040"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.3704,  1.9303,  1.9161,  ..., -0.1143,  1.5672, -0.9027],\n",
       "        [-1.8589, -2.6053,  0.1639,  ..., -0.4250, -2.5378, -0.8178],\n",
       "        [-0.0877,  0.3147,  2.4707,  ..., -0.7766, -0.8095, -1.2838],\n",
       "        ...,\n",
       "        [-0.5727,  0.1162,  1.5556,  ..., -1.8936,  1.4968, -1.3920],\n",
       "        [-0.3458, -1.6828, -0.5697,  ..., -1.6742,  3.2049,  0.6597],\n",
       "        [-1.2831,  0.7636,  1.0783,  ..., -1.1996, -0.0189,  1.4718]],\n",
       "       device='cuda:0')"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "user_embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([6040, 300])"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "user_embeddings.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can access user embeddings directly with `SasRecModel` class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-4.3240e-01, -1.0684e-01, -8.1389e-02,  ...,  6.6867e-01,\n",
       "          4.7710e-01, -1.7522e+00],\n",
       "        [ 7.8765e-02,  8.0354e-01, -9.5730e-01,  ..., -1.9753e+00,\n",
       "          1.4708e+00, -7.7084e-01],\n",
       "        [ 1.0037e+00,  6.5928e-01, -1.4148e+00,  ...,  1.0593e+00,\n",
       "         -7.4484e-01, -2.7631e-01],\n",
       "        ...,\n",
       "        [-1.4800e+00, -1.9950e-01, -3.6084e-01,  ...,  1.5323e+00,\n",
       "          2.6428e-01, -8.7467e-01],\n",
       "        [ 1.1612e+00, -6.3325e-04,  4.0772e-01,  ...,  2.5503e-01,\n",
       "         -3.4220e-01,  5.4465e-01],\n",
       "        [-6.3809e-01,  1.0407e+00, -2.1736e+00,  ...,  8.7484e-01,\n",
       "          1.4989e+00, -1.1730e+00]], device='cuda:0', grad_fn=<SliceBackward0>)"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "\n",
    "core_model = SasRecModel(\n",
    "    tensor_schema,\n",
    "    num_blocks=2,\n",
    "    num_heads=2,\n",
    "    max_len=MAX_SEQ_LEN,\n",
    "    hidden_size=300,\n",
    "    dropout=0.5\n",
    ")\n",
    "core_model.eval()\n",
    "core_model = core_model.to(device)\n",
    "\n",
    "# Get first batch of data \n",
    "data = next(iter(prediction_dataloader))\n",
    "tensor_map, padding_mask = data.features, data.padding_mask\n",
    "\n",
    "# Ensure everything is on the same device\n",
    "padding_mask = padding_mask.to(device)\n",
    "tensor_map[\"item_id_seq\"] = tensor_map[\"item_id_seq\"].to(device)\n",
    "\n",
    "# Get user embeddings\n",
    "user_embeddings_batch = core_model.get_query_embeddings(tensor_map, padding_mask)\n",
    "user_embeddings_batch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([512, 300])"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "user_embeddings_batch.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Item embeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`get_all_embeddings()` method in transformers can be used to get copies of all embeddings that are presented in model as a dict."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'item_embedding': tensor([[-0.0238,  0.0208, -0.0037,  ...,  0.0340,  0.0224, -0.0369],\n",
       "         [-0.0103,  0.0099,  0.0112,  ...,  0.0112,  0.0255, -0.0349],\n",
       "         [ 0.0090, -0.0228, -0.0544,  ...,  0.0024,  0.0438,  0.0060],\n",
       "         ...,\n",
       "         [-0.0220,  0.0150, -0.0027,  ...,  0.0562,  0.0170, -0.0025],\n",
       "         [ 0.0147, -0.0004,  0.0214,  ..., -0.0098, -0.0017,  0.0299],\n",
       "         [-0.0095,  0.0142, -0.0030,  ...,  0.0435, -0.0034,  0.0077]]),\n",
       " 'positional_embedding': tensor([[ 0.1256,  0.0453, -0.0112,  ..., -0.1254,  0.0047, -0.0691],\n",
       "         [-0.0208, -0.0014,  0.0146,  ...,  0.0680, -0.0909, -0.0365],\n",
       "         [-0.0839,  0.0381,  0.0081,  ..., -0.0139, -0.0645, -0.0350],\n",
       "         ...,\n",
       "         [ 0.0322,  0.0308, -0.0525,  ...,  0.0426,  0.0454,  0.0101],\n",
       "         [-0.0335, -0.0006,  0.0082,  ..., -0.0202, -0.0435, -0.0786],\n",
       "         [-0.0537, -0.0550,  0.1043,  ...,  0.0992,  0.0194,  0.0321]])}"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "all_embeddings = best_model.get_all_embeddings()\n",
    "all_embeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can access item embeddings from this dictionary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.0238,  0.0208, -0.0037,  ...,  0.0340,  0.0224, -0.0369],\n",
       "        [-0.0103,  0.0099,  0.0112,  ...,  0.0112,  0.0255, -0.0349],\n",
       "        [ 0.0090, -0.0228, -0.0544,  ...,  0.0024,  0.0438,  0.0060],\n",
       "        ...,\n",
       "        [-0.0220,  0.0150, -0.0027,  ...,  0.0562,  0.0170, -0.0025],\n",
       "        [ 0.0147, -0.0004,  0.0214,  ..., -0.0098, -0.0017,  0.0299],\n",
       "        [-0.0095,  0.0142, -0.0030,  ...,  0.0435, -0.0034,  0.0077]])"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "item_embeddings = all_embeddings[\"item_embedding\"]\n",
    "item_embeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Item embeddings shape is (N_ITEMS, HIDDEN_SIZE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([3883, 300])"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "item_embeddings.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Ensure we got correct dimension and ensure we got the copy of tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert item_embeddings.shape[0] == len(tokenizer.item_id_encoder.mapping[\"item_id\"])\n",
    "assert id(item_embeddings) != id(best_model._model.item_embedder.item_emb.weight.data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For example we observe one new item id in our training data. We can easily expand our item embedder by one element"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to expand item embeddings by new size `set_item_embeddings_by_size` method is applied"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "best_model.set_item_embeddings_by_size(item_embeddings.shape[0] + 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now our new item embeddings have one extra embedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_size = best_model.get_all_embeddings()[\"item_embedding\"].shape[0]\n",
    "old_size = item_embeddings.shape[0]\n",
    "\n",
    "assert new_size == old_size + 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Alternatively, we can pass our item embeddings that replace the existing ones by calling `set_item_embeddings_by_tensor`.\n",
    "\n",
    "If tensor contains new items, they will be added to item embedder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_embeddings_weights = torch.rand((new_size + 1, 300))    # randint used for example only\n",
    "\n",
    "best_model.set_item_embeddings_by_tensor(new_embeddings_weights)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "At the moment we expanded our item embeddings by one more item and replace weights by passing `new_embeddings_weights`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "old_size = new_size\n",
    "new_size = best_model.get_all_embeddings()[\"item_embedding\"].shape[0]\n",
    "\n",
    "assert new_size == old_size + 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Similarly, we can append tensor for only new items with no replace for existing by calling `append_item_embeddings`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "new_item_weights = torch.rand((1, 300))    # randint used for example only\n",
    "\n",
    "best_model.append_item_embeddings(new_item_weights)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We passed one new example and its weights to item embeddings, thus expanded our vocabulary by one item again"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "old_size = new_size\n",
    "new_size = best_model.get_all_embeddings()[\"item_embedding\"].shape[0]\n",
    "\n",
    "assert new_size == old_size + 1"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example of launching an inference for a single user without using a trainer (in order to speed up)\n",
    "An example for the production of an online script\n",
    "\n",
    "Let's assume that the user's sequence consisted of a sequence of items [1, 2, 3, 4, 5]. \n",
    "Сreate a padding mask corresponding to the sequence of items.\n",
    "\n",
    "It is important to take only the latest MAX_SEQ_LEN or less items."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [],
   "source": [
    "item_sequence = torch.arange(1, 5).unsqueeze(0)[:, -MAX_SEQ_LEN:]\n",
    "padding_mask = torch.ones_like(item_sequence, dtype=torch.bool)\n",
    "sequence_item_count = item_sequence.shape[1]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Wrapping created tensors in the SasRecPredictionBatch entity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch = SasRecPredictionBatch(\n",
    "    query_id=torch.arange(0, item_sequence.shape[0], 1).long(),\n",
    "    padding_mask=padding_mask.bool(),\n",
    "    features={ITEM_FEATURE_NAME: item_sequence.long()}\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run predict step of the SasRec and get scores from the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 2.7048,  2.2736,  2.1556,  ..., -0.2144, -1.3674,  1.9593]])"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with torch.no_grad():\n",
    "    scores = best_model.predict_step(batch, 0)\n",
    "scores"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Getting three items with the highest score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[   4, 2013, 3032]])"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.topk(scores, k=3).indices"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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,
  "vscode": {
   "interpreter": {
    "hash": "0857f111b041889635bea848a6a183706c3f1c18c9dafdb447caa5e8bea01452"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}