From 2f29eb4d46242fcfd8bc97aa705919e81ea6ae40 Mon Sep 17 00:00:00 2001
From: Nikita Kulikov <neilkulikov@gmail.com>
Date: Mon, 2 Feb 2026 07:11:40 +0300
Subject: [PATCH 1/5] Gitignore updated to recent examples

---
 examples/09_sasrec_example_freq_sampled.ipynb | 4768 +++++++++++++++++
 examples/09_sasrec_example_sampled.ipynb      | 4744 ++++++++++++++++
 examples/09_sasrec_example_thr_sampled.ipynb  | 4768 +++++++++++++++++
 replay/nn/transform/__init__.py               |   10 +
 replay/nn/transform/negative_sampling.py      |  224 +-
 5 files changed, 14500 insertions(+), 14 deletions(-)
 create mode 100644 examples/09_sasrec_example_freq_sampled.ipynb
 create mode 100644 examples/09_sasrec_example_sampled.ipynb
 create mode 100644 examples/09_sasrec_example_thr_sampled.ipynb

diff --git a/examples/09_sasrec_example_freq_sampled.ipynb b/examples/09_sasrec_example_freq_sampled.ipynb
new file mode 100644
index 000000000..4988e2512
--- /dev/null
+++ b/examples/09_sasrec_example_freq_sampled.ipynb
@@ -0,0 +1,4768 @@
+{
+ "cells": [
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Example of SasRec training/inference with Parquet Module"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Seed set to 42\n"
+     ]
+    }
+   ],
+   "source": [
+    "from typing import Optional\n",
+    "\n",
+    "import lightning as L\n",
+    "import pandas as pd\n",
+    "\n",
+    "L.seed_everything(42)\n",
+    "\n",
+    "import warnings\n",
+    "warnings.filterwarnings(\"ignore\")"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Preparing data\n",
+    "In this example, we will be using the MovieLens dataset, namely the 1m subset. It's demonstrated a simple case, so only item ids will be used as model input.\n",
+    "\n",
+    "---\n",
+    "**NOTE**\n",
+    "\n",
+    "Current implementation of SasRec is able to handle item and interactions features. It does not take into account user features. \n",
+    "\n",
+    "---"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "interactions = pd.read_csv(\"./data/ml1m_ratings.dat\", sep=\"\\t\", names=[\"user_id\", \"item_id\",\"rating\",\"timestamp\"])\n",
+    "interactions = interactions.drop(columns=[\"rating\"])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "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>timestamp</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>1000138</th>\n",
+       "      <td>6040</td>\n",
+       "      <td>858</td>\n",
+       "      <td>0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1000153</th>\n",
+       "      <td>6040</td>\n",
+       "      <td>2384</td>\n",
+       "      <td>1</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>999873</th>\n",
+       "      <td>6040</td>\n",
+       "      <td>593</td>\n",
+       "      <td>2</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1000007</th>\n",
+       "      <td>6040</td>\n",
+       "      <td>1961</td>\n",
+       "      <td>3</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1000192</th>\n",
+       "      <td>6040</td>\n",
+       "      <td>2019</td>\n",
+       "      <td>4</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>...</th>\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>446</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>825438</th>\n",
+       "      <td>4958</td>\n",
+       "      <td>1407</td>\n",
+       "      <td>447</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>825724</th>\n",
+       "      <td>4958</td>\n",
+       "      <td>3264</td>\n",
+       "      <td>448</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>825731</th>\n",
+       "      <td>4958</td>\n",
+       "      <td>2634</td>\n",
+       "      <td>449</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>825603</th>\n",
+       "      <td>4958</td>\n",
+       "      <td>1924</td>\n",
+       "      <td>450</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "<p>1000209 rows × 3 columns</p>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "         user_id  item_id  timestamp\n",
+       "1000138     6040      858          0\n",
+       "1000153     6040     2384          1\n",
+       "999873      6040      593          2\n",
+       "1000007     6040     1961          3\n",
+       "1000192     6040     2019          4\n",
+       "...          ...      ...        ...\n",
+       "825793      4958     2399        446\n",
+       "825438      4958     1407        447\n",
+       "825724      4958     3264        448\n",
+       "825731      4958     2634        449\n",
+       "825603      4958     1924        450\n",
+       "\n",
+       "[1000209 rows x 3 columns]"
+      ]
+     },
+     "execution_count": 3,
+     "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"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Encode catagorical data.\n",
+    "To ensure all categorical data is fit for training, it needs to be encoded using the `LabelEncoder` class. Create an instance of the encoder, providing a `LabelEncodingRule` for each categorcial column in the dataset that will be used in model. Note that ids of users and ids of items are always used."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "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>timestamp</th>\n",
+       "      <th>user_id</th>\n",
+       "      <th>item_id</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>12</td>\n",
+       "      <td>0</td>\n",
+       "      <td>0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>68</td>\n",
+       "      <td>1</td>\n",
+       "      <td>0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>67</td>\n",
+       "      <td>2</td>\n",
+       "      <td>0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>12</td>\n",
+       "      <td>3</td>\n",
+       "      <td>0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>140</td>\n",
+       "      <td>4</td>\n",
+       "      <td>0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>...</th>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1000204</th>\n",
+       "      <td>14</td>\n",
+       "      <td>4555</td>\n",
+       "      <td>3705</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1000205</th>\n",
+       "      <td>90</td>\n",
+       "      <td>2813</td>\n",
+       "      <td>3705</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1000206</th>\n",
+       "      <td>70</td>\n",
+       "      <td>2404</td>\n",
+       "      <td>3705</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1000207</th>\n",
+       "      <td>25</td>\n",
+       "      <td>5835</td>\n",
+       "      <td>3705</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1000208</th>\n",
+       "      <td>380</td>\n",
+       "      <td>979</td>\n",
+       "      <td>3705</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "<p>1000209 rows × 3 columns</p>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "         timestamp  user_id  item_id\n",
+       "0               12        0        0\n",
+       "1               68        1        0\n",
+       "2               67        2        0\n",
+       "3               12        3        0\n",
+       "4              140        4        0\n",
+       "...            ...      ...      ...\n",
+       "1000204         14     4555     3705\n",
+       "1000205         90     2813     3705\n",
+       "1000206         70     2404     3705\n",
+       "1000207         25     5835     3705\n",
+       "1000208        380      979     3705\n",
+       "\n",
+       "[1000209 rows x 3 columns]"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from replay.preprocessing.label_encoder import LabelEncoder, LabelEncodingRule\n",
+    "\n",
+    "encoder = LabelEncoder(\n",
+    "    [\n",
+    "        LabelEncodingRule(\"user_id\", default_value=\"last\"),\n",
+    "        LabelEncodingRule(\"item_id\", default_value=\"last\"),\n",
+    "    ]\n",
+    ")\n",
+    "interactions = interactions.sort_values(by=\"item_id\", ascending=True)\n",
+    "encoded_interactions = encoder.fit_transform(interactions)\n",
+    "encoded_interactions"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Split interactions into the train, validation and test datasets using LastNSplitter\n",
+    "We use widespread splitting strategy Last-One-Out. We filter out cold items and users for simplicity."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from replay.splitters import LastNSplitter\n",
+    "\n",
+    "splitter = LastNSplitter(\n",
+    "    N=1,\n",
+    "    divide_column=\"user_id\",\n",
+    "    query_column=\"user_id\",\n",
+    "    strategy=\"interactions\",\n",
+    "    drop_cold_users=True,\n",
+    "    drop_cold_items=True\n",
+    ")\n",
+    "\n",
+    "test_events, test_gt = splitter.split(encoded_interactions)\n",
+    "validation_events, validation_gt = splitter.split(test_events)\n",
+    "train_events = validation_events"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Dataset preprocessing (\"baking\")\n",
+    "SasRec expects each user in the batch to provide their events in form of a sequence. For this reason, the event splits must be properly processed using the `groupby_sequences` function provided by RePlay."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from replay.data.nn.utils import groupby_sequences\n",
+    "\n",
+    "\n",
+    "def bake_data(full_data):\n",
+    "    grouped_interactions = groupby_sequences(events=full_data, groupby_col=\"user_id\", sort_col=\"timestamp\")\n",
+    "    return grouped_interactions"
+   ]
+  },
+  {
+   "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>timestamp</th>\n",
+       "      <th>item_id</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>0</td>\n",
+       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
+       "      <td>[2426, 822, 2733, 2587, 2937, 3618, 2943, 708,...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>1</td>\n",
+       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
+       "      <td>[3272, 3026, 2760, 851, 346, 3393, 1107, 515, ...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>2</td>\n",
+       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
+       "      <td>[579, 1140, 1154, 2426, 1524, 1260, 2160, 2621...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>3</td>\n",
+       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
+       "      <td>[1781, 2940, 2468, 890, 948, 106, 593, 309, 49...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>4</td>\n",
+       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
+       "      <td>[1108, 2229, 21, 2435, 2142, 106, 1167, 593, 1...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>...</th>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6035</th>\n",
+       "      <td>6035</td>\n",
+       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
+       "      <td>[2426, 1279, 3151, 3321, 1178, 3301, 2501, 278...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6036</th>\n",
+       "      <td>6036</td>\n",
+       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
+       "      <td>[1592, 2302, 1633, 1813, 2879, 1482, 2651, 250...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6037</th>\n",
+       "      <td>6037</td>\n",
+       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
+       "      <td>[1971, 3500, 2077, 1666, 1399, 2651, 2748, 283...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6038</th>\n",
+       "      <td>6038</td>\n",
+       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
+       "      <td>[1486, 1485, 3384, 3512, 3302, 3126, 3650, 330...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6039</th>\n",
+       "      <td>6039</td>\n",
+       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
+       "      <td>[2432, 2960, 1848, 2114, 2142, 3091, 3248, 317...</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "<p>6040 rows × 3 columns</p>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "      user_id                                          timestamp  \\\n",
+       "0           0  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
+       "1           1  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
+       "2           2  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
+       "3           3  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
+       "4           4  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
+       "...       ...                                                ...   \n",
+       "6035     6035  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
+       "6036     6036  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
+       "6037     6037  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
+       "6038     6038  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
+       "6039     6039  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
+       "\n",
+       "                                                item_id  \n",
+       "0     [2426, 822, 2733, 2587, 2937, 3618, 2943, 708,...  \n",
+       "1     [3272, 3026, 2760, 851, 346, 3393, 1107, 515, ...  \n",
+       "2     [579, 1140, 1154, 2426, 1524, 1260, 2160, 2621...  \n",
+       "3     [1781, 2940, 2468, 890, 948, 106, 593, 309, 49...  \n",
+       "4     [1108, 2229, 21, 2435, 2142, 106, 1167, 593, 1...  \n",
+       "...                                                 ...  \n",
+       "6035  [2426, 1279, 3151, 3321, 1178, 3301, 2501, 278...  \n",
+       "6036  [1592, 2302, 1633, 1813, 2879, 1482, 2651, 250...  \n",
+       "6037  [1971, 3500, 2077, 1666, 1399, 2651, 2748, 283...  \n",
+       "6038  [1486, 1485, 3384, 3512, 3302, 3126, 3650, 330...  \n",
+       "6039  [2432, 2960, 1848, 2114, 2142, 3091, 3248, 317...  \n",
+       "\n",
+       "[6040 rows x 3 columns]"
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "train_events = bake_data(train_events)\n",
+    "\n",
+    "validation_events = bake_data(validation_events)\n",
+    "validation_gt = bake_data(validation_gt)\n",
+    "\n",
+    "test_events = bake_data(test_events)\n",
+    "test_gt = bake_data(test_gt)\n",
+    "\n",
+    "train_events"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "To ensure we don't have unknown users in ground truth, we join validation events and validation ground truth (also join test events and test ground truth correspondingly) by user ids to leave only the common ones.  "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def add_gt_to_events(events_df, gt_df):\n",
+    "    gt_to_join = gt_df[[\"user_id\", \"item_id\"]].rename(columns={\"item_id\": \"ground_truth\"})\n",
+    "\n",
+    "    events_df = events_df.merge(gt_to_join, on=\"user_id\", how=\"inner\")\n",
+    "    return events_df\n",
+    "\n",
+    "validation_events = add_gt_to_events(validation_events, validation_gt)\n",
+    "test_events = add_gt_to_events(test_events, test_gt)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pathlib import Path\n",
+    "\n",
+    "data_dir = Path(\"temp/data/\")\n",
+    "data_dir.mkdir(parents=True, exist_ok=True)\n",
+    "\n",
+    "TRAIN_PATH = data_dir / \"train.parquet\"\n",
+    "VAL_PATH = data_dir / \"val.parquet\"\n",
+    "PREDICT_PATH = data_dir / \"test.parquet\"\n",
+    "\n",
+    "ENCODER_PATH = data_dir / \"encoder\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_events.to_parquet(TRAIN_PATH)\n",
+    "validation_events.to_parquet(VAL_PATH)\n",
+    "test_events.to_parquet(PREDICT_PATH)\n",
+    "\n",
+    "encoder.save(ENCODER_PATH)"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Prepare to model training\n",
+    "### 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. It is required by the SasRec model to correctly create embeddings for every source column. Note that user_id does not required in `TensorSchema`.\n",
+    "\n",
+    "Note that the **padding value** is the next value (item_id) after the last one. **Cardinality** is the number of unique values ​​given the padding value."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from replay.data import FeatureHint, FeatureType\n",
+    "from replay.data.nn import TensorFeatureInfo, TensorSchema\n",
+    "\n",
+    "\n",
+    "EMBEDDING_DIM = 64\n",
+    "\n",
+    "encoder = encoder.load(ENCODER_PATH)\n",
+    "NUM_UNIQUE_ITEMS = len(encoder.mapping[\"item_id\"])\n",
+    "\n",
+    "tensor_schema = TensorSchema(\n",
+    "    [\n",
+    "        TensorFeatureInfo(\n",
+    "            name=\"item_id\",\n",
+    "            is_seq=True,\n",
+    "            padding_value=NUM_UNIQUE_ITEMS,\n",
+    "            cardinality=NUM_UNIQUE_ITEMS + 1,  # taking into account padding\n",
+    "            embedding_dim=EMBEDDING_DIM,\n",
+    "            feature_type=FeatureType.CATEGORICAL,\n",
+    "            feature_hint=FeatureHint.ITEM_ID,\n",
+    "        )\n",
+    "    ]\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Configure ParquetModule and transformation pipelines\n",
+    "\n",
+    "The `ParquetModule` class enables training of models on large datasets by reading data in batch-wise way. This class initialized with **paths to every data split, a metadata dict containing information about shape and padding value of every column and a dict of transforms**. `ParquetModule`'s  \"transform pipelines\" are stage-specific modules implementing additional preprocessing to be performed on batch level right before the forward pass.  \n",
+    "\n",
+    "For SasRec model, RePlay provides a function that generates a sequence of appropriate transforms for each data split named **make_default_sasrec_transforms**.\n",
+    "\n",
+    "Internally this function creates the following transforms:\n",
+    "1) Training:\n",
+    "    1. Create a target, which contains the shifted item sequence that represents the next item in the sequence (for the next item prediction task).\n",
+    "    2. Rename features to match it with expected format by the model during training.\n",
+    "    3. Unsqueeze target (*positive_labels*) and it's padding mask (*target_padding_mask*) for getting required shape of this tensors for loss computation.\n",
+    "    4. Group input features to be embed in expected format.\n",
+    "\n",
+    "2) Validation/Inference:\n",
+    "    1. Rename/group features to match it with expected format by the model during valdiation/inference.\n",
+    "\n",
+    "If a different set of transforms is required, you can create them yourself and submit them to the ParquetModule in the form of a dictionary where the key is the name of the split, and the value is the list of transforms. Available transforms are in the replay/nn/transforms/.\n",
+    "\n",
+    "**Note:** One of the transforms for the training data prepares the initial sequence for the task of Next Item Prediction so it shifts the sequence of items. For the final sequence length to be correct, you need to set shape of item_id in metadata as **model sequence length + shift**. Default shift value is 1."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import copy\n",
+    "\n",
+    "import torch\n",
+    "\n",
+    "from replay.data.nn import TensorSchema\n",
+    "from replay.nn.transform import GroupTransform, NextTokenTransform, RenameTransform, UnsqueezeTransform, FrequencyNegativeSamplingTransform"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def make_sasrec_transforms(\n",
+    "    tensor_schema: TensorSchema, query_column: str = \"query_id\", num_negative_samples: int = 128,\n",
+    ") -> dict[str, list[torch.nn.Module]]:\n",
+    "    item_column = tensor_schema.item_id_feature_name\n",
+    "    vocab_size = tensor_schema[item_column].cardinality\n",
+    "    train_transforms = [\n",
+    "        FrequencyNegativeSamplingTransform(vocab_size, num_negative_samples),\n",
+    "        NextTokenTransform(label_field=item_column, query_features=query_column, shift=1),\n",
+    "        RenameTransform(\n",
+    "            {\n",
+    "                query_column: \"query_id\",\n",
+    "                f\"{item_column}_mask\": \"padding_mask\",\n",
+    "                \"positive_labels_mask\": \"target_padding_mask\",\n",
+    "            }\n",
+    "        ),\n",
+    "        UnsqueezeTransform(\"target_padding_mask\", -1),\n",
+    "        UnsqueezeTransform(\"positive_labels\", -1),\n",
+    "        GroupTransform({\"feature_tensors\": [item_column]}),\n",
+    "    ]\n",
+    "\n",
+    "    val_transforms = [\n",
+    "        RenameTransform({query_column: \"query_id\", f\"{item_column}_mask\": \"padding_mask\"}),\n",
+    "        GroupTransform({\"feature_tensors\": [item_column]}),\n",
+    "    ]\n",
+    "    test_transforms = copy.deepcopy(val_transforms)\n",
+    "\n",
+    "    predict_transforms = copy.deepcopy(val_transforms)\n",
+    "\n",
+    "    transforms = {\n",
+    "        \"train\": train_transforms,\n",
+    "        \"validate\": val_transforms,\n",
+    "        \"test\": test_transforms,\n",
+    "        \"predict\": predict_transforms,\n",
+    "    }\n",
+    "\n",
+    "    return transforms\n",
+    "\n",
+    "transforms = make_sasrec_transforms(tensor_schema, query_column=\"user_id\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "MAX_SEQ_LEN = 50\n",
+    "\n",
+    "def create_meta(shape: int, gt_shape: Optional[int] = None):\n",
+    "    meta = {\n",
+    "        \"user_id\": {},\n",
+    "        \"item_id\": {\"shape\": shape, \"padding\": tensor_schema[\"item_id\"].padding_value},\n",
+    "    }\n",
+    "    if gt_shape is not None:\n",
+    "        meta.update({\"ground_truth\": {\"shape\": gt_shape, \"padding\": -1}})\n",
+    "\n",
+    "    return meta\n",
+    "\n",
+    "train_metadata = {\n",
+    "    \"train\": create_meta(shape=MAX_SEQ_LEN+1),\n",
+    "    \"validate\": create_meta(shape=MAX_SEQ_LEN, gt_shape=1),\n",
+    "}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from replay.data.nn import ParquetModule\n",
+    "\n",
+    "BATCH_SIZE = 32\n",
+    "\n",
+    "parquet_module = ParquetModule(\n",
+    "    train_path=TRAIN_PATH,\n",
+    "    validate_path=VAL_PATH,\n",
+    "    batch_size=BATCH_SIZE,\n",
+    "    metadata=train_metadata,\n",
+    "    transforms=transforms,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Train model\n",
+    "### Create SasRec model instance and run the training stage using lightning\n",
+    "We may now train the model using the Lightning trainer class. \n",
+    "\n",
+    "RePlay's implementation of SasRec is designed in a modular, **block-based approach**. Instead of passing configuration parameters to the constructor, SasRec is now built by providing fully initialized components that makes the model more flexible and easier to extend.\n",
+    "\n",
+    "#### Default Configuration\n",
+    "\n",
+    "Default SasRec model may be created quickly via method **from_params**. Default model instance has CE loss, original SasRec transformer layes, and embeddings are aggregated via sum."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from replay.nn.sequential import SasRec\n",
+    "from typing import Literal\n",
+    "def make_sasrec(\n",
+    "    schema: TensorSchema,\n",
+    "    embedding_dim: int = 192,\n",
+    "    num_heads: int = 4,\n",
+    "    num_blocks: int = 2,\n",
+    "    max_sequence_length: int = 50,\n",
+    "    dropout: float = 0.3,\n",
+    "    excluded_features: Optional[list[str]] = None,\n",
+    "    categorical_list_feature_aggregation_method: Literal[\"sum\", \"mean\", \"max\"] = \"sum\",\n",
+    ") -> SasRec:\n",
+    "    from replay.nn.sequential.sasrec import SasRecBody, SasRecTransformerLayer\n",
+    "    from replay.nn.agg import SumAggregator\n",
+    "    from replay.nn.embedding import SequenceEmbedding\n",
+    "    from replay.nn.loss import CE, CESampled\n",
+    "    from replay.nn.mask import DefaultAttentionMask\n",
+    "    from replay.nn.sequential.sasrec.agg import PositionAwareAggregator\n",
+    "    from replay.nn.sequential.sasrec.transformer import SasRecTransformerLayer\n",
+    "    excluded_features = [\n",
+    "        schema.query_id_feature_name,\n",
+    "        schema.timestamp_feature_name,\n",
+    "        *(excluded_features or []),\n",
+    "    ]\n",
+    "    excluded_features = list(set(excluded_features))\n",
+    "    body = SasRecBody(\n",
+    "        embedder=SequenceEmbedding(\n",
+    "            schema=schema,\n",
+    "            categorical_list_feature_aggregation_method=categorical_list_feature_aggregation_method,\n",
+    "            excluded_features=excluded_features,\n",
+    "        ),\n",
+    "        embedding_aggregator=PositionAwareAggregator(\n",
+    "            embedding_aggregator=SumAggregator(embedding_dim=embedding_dim),\n",
+    "            max_sequence_length=max_sequence_length,\n",
+    "            dropout=dropout,\n",
+    "        ),\n",
+    "        attn_mask_builder=DefaultAttentionMask(\n",
+    "            reference_feature_name=schema.item_id_feature_name,\n",
+    "            num_heads=num_heads,\n",
+    "        ),\n",
+    "        encoder=SasRecTransformerLayer(\n",
+    "            embedding_dim=embedding_dim,\n",
+    "            num_heads=num_heads,\n",
+    "            num_blocks=num_blocks,\n",
+    "            dropout=dropout,\n",
+    "            activation=\"relu\",\n",
+    "        ),\n",
+    "        output_normalization=torch.nn.LayerNorm(embedding_dim),\n",
+    "    )\n",
+    "    padding_idx = schema.item_id_features.item().padding_value\n",
+    "    return SasRec(\n",
+    "        body=body,\n",
+    "        loss=CESampled(padding_idx=padding_idx),\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "NUM_BLOCKS = 2\n",
+    "NUM_HEADS = 2\n",
+    "DROPOUT = 0.3\n",
+    "\n",
+    "sasrec = make_sasrec(\n",
+    "    schema=tensor_schema,\n",
+    "    embedding_dim=EMBEDDING_DIM,\n",
+    "    max_sequence_length=MAX_SEQ_LEN,\n",
+    "    num_heads=NUM_HEADS,\n",
+    "    num_blocks=NUM_BLOCKS,\n",
+    "    dropout=DROPOUT,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "A universal PyTorch Lightning module is provided. It can work with any NN model."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from replay.nn.lightning.optimizer import OptimizerFactory\n",
+    "from replay.nn.lightning.scheduler import LRSchedulerFactory\n",
+    "from replay.nn.lightning import LightningModule\n",
+    "\n",
+    "model = LightningModule(\n",
+    "    sasrec,\n",
+    "    optimizer_factory=OptimizerFactory(),\n",
+    "    lr_scheduler_factory=LRSchedulerFactory(),\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "To facilitate training, we add the following callbacks:\n",
+    "1) `ModelCheckpoint` - to save the best trained model based on its Recall metric. It's a default Lightning Callback.\n",
+    "1) `ComputeMetricsCallback` - to display a detailed validation metric matrix after each epoch. It's a custom RePlay callback for computing recsys metrics on validation and test stages.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "GPU available: True (cuda), used: True\n",
+      "TPU available: False, using: 0 TPU cores\n",
+      "You are using a CUDA device ('NVIDIA GeForce RTX 3060 Laptop GPU') 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 | Mode  | FLOPs\n",
+      "-------------------------------------------------\n",
+      "0 | model | SasRec | 291 K  | train | 0    \n",
+      "-------------------------------------------------\n",
+      "291 K     Trainable params\n",
+      "0         Non-trainable params\n",
+      "291 K     Total params\n",
+      "1.164     Total estimated model params size (MB)\n",
+      "39        Modules in train mode\n",
+      "0         Modules in eval mode\n",
+      "0         Total Flops\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "45904c0e26294a109f736787e2831727",
+       "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": "16229579b0fa4c3eadc8eeea029b139c",
+       "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": "b95d0d9187cd4511ba6122995385977c",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 0, global step 189: 'recall@10' reached 0.01921 (best 0.01921), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=0-step=189-v3.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.000662  0.004889  0.005938  0.003535\n",
+      "ndcg    0.000662  0.008170  0.012093  0.004836\n",
+      "recall  0.000662  0.019208  0.034940  0.008776\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "f396e2c9f6d1415181d9cf9ad119e483",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 1, global step 378: 'recall@10' reached 0.02335 (best 0.02335), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=1-step=378-v2.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k             1        10        20         5\n",
+      "map     0.00149  0.006367  0.007697  0.005037\n",
+      "ndcg    0.00149  0.010302  0.015430  0.007117\n",
+      "recall  0.00149  0.023348  0.044213  0.013578\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "5efd6eba4c714bc7a47bab506d96d083",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 2, global step 567: 'recall@10' reached 0.03809 (best 0.03809), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=2-step=567.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.004968  0.013050  0.015700  0.011130\n",
+      "ndcg    0.004968  0.018854  0.028575  0.014137\n",
+      "recall  0.004968  0.038086  0.076668  0.023348\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "14b9ffb84489441fbb65ecd2bffab125",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 3, global step 756: 'recall@10' reached 0.05448 (best 0.05448), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=3-step=756-v3.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.004968  0.015252  0.018691  0.011594\n",
+      "ndcg    0.004968  0.024178  0.036971  0.015113\n",
+      "recall  0.004968  0.054479  0.105647  0.025998\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "4e6aea05fd6e4fd796fdcc9c8432ffe9",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 4, global step 945: 'recall@10' reached 0.08130 (best 0.08130), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=4-step=945-v4.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.008942  0.025228  0.029602  0.020003\n",
+      "ndcg    0.008942  0.038097  0.054455  0.025390\n",
+      "recall  0.008942  0.081305  0.146879  0.041894\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "d79a151cb404477283713ea09c68f907",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 5, global step 1134: 'recall@10' reached 0.09571 (best 0.09571), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=5-step=1134-v1.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.010101  0.029038  0.033716  0.022943\n",
+      "ndcg    0.010101  0.044299  0.061645  0.029304\n",
+      "recall  0.010101  0.095711  0.164928  0.048849\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "ab01d88752a54a229fa429899a4e2552",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 6, global step 1323: 'recall@10' reached 0.10714 (best 0.10714), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=6-step=1323-v2.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.012254  0.033783  0.039462  0.027742\n",
+      "ndcg    0.012254  0.050622  0.071662  0.035493\n",
+      "recall  0.012254  0.107137  0.191091  0.059281\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "14a8ee885fe8477fa24673a3227e7e9c",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 7, global step 1512: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.010929  0.031616  0.037486  0.025432\n",
+      "ndcg    0.010929  0.048030  0.069781  0.032787\n",
+      "recall  0.010929  0.103163  0.189932  0.055473\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "6b9e020670534db1879e97dd7a566422",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 8, global step 1701: 'recall@10' reached 0.12353 (best 0.12353), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=8-step=1701-v1.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.012916  0.038435  0.044025  0.031021\n",
+      "ndcg    0.012916  0.058060  0.078630  0.040108\n",
+      "recall  0.012916  0.123530  0.205332  0.068058\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "bed94580dcbd4656951d1efda7ba8a48",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 9, global step 1890: 'recall@10' reached 0.12899 (best 0.12899), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=9-step=1890-v2.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.013082  0.040335  0.046561  0.032693\n",
+      "ndcg    0.013082  0.060782  0.083663  0.042103\n",
+      "recall  0.013082  0.128995  0.219904  0.070873\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "e05eafc1c8854f8787b5e39e2c47dd76",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 10, global step 2079: 'recall@10' reached 0.12966 (best 0.12966), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=10-step=2079.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.014738  0.041342  0.047593  0.033750\n",
+      "ndcg    0.014738  0.061675  0.084694  0.043076\n",
+      "recall  0.014738  0.129657  0.221229  0.071701\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "f9eee10fd1af441db81b11a8926c9d14",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 11, global step 2268: 'recall@10' reached 0.13396 (best 0.13396), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=11-step=2268.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.014075  0.040714  0.047416  0.032690\n",
+      "ndcg    0.014075  0.062142  0.086667  0.042514\n",
+      "recall  0.014075  0.133963  0.231164  0.072860\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "189ec86d37e0416583428616bbce96fc",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 12, global step 2457: 'recall@10' reached 0.14340 (best 0.14340), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=12-step=2457-v1.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016559  0.043821  0.050222  0.035111\n",
+      "ndcg    0.016559  0.066609  0.090205  0.045130\n",
+      "recall  0.016559  0.143401  0.237291  0.076172\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "c7b5622a9ed74116bfa6d8f35cba8025",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 13, global step 2646: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.014406  0.043623  0.050502  0.035351\n",
+      "ndcg    0.014406  0.066028  0.091376  0.045681\n",
+      "recall  0.014406  0.140917  0.241762  0.077331\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "72432093eb964a6183fc1dbd640a5450",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 14, global step 2835: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.015565  0.044625  0.051633  0.036168\n",
+      "ndcg    0.015565  0.067321  0.093215  0.046660\n",
+      "recall  0.015565  0.143236  0.246398  0.078987\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "4b1a70c1dcd4427da45471a9610505f7",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 15, global step 3024: 'recall@10' reached 0.15367 (best 0.15367), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=15-step=3024.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.015069  0.046444  0.053260  0.037230\n",
+      "ndcg    0.015069  0.071060  0.096111  0.048327\n",
+      "recall  0.015069  0.153668  0.253188  0.082464\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "dc15f2eb71fe47a28ca54c0bd9cf42f3",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 16, global step 3213: 'recall@10' reached 0.15946 (best 0.15946), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=16-step=3213.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.018546  0.050829  0.057741  0.041376\n",
+      "ndcg    0.018546  0.075784  0.101199  0.052429\n",
+      "recall  0.018546  0.159463  0.260474  0.086273\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "134594f99f4a4711a91047c3b4c1c1c0",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 17, global step 3402: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.018546  0.050609  0.057747  0.041723\n",
+      "ndcg    0.018546  0.074945  0.101494  0.053166\n",
+      "recall  0.018546  0.156152  0.262295  0.088260\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "952d8b7c4ecc4b8f845681f333e60e95",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 18, global step 3591: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.015565  0.048216  0.055341  0.039681\n",
+      "ndcg    0.015565  0.072890  0.099398  0.052029\n",
+      "recall  0.015565  0.154993  0.260970  0.090081\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "2e47543bfb714dcfb9ab93773989c6e3",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 19, global step 3780: 'recall@10' reached 0.16162 (best 0.16162), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=19-step=3780.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017553  0.050086  0.057070  0.040307\n",
+      "ndcg    0.017553  0.075715  0.101411  0.051752\n",
+      "recall  0.017553  0.161616  0.263785  0.086935\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "0186cb8fe6314be29e5c27ff61a4efb4",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 20, global step 3969: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016559  0.048569  0.055774  0.038668\n",
+      "ndcg    0.016559  0.074165  0.100638  0.049831\n",
+      "recall  0.016559  0.160126  0.265276  0.084120\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "ff8cbda5a5ad4131bc24bf0e21701493",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 21, global step 4158: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k             1        10        20         5\n",
+      "map     0.01689  0.050299  0.057754  0.041229\n",
+      "ndcg    0.01689  0.075501  0.103081  0.053162\n",
+      "recall  0.01689  0.159629  0.269581  0.089750\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "c7efc2e3954042dd8c399e271650103b",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 22, global step 4347: 'recall@10' reached 0.16360 (best 0.16360), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=22-step=4347-v1.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017718  0.051337  0.058542  0.041903\n",
+      "ndcg    0.017718  0.077220  0.103802  0.054198\n",
+      "recall  0.017718  0.163603  0.269415  0.092068\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "8fdf5663208e4b7189cb87eb315dd714",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 23, global step 4536: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017553  0.049966  0.057488  0.040051\n",
+      "ndcg    0.017553  0.075609  0.103244  0.051355\n",
+      "recall  0.017553  0.161616  0.271403  0.086107\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "7b221483a15448a186b77b3329be6ff6",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 24, global step 4725: 'recall@10' reached 0.16956 (best 0.16956), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=24-step=4725.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016559  0.050711  0.057813  0.040766\n",
+      "ndcg    0.016559  0.078002  0.104261  0.053405\n",
+      "recall  0.016559  0.169565  0.274218  0.092399\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "54f8fc2699aa419fa209a5ba25852555",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 25, global step 4914: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.018381  0.050549  0.058444  0.041033\n",
+      "ndcg    0.018381  0.075701  0.104654  0.052507\n",
+      "recall  0.018381  0.159795  0.274714  0.087763\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "399879ba18f7499d846c1f937a714ae6",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 26, global step 5103: 'recall@10' reached 0.17122 (best 0.17122), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=26-step=5103.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.022686  0.057648  0.064919  0.048236\n",
+      "ndcg    0.022686  0.083871  0.110727  0.060845\n",
+      "recall  0.022686  0.171220  0.278192  0.099520\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "7fd4a45308704782b46d80a418bbe30e",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 27, global step 5292: 'recall@10' reached 0.17188 (best 0.17188), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=27-step=5292.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.018049  0.053604  0.060891  0.043368\n",
+      "ndcg    0.018049  0.080857  0.107818  0.055870\n",
+      "recall  0.018049  0.171883  0.279351  0.094221\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "fa5f1113fac746b0ac41368de2eca760",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 28, global step 5481: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016228  0.052327  0.059730  0.042565\n",
+      "ndcg    0.016228  0.079886  0.107125  0.055913\n",
+      "recall  0.016228  0.171883  0.280179  0.097036\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "a34e71864baf44dab20ed841ae6e720e",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 29, global step 5670: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.015731  0.050238  0.057931  0.040238\n",
+      "ndcg    0.015731  0.077079  0.105369  0.052521\n",
+      "recall  0.015731  0.166915  0.279351  0.090247\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "f9cb3a2036894111b2ea320f1b666937",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 30, global step 5859: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016559  0.051478  0.059405  0.041271\n",
+      "ndcg    0.016559  0.078418  0.107606  0.053619\n",
+      "recall  0.016559  0.168405  0.284484  0.091571\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "dd15b49180b14ee9bdfc03e1d9f1440c",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 31, global step 6048: 'recall@10' reached 0.17519 (best 0.17519), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=31-step=6048.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016062  0.051997  0.059447  0.041563\n",
+      "ndcg    0.016062  0.080327  0.107738  0.054703\n",
+      "recall  0.016062  0.175195  0.284153  0.095214\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "2087d0ec66e245d68755fe0d4b38be3a",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 32, global step 6237: 'recall@10' reached 0.17834 (best 0.17834), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=32-step=6237.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016725  0.053134  0.060705  0.042214\n",
+      "ndcg    0.016725  0.081918  0.109782  0.055210\n",
+      "recall  0.016725  0.178341  0.289121  0.095214\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "fd91e9acc7094ed1b7baa2f1c257a447",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 33, global step 6426: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.019208  0.053199  0.061121  0.042943\n",
+      "ndcg    0.019208  0.080245  0.109425  0.055177\n",
+      "recall  0.019208  0.170889  0.286968  0.092896\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "b41a72bd49c541c88a0068c1eb3c8940",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 34, global step 6615: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016228  0.052733  0.060214  0.042899\n",
+      "ndcg    0.016228  0.080353  0.108022  0.056273\n",
+      "recall  0.016228  0.172380  0.282663  0.097367\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "7574b671f1144c35ad72d2bb149072a4",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 35, global step 6804: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k             1        10        20         5\n",
+      "map     0.01954  0.055007  0.062359  0.044064\n",
+      "ndcg    0.01954  0.083274  0.110562  0.056605\n",
+      "recall  0.01954  0.178010  0.286968  0.095214\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "b47e4be5150e4f8bacaf33c8b9f16b16",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 36, global step 6993: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.020036  0.056306  0.064246  0.046409\n",
+      "ndcg    0.020036  0.084078  0.113253  0.059811\n",
+      "recall  0.020036  0.176685  0.292598  0.101010\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "3dcb7bbf088b490ebd44a7421f69b641",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 37, global step 7182: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.019871  0.057126  0.065215  0.047538\n",
+      "ndcg    0.019871  0.084864  0.114584  0.061411\n",
+      "recall  0.019871  0.177016  0.295082  0.103991\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "96902e23ab3d49b1beca203004cbab7d",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 38, global step 7371: 'recall@10' reached 0.18132 (best 0.18132), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=38-step=7371.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k            1        10        20         5\n",
+      "map     0.0154  0.054742  0.062340  0.044704\n",
+      "ndcg    0.0154  0.084020  0.111946  0.059455\n",
+      "recall  0.0154  0.181321  0.292267  0.104819\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "fd9bd09986c240529c23d02d6dca3201",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 39, global step 7560: 'recall@10' reached 0.18231 (best 0.18231), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=39-step=7560.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017387  0.054622  0.062271  0.043890\n",
+      "ndcg    0.017387  0.083950  0.112185  0.057401\n",
+      "recall  0.017387  0.182315  0.294751  0.099023\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "0aba8305e71f4f1c957ec35d79fc0acb",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 40, global step 7749: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k            1        10        20         5\n",
+      "map     0.0154  0.052420  0.060409  0.041273\n",
+      "ndcg    0.0154  0.082044  0.111393  0.054579\n",
+      "recall  0.0154  0.181487  0.298063  0.095546\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "11610e6768a442a7aac578162ca26123",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 41, global step 7938: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017387  0.055064  0.063020  0.044130\n",
+      "ndcg    0.017387  0.084157  0.113600  0.057531\n",
+      "recall  0.017387  0.181321  0.298725  0.098692\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "afac63fbd8b54249bc6e0ebf214ae721",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 42, global step 8127: 'recall@10' reached 0.18364 (best 0.18364), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=42-step=8127.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.019374  0.056773  0.064316  0.046059\n",
+      "ndcg    0.019374  0.085958  0.113787  0.059584\n",
+      "recall  0.019374  0.183640  0.294420  0.101176\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "ef019fe931f742b3b811e73fd5404c52",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 43, global step 8316: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.018381  0.055013  0.062569  0.044560\n",
+      "ndcg    0.018381  0.083226  0.110987  0.057615\n",
+      "recall  0.018381  0.177513  0.287796  0.097698\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "6ec9a487bb04449db455a3b620156a77",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 44, global step 8505: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.019871  0.056811  0.064680  0.046062\n",
+      "ndcg    0.019871  0.085587  0.114373  0.059273\n",
+      "recall  0.019871  0.181818  0.295910  0.099851\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "5b06405c1da54f31823138198940eda6",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 45, global step 8694: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016062  0.053028  0.060992  0.042634\n",
+      "ndcg    0.016062  0.081447  0.110764  0.055989\n",
+      "recall  0.016062  0.176354  0.292929  0.097036\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "9f9ccc6e42734a7eab2c3b6097f56427",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 46, global step 8883: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.015565  0.054815  0.062818  0.044467\n",
+      "ndcg    0.015565  0.084146  0.113348  0.058769\n",
+      "recall  0.015565  0.181818  0.297400  0.102666\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "8f0b2c3bfb9b4dfab3b788b72ba05835",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 47, global step 9072: 'recall@10' reached 0.18662 (best 0.18662), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=47-step=9072.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.018877  0.058107  0.065509  0.047707\n",
+      "ndcg    0.018877  0.087807  0.115050  0.062437\n",
+      "recall  0.018877  0.186620  0.294916  0.107799\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "c5a662fbe23e433bb394c0ea3fe13f0c",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 48, global step 9261: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016725  0.053884  0.061396  0.043225\n",
+      "ndcg    0.016725  0.083373  0.110910  0.057023\n",
+      "recall  0.016725  0.182149  0.291439  0.099520\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "a05aba0cde6e40b584a4578dfe033a85",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 49, global step 9450: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.019705  0.057117  0.064926  0.046454\n",
+      "ndcg    0.019705  0.086163  0.114734  0.060090\n",
+      "recall  0.019705  0.183143  0.296407  0.102004\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "2efa59edb69541259ad8a15ce10f5b1e",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 50, global step 9639: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.022189  0.059775  0.067521  0.049652\n",
+      "ndcg    0.022189  0.088257  0.116693  0.063292\n",
+      "recall  0.022189  0.183308  0.296241  0.105150\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "f85b90dc5a9b49da8739cf0df654d539",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 51, global step 9828: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.020864  0.058991  0.066927  0.049070\n",
+      "ndcg    0.020864  0.087316  0.116491  0.062868\n",
+      "recall  0.020864  0.181653  0.297566  0.105150\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "474196c5121f4c1eb1eed65c9f975492",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 52, global step 10017: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.019208  0.054836  0.063047  0.044936\n",
+      "ndcg    0.019208  0.082000  0.111981  0.057885\n",
+      "recall  0.019208  0.172545  0.291273  0.097698\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "3c85df11f2ea4dbbae2be2ff3327b1c0",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 53, global step 10206: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.015234  0.051760  0.059551  0.041094\n",
+      "ndcg    0.015234  0.080313  0.109110  0.054204\n",
+      "recall  0.015234  0.175857  0.290611  0.094552\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "6935921a319d4f9dbb3322c7bfade032",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 54, global step 10395: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016393  0.053267  0.061428  0.042606\n",
+      "ndcg    0.016393  0.082624  0.112685  0.056081\n",
+      "recall  0.016393  0.181156  0.300712  0.097533\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "a20b778f33914f6a859ac8c9af047756",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 55, global step 10584: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.018546  0.056375  0.064499  0.045653\n",
+      "ndcg    0.018546  0.085650  0.115357  0.059239\n",
+      "recall  0.018546  0.183640  0.301374  0.101010\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "74efdae045ca49fdbdac049832ffe873",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 56, global step 10773: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k             1        10        20         5\n",
+      "map     0.01689  0.054781  0.063420  0.044130\n",
+      "ndcg    0.01689  0.083637  0.115169  0.057295\n",
+      "recall  0.01689  0.180162  0.305017  0.097533\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "c9fed52a17f74d13aaab5606b8c22b21",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 57, global step 10962: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.014903  0.052826  0.061114  0.042604\n",
+      "ndcg    0.014903  0.081359  0.111826  0.056267\n",
+      "recall  0.014903  0.176519  0.297566  0.098195\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "8385fa92b9fa4fbbad147fe2276ecd05",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 58, global step 11151: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.018546  0.057551  0.065000  0.046760\n",
+      "ndcg    0.018546  0.087028  0.114631  0.060664\n",
+      "recall  0.018546  0.185296  0.295413  0.103328\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "1e41ff6e0789411d94ea038d498d9e1f",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 59, global step 11340: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016725  0.054240  0.062219  0.043672\n",
+      "ndcg    0.016725  0.083596  0.112902  0.057523\n",
+      "recall  0.016725  0.181818  0.298228  0.100182\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "19d7ff7655534dcb83a44980e6db8490",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 60, global step 11529: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.018877  0.056656  0.064225  0.045791\n",
+      "ndcg    0.018877  0.085705  0.113617  0.059086\n",
+      "recall  0.018877  0.182812  0.293923  0.099851\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "d138b9ee480745289b9fb32d37479313",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 61, global step 11718: 'recall@10' reached 0.18894 (best 0.18894), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=61-step=11718.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.020202  0.058858  0.066230  0.047933\n",
+      "ndcg    0.020202  0.088808  0.116065  0.061998\n",
+      "recall  0.020202  0.188939  0.297566  0.105315\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "ab3628938d5549658e2072ffac8c8cd3",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 62, global step 11907: 'recall@10' reached 0.19126 (best 0.19126), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=62-step=11907.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.018546  0.059156  0.066260  0.047778\n",
+      "ndcg    0.018546  0.089650  0.115582  0.061979\n",
+      "recall  0.018546  0.191257  0.293923  0.105481\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "5fb68ad004cd45dcaed411451fccdd7a",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 63, global step 12096: 'recall@10' reached 0.19424 (best 0.19424), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=63-step=12096.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017056  0.058524  0.066103  0.047138\n",
+      "ndcg    0.017056  0.089873  0.117581  0.062228\n",
+      "recall  0.017056  0.194237  0.304024  0.108627\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "71e338441644421ea8b6cc32c086b71b",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 64, global step 12285: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.018381  0.057568  0.065550  0.046550\n",
+      "ndcg    0.018381  0.087703  0.117039  0.060806\n",
+      "recall  0.018381  0.188276  0.304852  0.104653\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "0381d441937f40219cbedef04f8235bc",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 65, global step 12474: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k             1        10        20         5\n",
+      "map     0.01689  0.057109  0.065127  0.046473\n",
+      "ndcg    0.01689  0.087242  0.116738  0.061132\n",
+      "recall  0.01689  0.187614  0.304852  0.106143\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "5120fc8bdf784305890f9db8a9c6fa80",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 66, global step 12663: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.018049  0.057167  0.065304  0.046158\n",
+      "ndcg    0.018049  0.087499  0.117282  0.060454\n",
+      "recall  0.018049  0.188939  0.307004  0.104488\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "189c0b78fcc24681b130c8c53c0aa0cb",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 67, global step 12852: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017387  0.056557  0.064928  0.046536\n",
+      "ndcg    0.017387  0.085917  0.116763  0.061370\n",
+      "recall  0.017387  0.183474  0.306177  0.106971\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "e15930826aa840c4bc6e03a49be0848d",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 68, global step 13041: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.018712  0.058088  0.066686  0.047259\n",
+      "ndcg    0.018712  0.088043  0.119697  0.061554\n",
+      "recall  0.018712  0.187945  0.313794  0.105481\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "0c8d861c579a4e43ac269be8985f0af3",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 69, global step 13230: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.018049  0.058269  0.066586  0.047334\n",
+      "ndcg    0.018049  0.088878  0.119519  0.062153\n",
+      "recall  0.018049  0.190926  0.312800  0.107799\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "ea017f9572c046a5ab02de1d9196badc",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 70, global step 13419: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016725  0.057414  0.065719  0.046779\n",
+      "ndcg    0.016725  0.088837  0.119433  0.062701\n",
+      "recall  0.016725  0.193575  0.315284  0.111939\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "946138e70b474424ac7b15fabbcb1b15",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 71, global step 13608: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017884  0.056873  0.064915  0.045822\n",
+      "ndcg    0.017884  0.087012  0.116591  0.060047\n",
+      "recall  0.017884  0.187614  0.305183  0.103825\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "37b783951bcd4fe49afa8ad617c72fdb",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 72, global step 13797: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.019043  0.057263  0.065566  0.046616\n",
+      "ndcg    0.019043  0.086345  0.116833  0.060380\n",
+      "recall  0.019043  0.183308  0.304355  0.102666\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "89240546fba2486eace45af1403ffc4e",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 73, global step 13986: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.018215  0.057662  0.066092  0.047373\n",
+      "ndcg    0.018215  0.086925  0.117706  0.061920\n",
+      "recall  0.018215  0.184136  0.306011  0.106640\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "34f2d52527524ccb9e4caafdbf45a8a7",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 74, global step 14175: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.018215  0.058467  0.066697  0.047516\n",
+      "ndcg    0.018215  0.089375  0.119492  0.062387\n",
+      "recall  0.018215  0.192582  0.311972  0.108130\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "9eb272372ba74074979223fb36988b4b",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 75, global step 14364: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017387  0.057744  0.065851  0.046274\n",
+      "ndcg    0.017387  0.088852  0.118646  0.060779\n",
+      "recall  0.017387  0.192747  0.311144  0.105315\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "e7fb52f34f9849fe80f317eff347ac59",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 76, global step 14553: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.015897  0.055235  0.063557  0.044367\n",
+      "ndcg    0.015897  0.085566  0.116100  0.058943\n",
+      "recall  0.015897  0.186786  0.307998  0.103825\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "e3905ba364ce4b79bec284d9f3ba399a",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 77, global step 14742: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016725  0.057168  0.065456  0.046627\n",
+      "ndcg    0.016725  0.087777  0.118245  0.061825\n",
+      "recall  0.016725  0.189767  0.310813  0.108627\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "5d9d0a90829a47be8a39188f73fcb299",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 78, global step 14931: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016725  0.054820  0.063159  0.043650\n",
+      "ndcg    0.016725  0.084686  0.115602  0.057321\n",
+      "recall  0.016725  0.184633  0.307998  0.099354\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "d1f3a1f2e246405eaa8a9db0a95a511a",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 79, global step 15120: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k            1        10        20         5\n",
+      "map     0.0154  0.056311  0.064352  0.045432\n",
+      "ndcg    0.0154  0.087242  0.116821  0.060765\n",
+      "recall  0.0154  0.190098  0.307667  0.107965\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "9fa3134251a44d71b04280c72af0ddc5",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 80, global step 15309: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017056  0.056116  0.064667  0.045543\n",
+      "ndcg    0.017056  0.085447  0.116845  0.059621\n",
+      "recall  0.017056  0.183143  0.307832  0.102832\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "837c34003f17463cb66971c10ddce1e6",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 81, global step 15498: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017221  0.055801  0.064094  0.044936\n",
+      "ndcg    0.017221  0.085296  0.115712  0.058833\n",
+      "recall  0.017221  0.183640  0.304355  0.101507\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "bfab9a3923474fd4930e468a5144466f",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 82, global step 15687: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016062  0.054903  0.063218  0.044500\n",
+      "ndcg    0.016062  0.083626  0.114177  0.058224\n",
+      "recall  0.016062  0.179169  0.300546  0.100182\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "ce807241730e4067b08757e709957866",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 83, global step 15876: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.015897  0.057228  0.065394  0.046249\n",
+      "ndcg    0.015897  0.088010  0.117992  0.060994\n",
+      "recall  0.015897  0.190594  0.309654  0.106143\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "da634848c3d24e2e94fc916a025e7bbe",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 84, global step 16065: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016725  0.056139  0.064478  0.044845\n",
+      "ndcg    0.016725  0.086567  0.117230  0.059108\n",
+      "recall  0.016725  0.188111  0.309985  0.102997\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "1a6a134c1f7a47d7845ac0b7070b2c7b",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 85, global step 16254: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017056  0.056356  0.064636  0.046164\n",
+      "ndcg    0.017056  0.085823  0.116328  0.060727\n",
+      "recall  0.017056  0.183971  0.305349  0.105481\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "8c0c921e05674658b3bb48f8f101abc9",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 86, global step 16443: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.018877  0.057997  0.065875  0.046260\n",
+      "ndcg    0.018877  0.088555  0.117546  0.059837\n",
+      "recall  0.018877  0.190926  0.306177  0.101507\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "19135c0fba704ad0ae7838c12ed0fed5",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 87, global step 16632: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017221  0.056635  0.064463  0.046288\n",
+      "ndcg    0.017221  0.087059  0.116068  0.061620\n",
+      "recall  0.017221  0.188442  0.304189  0.108958\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "5bf0b56c43d24c5d930359679b722e04",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 88, global step 16821: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016725  0.057663  0.065748  0.046454\n",
+      "ndcg    0.016725  0.088640  0.118382  0.061238\n",
+      "recall  0.016725  0.191919  0.310151  0.106640\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "0223aa1f4bc645d9aab2100e1c4f7384",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 89, global step 17010: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016062  0.057177  0.065153  0.045805\n",
+      "ndcg    0.016062  0.088354  0.117906  0.060564\n",
+      "recall  0.016062  0.192250  0.310151  0.105812\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "28bf241cbefb422281ec0406264c576a",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 90, global step 17199: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k             1        10        20         5\n",
+      "map     0.01689  0.057756  0.066077  0.046741\n",
+      "ndcg    0.01689  0.088707  0.119294  0.061811\n",
+      "recall  0.01689  0.191754  0.313297  0.108130\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "d12545edc74045718845f0fd8bff93cd",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 91, global step 17388: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016228  0.057542  0.065642  0.046476\n",
+      "ndcg    0.016228  0.088482  0.118328  0.061535\n",
+      "recall  0.016228  0.191422  0.310151  0.107799\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "b6753a36ce234372b66a2d48f49f1035",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 92, global step 17577: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016559  0.056973  0.064812  0.045885\n",
+      "ndcg    0.016559  0.087680  0.116507  0.060459\n",
+      "recall  0.016559  0.190098  0.304686  0.105150\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "dbdee467d75c4e36ad40a23e875994b6",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 93, global step 17766: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016559  0.058180  0.066174  0.047397\n",
+      "ndcg    0.016559  0.089539  0.118949  0.063218\n",
+      "recall  0.016559  0.193741  0.310647  0.111939\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "2399141e31cf415f9ee3a4ad926ce1dd",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 94, global step 17955: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.015731  0.056793  0.065030  0.045725\n",
+      "ndcg    0.015731  0.087689  0.117979  0.060803\n",
+      "recall  0.015731  0.190429  0.310813  0.107137\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "39f19e1dc4174c0bb397e17bd243bdb6",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 95, global step 18144: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.018546  0.057964  0.066046  0.047199\n",
+      "ndcg    0.018546  0.088238  0.118000  0.061845\n",
+      "recall  0.018546  0.189270  0.307667  0.106971\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "1c38e55824ff41bbaf7520ddfa469142",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 96, global step 18333: 'recall@10' reached 0.19622 (best 0.19622), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=96-step=18333.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k             1        10        20         5\n",
+      "map     0.01954  0.060201  0.068178  0.048722\n",
+      "ndcg    0.01954  0.091547  0.121043  0.063432\n",
+      "recall  0.01954  0.196225  0.313794  0.108627\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "04808d64d1004f1aa9ce59a5e66148f4",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 97, global step 18522: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017221  0.057618  0.065656  0.046542\n",
+      "ndcg    0.017221  0.088212  0.118036  0.061226\n",
+      "recall  0.017221  0.190098  0.309157  0.106309\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "57c0e23d1c8945d4b7f02d50241b88fb",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 98, global step 18711: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.014903  0.055248  0.063857  0.043799\n",
+      "ndcg    0.014903  0.086091  0.117646  0.058227\n",
+      "recall  0.014903  0.188939  0.314125  0.102500\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "9da10aefb0d24ca7be02704f6c74a4c6",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 99, global step 18900: 'recall@10' was not in top 1\n",
+      "`Trainer.fit` stopped: `max_epochs=100` reached.\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k             1        10        20         5\n",
+      "map     0.01689  0.057811  0.065843  0.047138\n",
+      "ndcg    0.01689  0.088620  0.118022  0.062352\n",
+      "recall  0.01689  0.191257  0.307832  0.109124\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "from lightning.pytorch.callbacks import ModelCheckpoint\n",
+    "from lightning.pytorch.loggers import CSVLogger\n",
+    "\n",
+    "from replay.nn.lightning.callback import ComputeMetricsCallback\n",
+    "\n",
+    "\n",
+    "checkpoint_callback = ModelCheckpoint(\n",
+    "    dirpath=\"sasrec/checkpoints\",\n",
+    "    save_top_k=1,\n",
+    "    verbose=True,\n",
+    "    monitor=\"recall@10\",\n",
+    "    mode=\"max\",\n",
+    ")\n",
+    "\n",
+    "validation_metrics_callback = ComputeMetricsCallback(\n",
+    "    metrics=[\"map\", \"ndcg\", \"recall\"],\n",
+    "    ks=[1, 5, 10, 20],\n",
+    "    item_count=NUM_UNIQUE_ITEMS,\n",
+    ")\n",
+    "\n",
+    "csv_logger = CSVLogger(save_dir=\"sasrec/logs/train\", name=\"SasRec-example\")\n",
+    "\n",
+    "trainer = L.Trainer(\n",
+    "    max_epochs=100,\n",
+    "    callbacks=[checkpoint_callback, validation_metrics_callback],\n",
+    "    logger=csv_logger,\n",
+    ")\n",
+    "\n",
+    "trainer.fit(model, datamodule=parquet_module)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAhQAAAGnCAYAAADvzM6jAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjgsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvwVt1zgAAAAlwSFlzAAASdAAAEnQB3mYfeAAALWJJREFUeJzt3Xt8VPWd//H3hEAmZBoEIWS4La4BpYE0UDOUPBCWooiggpcUYZHCClRDpaB4iaxr0YpRQLeWUlC7hXLxAnVBagW1jwR5CCTFgtUVGyKLEEwMFCEkJAGS8/uDX2YZJpfJfJOZyczr+XjkEfI93/Od7yczZ3jnnDPn2CzLsgQAAGAgKtgTAAAAbR+BAgAAGCNQAAAAYwQKAABgjEABAACMESgAAIAxAgUAADBGoAAAAMYIFAAAwBiBAgAAGIsO9gRay6lTp7Rjxw717t1bMTExwZ4OAABtRnV1tY4ePaqRI0fqiiuu8GmdsA0UO3bs0MSJE4M9DQAA2qzNmzdrwoQJPvUN20DRu3dvSRd/GUlJSUGeTcsqLy9Xfn6+XC6XHA5HsKcTENRMzeEo0uqVqLmt1FxYWKiJEye6/y/1RdgGirrDHElJSUpOTg7ybFpWWVmZSkpKNGDAAMXHxwd7OgFBzdQcjiKtXoma21rNzTllgJMyAQCAMQIFAAAwRqAAAADGCBQAAMAYgQIAABgjUAAAAGMECgAAYMyvQPGXv/xFP/3pT5WcnKy4uDj16dNHP/rRj1RQUODRb/r06bLZbF5f1157rdeYtbW1ev7553XVVVfJbrcrJSVFr732mn9VAQCAgPLrwlbPPfecPvroI2VkZCglJUUlJSVavny5hgwZoj179mjgwIHuvjExMXr11Vc91u/UqZPXmAsXLlR2drZmzZqltLQ0bdmyRVOmTJHNZtPdd9/tzzQBAECA+BUoHnzwQW3YsEEdOnRwt02aNEmDBg1Sdna21q1b938PEB2tqVOnNjresWPHtGzZMs2ZM0fLly+XJM2cOVMjR47Uww8/rIyMDLVr186fqQIAgADw65BHenq6R5iQpH79+ik5OVkHDhzw6l9TU6OysrIGx9uyZYvOnz+vzMxMd5vNZtP999+voqIi7d69259pAgCAAGmxkzIty9I333yjrl27erSfPXtW8fHx6tSpk7p06aI5c+aovLzco8++ffsUFxenAQMGeLS7XC73cgAAELpa7OZg69ev17Fjx/TUU0+525xOpx555BENGTJEtbW12rZtm1asWKFPPvlEubm5io6++PDFxcXq3r27bDabx5hOp1OS9PXXXzf62KWlpTp+/LhHW2FhoaSLd3lrbO9IW1RRUeHxPRJQc2SItJojrV6JmtuKy//w90WLBIovvvhCc+bM0bBhw/TjH//Y3f7ss8969Lv77rvVv39/LVy4UJs2bXKfbFlZWVnvHc3sdrt7eWNWrFihRYsW1bssPz9fJSUlzaqnrcjPzw/2FAKOmiNDpNUcafVK1Bzqjhw50ux1jANFSUmJxo8fr06dOmnTpk1Nnjw5f/58PfHEE/rggw/cgSI2NlbV1dVefauqqtzLG5OZmamMjAyPtrp7ubtcLq9DKW1dRUWF8vPz5XK5FBcXF+zpBAQ1h2fNY1780OPnrnZLs66t1StfROlElU3vzR8RpJkFRiQ8x5ej5rZRc33nQzbFKFCcPn1aN998s06dOqWdO3eqR48eTa4TGxurK6+8UidPnnS3OZ1O5eTkyLIsj8MexcXFktTkuAkJCUpISKh3mcPhaHP3n/dVXFxc2NbWEGoOLyWVtnrbT1TZVFJpC9u6LxfOz3FDqDm0ORyOZq/j90mZVVVVuvXWW1VQUKA//vGP+u53v+vTemfOnNGJEyfUrVs3d1tqaqrOnj3rlYjy8vLcywEAQOjyK1DU1NRo0qRJ2r17tzZu3Khhw4Z59amqqtKZM2e82p9++mlZlqWxY8e62yZMmKD27dtrxYoV7jbLsrRy5Ur17NlT6enp/kwTAAAEiF+HPB566CG9/fbbuvXWW3Xy5EmPC1lJ0tSpU1VSUqLBgwdr8uTJ7kttb9++XX/60580duxYTZgwwd2/V69emjdvnpYsWaLz588rLS1Nmzdv1s6dO7V+/XouagUAQIjzK1Ds379fkrR161Zt3brVa/nUqVN1xRVX6JZbbtH777+vNWvWqKamRklJSVq8eLEWLFigqCjPnSPZ2dnq3LmzVq1apdWrV6tfv35at26dpkyZ4s8UAQBAAPkVKHJzc5vsc8UVV2jt2rU+jxkVFaWsrCxlZWX5MyUAABBE3L4cAAAYI1AAAABjBAoAAGCMQAEAAIwRKAAAgDECBQAAMEagAAAAxggUAADAGIECAAAYI1AAAABjfl16GwBCRd/H3ml0+eHs8QGaCRDZ2EMBAACMESgAAIAxAgUAADBGoAAAAMYIFAAAwBiBAgAAGCNQAAAAYwQKAABgjEABAACMESgAAIAxAgUAADBGoAAAAMYIFAAAwBiBAgAAGCNQAAAAYwQKAABgjEABAACMESgAAIAxAgUAADBGoAAAAMYIFAAAwBiBAgAAGCNQAAAAYwQKAABgjEABAACMESgAAIAxAgUAADBGoAAAAMYIFAAAwBiBAgAAGCNQAAAAYwQKAABgjEABAACMESgAAIAxAgUAADBGoAAAAMYIFAAAwBiBAgAAGCNQAAAAYwQKAABgjEABAACMESgAAIAxAgUAADBGoAAAAMYIFAAAwBiBAgAAGPMrUPzlL3/RT3/6UyUnJysuLk59+vTRj370IxUUFHj1PXDggMaOHSuHw6EuXbronnvu0fHjx7361dbW6vnnn9dVV10lu92ulJQUvfbaa/5MDwAABFi0Pys999xz+uijj5SRkaGUlBSVlJRo+fLlGjJkiPbs2aOBAwdKkoqKijRixAh16tRJixcvVnl5uZYuXapPP/1U+fn56tChg3vMhQsXKjs7W7NmzVJaWpq2bNmiKVOmyGaz6e67726ZagEAQKvwK1A8+OCD2rBhg0cgmDRpkgYNGqTs7GytW7dOkrR48WJVVFTo448/Vp8+fSRJLpdLN954o1avXq3Zs2dLko4dO6Zly5Zpzpw5Wr58uSRp5syZGjlypB5++GFlZGSoXbt2RoUCAIDW49chj/T0dI8wIUn9+vVTcnKyDhw44G77wx/+oFtuucUdJiTphhtuUP/+/fXmm2+627Zs2aLz588rMzPT3Waz2XT//ferqKhIu3fv9meaAAAgQFrspEzLsvTNN9+oa9euki7udSgtLdV1113n1dflcmnfvn3un/ft26e4uDgNGDDAq1/dcgAAELr8OuRRn/Xr1+vYsWN66qmnJEnFxcWSJKfT6dXX6XTq5MmTqq6uVkxMjIqLi9W9e3fZbDavfpL09ddfN/rYpaWlXid6FhYWSpLKy8tVVlbmX1EhqqKiwuN7JKDm8JQYa3n83NVueXz3Zdu9fIzLhfL2HwnP8eWouW0oLy9v9jotEii++OILzZkzR8OGDdOPf/xjSVJlZaUkKSYmxqu/3W5394mJiXF/b6xfY1asWKFFixbVuyw/P18lJSW+F9OG5OfnB3sKAUfN4SUrtf72WdfWSpJycnL8HqOOL2MEWzg/xw2h5tB25MiRZq9jHChKSko0fvx4derUSZs2bXKfPBkbGytJqq6u9lqnqqrKo09sbKxP/RqSmZmpjIwMj7bCwkJNnDhRLpfL61BKW1dRUaH8/Hy5XC7FxcUFezoBQc3hWfOYFz/0+Lmr3dKsa2v1yhdROlFl03vzRzR7jMv5MkawRMJzfDlqbhs1X3o+pK+MAsXp06d1880369SpU9q5c6d69OjhXlZ3uKLu0MeliouL1aVLF/deCafTqZycHFmW5XHYo27dS8etT0JCghISEupd5nA4FB8f37zC2oi4uLiwra0h1BwYfR97p8k+h7PHGz9OSaWt3vYTVTaVVNp8qruhMeq0hdcLr+vI0JZqdjgczV7H75Myq6qqdOutt6qgoEB//OMf9d3vftdjec+ePdWtWzft3bvXa938/Hylpqa6f05NTdXZs2e9ElFeXp57OQAACF1+7aGoqanRpEmTtHv3bm3ZskXDhg2rt9+dd96pNWvW6OjRo+rdu7ck6c9//rMKCgo0f/58d78JEyZo/vz5WrFihfs6FJZlaeXKlerZs6fS09P9mSYA+KSpPTItsTcGCHd+BYqHHnpIb7/9tm699VadPHnSfSGrOlOnTpUkPf7449q4caNGjRqln/3sZyovL9eSJUs0aNAgzZgxw92/V69emjdvnpYsWaLz588rLS1Nmzdv1s6dO7V+/XouagXAb74cvgFgzq9AsX//fknS1q1btXXrVq/ldYGid+/e2rFjhx588EE99thj6tChg8aPH69ly5Z5faojOztbnTt31qpVq7R69Wr169dP69at05QpU/yZIgAACCC/AkVubq7PfZOTk7V9+/Ym+0VFRSkrK0tZWVn+TAkAAAQRty8HAADGCBQAAMAYgQIAABgjUAAAAGMtdnMwAIhkXMsCkY49FAAAwBh7KAD4hb/IAVyKPRQAAMAYgQIAABgjUAAAAGMECgAAYIxAAQAAjBEoAACAMQIFAAAwRqAAAADGCBQAAMAYV8oEELKauhongNDBHgoAAGCMQAEAAIwRKAAAgDECBQAAMEagAAAAxggUAADAGIECAAAYI1AAAABjBAoAAGCMQAEAAIwRKAAAgDECBQAAMEagAAAAxggUAADAGIECAAAYI1AAAABjBAoAAGCMQAEAAIwRKAAAgDECBQAAMEagAAAAxqKDPQEALa/vY+80uvxw9vgAzQRApGAPBQAAMEagAAAAxggUAADAGIECAAAYI1AAAABjBAoAAGCMQAEAAIwRKAAAgDECBQAAMEagAAAAxggUAADAGIECAAAYI1AAAABjBAoAAGCMQAEAAIwRKAAAgDECBQAAMEagAAAAxqL9XbG8vFxLlixRXl6e8vPz9e233+p3v/udpk+f7tFv+vTpWrNmjdf611xzjb744guPttraWi1dulS/+c1vVFxcrP79+ysrK0uTJ0/2d5oAgqTvY+8EewoAAsjvQHHixAk99dRT6tOnj773ve8pNze3wb4xMTF69dVXPdo6derk1W/hwoXKzs7WrFmzlJaWpi1btmjKlCmy2Wy6++67/Z0qAABoZX4HCqfTqeLiYiUmJmrv3r1KS0tr+EGiozV16tRGxzt27JiWLVumOXPmaPny5ZKkmTNnauTIkXr44YeVkZGhdu3a+TtdAADQivw+hyImJkaJiYk+96+pqVFZWVmDy7ds2aLz588rMzPT3Waz2XT//ferqKhIu3fv9neqAACglfm9h6I5zp49q/j4eJ09e1adO3fW5MmT9dxzz8nhcLj77Nu3T3FxcRowYIDHui6Xy718+PDh9Y5fWlqq48ePe7QVFhZKuniuR2NBpi2qqKjw+B4JqLl5EmOtRpc3tU00tX5r6Wq3PL6HCl/eQ/z5nfO6jgxtseby8vJmr9PqgcLpdOqRRx7RkCFDVFtbq23btmnFihX65JNPlJubq+joi1MoLi5W9+7dZbPZvNaXpK+//rrBx1ixYoUWLVpU77L8/HyVlJS0UDWhJT8/P9hTCDhq9k1WauPLc3JyjNZvbbOurQ3uBC7T1O9LMvud87qODG2p5iNHjjR7nVYPFM8++6zHz3fffbf69++vhQsXatOmTe6TLSsrKxUTE+O1vt1udy9vSGZmpjIyMjzaCgsLNXHiRLlcLq+9Hm1dRUWF8vPz5XK5FBcXF+zpBAQ1N6/mMS9+2Ojy9+aPMFq/tXS1W5p1ba1e+SJKJ6psTa8QIE39viT/fue8rqk5VB04cKDZ6wTkkMfl5s+fryeeeEIffPCBO1DExsaqurraq29VVZV7eUMSEhKUkJBQ7zKHw6H4+PgWmHXoiYuLC9vaGkLNvimpbPw/46bGa2r91naiyhb0OVzKl9+/ye+c13VkaEs1X3pKgq+CcmGr2NhYXXnllTp58qS7zel0qqSkRJbleRyyuLhYktSjR4+AzhEAAPguKIHizJkzOnHihLp16+ZuS01N1dmzZ712s+Tl5bmXAwCA0NSqgaKqqkpnzpzxan/66adlWZbGjh3rbpswYYLat2+vFStWuNssy9LKlSvVs2dPpaent+ZUAQCAAaNzKJYvX65Tp065P4GxdetWFRUVSZIeeOABffvttxo8eLAmT56sa6+9VpK0fft2/elPf9LYsWM1YcIE91i9evXSvHnztGTJEp0/f15paWnavHmzdu7cqfXr13NRKwAAQphRoFi6dKm++uor989vvfWW3nrrLUnS1KlTdcUVV+iWW27R+++/rzVr1qimpkZJSUlavHixFixYoKgozx0k2dnZ6ty5s1atWqXVq1erX79+WrdunaZMmWIyTQAA0MqMAsXhw4eb7LN27Vqfx4uKilJWVpaysrIMZgUAAAItKB8bBRBc3AkUQEsLyqc8AABAeCFQAAAAYwQKAABgjEABAACMESgAAIAxAgUAADBGoAAAAMYIFAAAwBiBAgAAGCNQAAAAYwQKAABgjEABAACMESgAAIAxAgUAADBGoAAAAMYIFAAAwBiBAgAAGCNQAAAAYwQKAABgjEABAACMESgAAIAxAgUAADBGoAAAAMaigz0BAIgEfR97x6stMdZSVqo05sUPtefJWwI/KaAFsYcCAAAYI1AAAABjBAoAAGCMQAEAAIwRKAAAgDECBQAAMMbHRoEQU/fxwks/UlhSaXMvP5w9PkgzA4CGESgAoAn1XUMCgCcOeQAAAGMECgAAYIxDHgAQAnw5rML5Mwhl7KEAAADGCBQAAMAYgQIAABgjUAAAAGMECgAAYIxAAQAAjBEoAACAMQIFAAAwxoWtgDaG+0oACEXsoQAAAMYIFAAAwBiBAgAAGCNQAAAAYwQKAABgjEABAACMESgAAIAxrkMBAGHCl2uUHM4eH4CZIBKxhwIAABgjUAAAAGMECgAAYMzvQFFeXq4nn3xSY8eOVZcuXWSz2bR69ep6+x44cEBjx46Vw+FQly5ddM899+j48eNe/Wpra/X888/rqquukt1uV0pKil577TV/pwgAAALE70Bx4sQJPfXUUzpw4IC+973vNdivqKhII0aMUGFhoRYvXqwFCxbonXfe0Y033qhz58559F24cKEeffRR3XjjjfrVr36lPn36aMqUKXr99df9nSYAAAgAvz/l4XQ6VVxcrMTERO3du1dpaWn19lu8eLEqKir08ccfq0+fPpIkl8ulG2+8UatXr9bs2bMlSceOHdOyZcs0Z84cLV++XJI0c+ZMjRw5Ug8//LAyMjLUrl07f6cLAABakd97KGJiYpSYmNhkvz/84Q+65ZZb3GFCkm644Qb1799fb775prtty5YtOn/+vDIzM91tNptN999/v4qKirR7925/pwoAAFpZq56UeezYMZWWluq6667zWuZyubRv3z73z/v27VNcXJwGDBjg1a9uOQAACE2temGr4uJiSRcPj1zO6XTq5MmTqq6uVkxMjIqLi9W9e3fZbDavfpL09ddfN/g4paWlXid5FhYWSrp48mhZWZlRHaGmoqLC43skiKSaE2MtSVJXu+f3SBBpNTe33qbey+peOyZjtLZI2pbrtMWay8vLm71OqwaKyspKSRcPj1zObre7+8TExLi/N9avIStWrNCiRYvqXZafn6+SkpJmz70tyM/PD/YUAi4Sas5K9fx51rW1QZlHMEVazb7Wm5OT0+jyy187/owRKJGwLV+uLdV85MiRZq/TqoEiNjZWklRdXe21rKqqyqNPbGysT/3qk5mZqYyMDI+2wsJCTZw4US6Xy+swSltXUVGh/Px8uVwuxcXFBXs6ARFJNY958UNJF/9qnXVtrV75IkonqmxNrBUeIq3m5tb73vwRjS6ve+2YjNHaImlbrtMWaz5w4ECz12nVQFF3uKLu0MeliouL1aVLF/deCafTqZycHFmW5XHYo27dHj16NPg4CQkJSkhIqHeZw+FQfHy83zWEsri4uLCtrSGRUHNJped/LCeqbF5t4S7Sava13qZe+y0xRqBEwrZ8ubZUs8PhaPY6rXpSZs+ePdWtWzft3bvXa1l+fr5SU1PdP6empurs2bNeqSgvL8+9HAAAhKZWv9vonXfeqTVr1ujo0aPq3bu3JOnPf/6zCgoKNH/+fHe/CRMmaP78+VqxYoX7OhSWZWnlypXq2bOn0tPTW3uqQKvz5W6QANAWGQWK5cuX69SpU+5PYGzdulVFRUWSpAceeECdOnXS448/ro0bN2rUqFH62c9+pvLyci1ZskSDBg3SjBkz3GP16tVL8+bN05IlS3T+/HmlpaVp8+bN2rlzp9avX89FrQAACGFGgWLp0qX66quv3D+/9dZbeuuttyRJU6dOVadOndS7d2/t2LFDDz74oB577DF16NBB48eP17Jly7w+1ZGdna3OnTtr1apVWr16tfr166d169ZpypQpJtMEgLDAHi6EMqNAcfjwYZ/6JScna/v27U32i4qKUlZWlrKyskymBQAAAozblwMAAGMECgAAYIxAAQAAjBEoAACAMQIFAAAwRqAAAADGCBQAAMAYgQIAABgjUAAAAGMECgAAYKzV7zYKRBLutQAgUrGHAgAAGCNQAAAAYwQKAABgjEABAACMESgAAIAxAgUAADBGoAAAAMYIFAAAwBiBAgAAGCNQAAAAYwQKAABgjEABAACMESgAAIAxAgUAADBGoAAAAMYIFAAAwBiBAgAAGCNQAAAAYwQKAABgjEABAACMESgAAIAxAgUAADBGoAAAAMYIFAAAwBiBAgAAGCNQAAAAYwQKAABgjEABAACMRQd7AgCA0NH3sXea7HM4e3wAZoK2hkAB/H9NvZHyJgoADeOQBwAAMEagAAAAxggUAADAGIECAAAYI1AAAABjBAoAAGCMQAEAAIwRKAAAgDEubAUAEcSXK2EC/mAPBQAAMEagAAAAxggUAADAGIECAAAY46RMAECLauzEz8RYS1mpgZsLAoc9FAAAwBiBAgAAGGv1QJGbmyubzVbv1549ezz67tq1S8OHD1fHjh2VmJiouXPnqry8vLWnCAAADAXsHIq5c+cqLS3Noy0pKcn97/3792v06NEaMGCAXnjhBRUVFWnp0qU6ePCg3n333UBNEwAA+CFggeL666/XXXfd1eDyxx9/XJ07d1Zubq7i4+MlSX379tWsWbP03nvvacyYMYGaKgAAaKaAnkNx5swZXbhwwau9rKxM77//vqZOneoOE5I0bdo0ORwOvfnmm4GcJgAAaKaABYoZM2YoPj5edrtdo0aN0t69e93LPv30U124cEHXXXedxzodOnRQamqq9u3bF6hpAgAAP7T6IY8OHTrozjvv1Lhx49S1a1d9/vnnWrp0qa6//nrt2rVLgwcPVnFxsSTJ6XR6re90OrVz585GH6O0tFTHjx/3aCssLJQklZeXq6ysrIWqCQ0VFRUe3yNBIGpOjLUaXe7L66ipMZqjq93y+B4JIq3mtlpvU9tCY9tBXa28f4U2fz4QYbMsK+Cv5MLCQqWkpGjEiBHatm2b1q5dq2nTpikvL08ul8uj77Rp0/T222/r1KlTDY7385//XIsWLap32UsvvaQ+ffq05PQBAAhrR44c0dy5c/XZZ58pOTnZp3WCcqXMpKQkTZgwQW+99ZZqamoUGxsrSaqurvbqW1VV5V7ekMzMTGVkZHi0FRYWauLEiXK5XBowYEDLTT4EVFRUKD8/Xy6XS3FxccGeTkAEouYxL37Y6PL35o8wHqM5utotzbq2Vq98EaUTVbYWGzeURVrNbbXepraFxraDupp5/wptBw4caPY6Qbv0du/evXXu3DlVVFS4D3XUHfq4VHFxsXr06NHoWAkJCUpISKh3mcPh8DjRM5zExcWFbW0Nac2aSyobf0P35XGbGsMfJ6psrTJuKIu0mttavU1tC77UwvtXaHM4HM1eJ2hXyjx06JDsdrscDocGDhyo6OhojxM1JencuXPav3+/UlNTgzNJAADgk1YPFJefLClJn3zyid5++22NGTNGUVFR6tSpk2644QatW7dOZ86ccfdbu3atysvLvQ5nAACA0NLqhzwmTZqk2NhYpaenKyEhQZ9//rlefvlldezYUdnZ2e5+zzzzjNLT0zVy5EjNnj1bRUVFWrZsmcaMGaOxY8e29jTRxjV2d0NJOpw9PkAzAYDI1Op7KCZOnKgTJ07ohRdeUGZmpt544w3dcccd2rt3r8fJkkOGDNEHH3yg2NhYzZ8/Xy+//LLuvfdebdq0qbWnCAAADLX6Hoq5c+dq7ty5PvUdPny4Pvroo1aeEeCfpvaCAEAk4/blAADAGIECAAAYI1AAAABjBAoAAGCMQAEAAIwRKAAAgDECBQAAMBa0m4MBANomrsmC+rCHAgAAGCNQAAAAYwQKAABgjEABAACMESgAAIAxAgUAADBGoAAAAMYIFAAAwBgXtkJE4EI8QNviyzZ7OHt8AGYCXxEoAAABN+bFD1VSaQv2NNCCOOQBAACMESgAAIAxAgUAADBGoAAAAMYIFAAAwBiBAgAAGCNQAAAAYwQKAABgjAtbIeT1fewdJcZaykrlYjgAEKrYQwEAAIwRKAAAgDECBQAAMEagAAAAxggUAADAGIECAAAYI1AAAABjBAoAAGCMQAEAAIwRKAAAgDEuvY2g6/vYO8GeAgDAEIECANAmNfXHyOHs8QGaCSQOeQAAgBZAoAAAAMYIFAAAwBiBAgAAGCNQAAAAYwQKAABgjI+NAgDCEh8rDSz2UAAAAGMECgAAYIxAAQAAjBEoAACAMU7KjGC+3JSrqZOWuLEXgHDGiZ2+Yw8FAAAwRqAAAADGCBQAAMAY51C0YWNe/FAllbYGl3NsDwAQKOyhAAAAxggUAADAWEgd8qiurtZ//Md/aO3atfr222+VkpKiX/ziF7rxxhuDPTUAALz48tH5vz1+fQBmEnwhFSimT5+uTZs2ad68eerXr59Wr16tcePGKScnR8OHDw/29AAAYSRQ19EZ8+KHykpt+Ly3cDnfLWQCRX5+vl5//XUtWbJECxYskCRNmzZNAwcO1COPPKJdu3YFeYYAAKAhIRMoNm3apHbt2mn27NnuNrvdrnvvvVePP/64jh49qt69ewdxhi2nJa5Q2VKPAwAIrkD9n9DaQiZQ7Nu3T/3791d8fLxHu8vlkiTt37+/wUBRWlqq48ePe7R9/vnnkqS//e1vKi8vb7F5zvr9XuMx4pvuory8vAaXVVZW6siRI4o5HaX4auPpNCplzq8bXe5LLS0h5px05EhtQGoOFdQc7Nm0vkirV6Jmf2tu6r34lWnX+TdwAw4dOiTp4rmNPrNCRHJysvXDH/7Qq/1//ud/LEnWypUrG1z3ySeftCTxxRdffPHFF18t+LV582af/x8PmT0UlZWViomJ8Wq32+3u5Q3JzMxURkaGR1tZWZkKCgo0aNCgesdtywoLCzVx4kRt3rxZSUlJwZ5OQFAzNYejSKtXoua2UnN1dbWOHj2qkSNH+rxOyASK2NjYenetVFVVuZc3JCEhQQkJCV7tw4YNa7kJhqCkpCQlJycHexoBRc2RIdJqjrR6JWpuC4YMGdKs/iFzYSun06ni4mKv9rq2Hj16BHpKAADARyETKFJTU1VQUKCysjKP9rqTE1NTU4MwKwAA4IuQCRR33XWXampq9PLLL7vbqqur9bvf/U5Dhw4Nm4+MAgAQjkLmHIqhQ4cqIyNDWVlZKi0tVVJSktasWaPDhw/rt7/9bbCnF1K6deumJ598Ut26dQv2VAKGmiNDpNUcafVK1BzObJZlWcGeRJ2qqio98cQTWrdunfteHk8//bRuuummYE8NAAA0IqQCBQAAaJtC5hwKAADQdhEoAACAMQIFAAAwRqAAAADGCBRBlpubK5vNVu/Xnj17JEmHDx9usI/NZtOsWbOaNV6o+Otf/6rbbrtNXbp0UceOHTVw4EC99NJLHn127dql4cOHq2PHjkpMTNTcuXPrvXtsdXW1Hn30UfXo0UOxsbEaOnSo3n///UCV4pPG6j179qx+/etfa8yYMXI6nfrOd76jwYMH6ze/+Y1qamo8xmns9fD6668Ho7QGNfUc/8u//Eu9dYwdO9ZrrLbwHEuN1xxu2/L06dMbrefYsWPuvuGwLftSb7huy74ImetQRLq5c+cqLS3No63uJjLdunXT2rVrvdbZtm2b1q9frzFjxjRrvFDw3nvv6dZbb9XgwYP1xBNPyOFw6Msvv1RRUZG7z/79+zV69GgNGDBAL7zwgoqKirR06VIdPHhQ7777rsd406dP16ZNmzRv3jz169dPq1ev1rhx45STk6Phw4cHujwvTdV76NAhPfDAAxo9erQefPBBxcfHa/v27crMzNSePXu0Zs0arzEnT56scePGebSF0v1rfHmOJalXr1569tlnPdrqu9R+qD/HUtM1h9u2/JOf/EQ33HCDR5tlWbrvvvvUt29f9ezZU1L4bMu+1PvZZ5+F3bbss2beZRwtLCcnx5Jkbdy4sdnrjh492oqPj7cqKytbZLxAOX36tNW9e3fr9ttvt2pqahrsd/PNN1tOp9M6ffq0u+2VV16xJFnbt293t+Xl5VmSrCVLlrjbKisrrauvvtoaNmxY6xTRDL7Ue/z4ceuzzz7zap8xY4YlyTp48KC77X//93+96g01vj7HI0eOtJKTk5scL9SfY8vyveb6tNVtuT47d+60JFnPPPOMuy1ctuX6XF5vuG3LzcEhjxBy5swZXbhwwae+xcXFysnJ0R133OG+xbvJeIG0YcMGffPNN3rmmWcUFRWliooK1dbWevQpKyvT+++/r6lTpyo+Pt7dPm3aNDkcDr355pvutk2bNqldu3aaPXu2u81ut+vee+/V7t27dfTo0dYvqhG+1Nu1a9d670J4++23S5IOHDhQ79gVFRU6d+5cy0/akC81X+rChQv17v6uE+rPsdT8muu05W25Phs2bJDNZtOUKVMkhde2XJ/L6w23bbk5CBQhYsaMGYqPj5fdbteoUaO0d+/eRvu//vrrqq2t1b/+67+2yHiB9MEHHyg+Pl7Hjh3TNddcI4fDofj4eN1///3u29V/+umnunDhgq677jqPdTt06KDU1FTt27fP3bZv3z7179/f481Kklwul6SLu1uDyZd6G1JSUiLp4pvU5RYtWiSHwyG73a60tDS99957rTJ/fzSn5oKCAsXFxek73/mOEhMT9cQTT+j8+fMefUL9OZb8f57b8rZ8ufPnz+vNN99Uenq6+vbtKym8tuXL1VdvQ9rqttwcnEMRZB06dNCdd96pcePGqWvXrvr888+1dOlSXX/99dq1a5cGDx5c73rr16+X0+nUD3/4wxYZL5AOHjyoCxcuaMKECbr33nv17LPPKjc3V7/61a906tQpvfbaa+7b1judTq/1nU6ndu7c6f65uLi4wX6S9PXXX7dSJb7xpd76nDt3Tv/5n/+pq666yuMYelRUlMaMGaPbb79dPXv21KFDh/TCCy/o5ptv1ttvv63x48cHqrQG+Vrz1VdfrVGjRmnQoEGqqKjQpk2b9Itf/EIFBQV644033OOF+nMs+f88t+Vt+XLbt2/XP/7xD49wFE7b8uXqq7c+bXlbbpZgH3OBt4MHD1qxsbHWTTfdVO/yv//975Yka/78+S0yXqD98z//syXJuu+++zzaf/KTn1iSrIKCAuv3v/+9JcnKy8vzWv+ee+6xOnXq5DHezTff7NXvyy+/tCRZL774YkuX0Cy+1FufWbNmWZKsd955p8nH+Mc//mF1797duuaaa1pkzqb8rdmy/q/u3bt3e4wXys+xZflXc1vfli83efJkq3379taJEyfcbeG0LV+uvnrr05a35ebgkEcISkpK0oQJE5STk+P1MSPp4l80kppMxb6OF2ixsbGSLp7ZfKm6Y5C7d+9296murvZav6qqyr28bryG+l36eMHiS72XW7JkiV555RU9/fTTXmd/16dLly6aMWOG/v73v3t9iiIY/Km5zkMPPSTp4iGES8cL5ef40jk0p+a2vi1fqry8XFu2bNFNN92kK6+80t0eTtvypRqq93JtfVtuDgJFiOrdu7fOnTuniooKr2UbNmzQNddco+9///stMl6g1X0ksHv37h7tCQkJkqRvv/3WvYuzbnfppYqLiz0+Vuh0Ohvsd+njBYsv9V5q9erVevTRR3Xffffp3//9331+nN69e0uSTp48aTLdFtHcmi9VXx2h/hxfOofm1NzWt+VLbd68WWfPnvUKR+G0LV+qoXovFQ7bcnMQKELUoUOHZLfb5XA4PNrz8vJUWFjo8180TY0XDHVvnpde9Eb6v+Oj3bp108CBAxUdHe11Atq5c+e0f/9+paamuttSU1NVUFCgsrIyj755eXnu5cHkS711tmzZopkzZ+qOO+7Qr3/962Y9zqFDh7zGC5bm1Hy5+uoI9edYan7N4bAtX2r9+vVyOBy67bbbPNrDaVu+VEP11gmXbblZgn3MJdKVlpZ6te3fv99q3769ddttt3ktmzt3riXJKiwsbJHxguGvf/2rJcmaMmWKR/vkyZOt6Oho69ixY5ZlWdbYsWMtp9NplZWVufu8+uqrliTr3Xffdbft2bPH67PcVVVVVlJSkjV06NBWrqZpvta7Y8cOy263W6NGjbKqqqoaHK++57ioqMjq3LmzlZKS0rKT95MvNZ8+fdqrztraWmvSpEmWJOvjjz92t4f6c2xZvj/PdcJhW65TWlpqRUdHW/fcc0+9y8NlW67TVL3htC03B5/yCLJJkyYpNjZW6enpSkhI0Oeff66XX35ZHTt2VHZ2tkffmpoavfHGG/rBD36gq6++2ni8YBk8eLD+7d/+Tf/1X/+lCxcuaOTIkcrNzdXGjRuVlZXl3q35zDPPKD09XSNHjtTs2bNVVFSkZcuWacyYMR6XZh46dKgyMjKUlZWl0tJSJSUlac2aNTp8+LB++9vfBqtMN1/q/eqrr3TbbbfJZrPprrvu0saNGz3GSElJUUpKiiTpkUce0ZdffqnRo0erR48eOnz4sFatWqWKigr98pe/DEaJXnypOTc3V5MnT9bkyZOVlJSkyspK/fd//7c++ugjzZ49W0OGDHGPF+rPseT761oKn225zhtvvKELFy40uLclXLblOo3VG27bcrMEO9FEul/+8peWy+WyunTpYkVHR1tOp9OaOnWqx9XU6mzbts2SZL300kstMl4wnTt3zvr5z39u/dM//ZPVvn17Kykpqd4zuHfu3Gmlp6dbdrvd6tatmzVnzhyPv3LqVFZWWgsWLLASExOtmJgYKy0tzdq2bVsAKvFNU/XWXRWxoa8nn3zS3XfDhg3WiBEjrG7dulnR0dFW165drdtvv93jL/pQ0FTNhw4dsjIyMqy+fftadrvd6tixo/X973/fWrlypVVbW+s1Xqg/x5bl++s6nLZly7KsH/zgB1ZCQoJ14cKFBvuEy7ZsWY3XG47bsq9slmVZgY8xAAAgnHBSJgAAMEagAAAAxggUAADAGIECAAAYI1AAAABjBAoAAGCMQAEAAIwRKAAAgDECBQAAMEagAAAAxggUAADAGIECAAAYI1AAAABjBAoAAGCMQAEAAIz9P5Cu58xcNqARAAAAAElFTkSuQmCC",
+      "text/plain": [
+       "<Figure size 600x480 with 1 Axes>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "plt.figure(figsize = (5, 4), dpi = 120)\n",
+    "plt.hist(transforms[\"train\"][0].frequencies.cpu().numpy(), bins = 50)\n",
+    "plt.grid()\n"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Now we can get the best model path stored in the checkpoint callback."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "'/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=96-step=18333.ckpt'"
+      ]
+     },
+     "execution_count": 21,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "best_model_path = checkpoint_callback.best_model_path\n",
+    "best_model_path"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Inference\n",
+    "\n",
+    "To obtain model scores, we will load the weights from the best checkpoint. To do this, we use the `LightningModule`, provide there the path to the checkpoint and the model instance."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import replay\n",
+    "torch.serialization.add_safe_globals([\n",
+    "    replay.nn.lightning.optimizer.OptimizerFactory,\n",
+    "    replay.nn.lightning.scheduler.LRSchedulerFactory\n",
+    "])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 23,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sasrec = make_sasrec(\n",
+    "    schema=tensor_schema,\n",
+    "    embedding_dim=EMBEDDING_DIM,\n",
+    "    max_sequence_length=MAX_SEQ_LEN,\n",
+    "    num_heads=NUM_HEADS,\n",
+    "    num_blocks=NUM_BLOCKS,\n",
+    "    dropout=DROPOUT,\n",
+    ")\n",
+    "\n",
+    "best_model = LightningModule.load_from_checkpoint(best_model_path, model=sasrec)\n",
+    "best_model.eval();"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Configure `ParquetModule` for inference"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 24,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "inference_metadata = {\"predict\": create_meta(shape=MAX_SEQ_LEN)}\n",
+    "\n",
+    "parquet_module = ParquetModule(\n",
+    "    predict_path=PREDICT_PATH,\n",
+    "    batch_size=BATCH_SIZE,\n",
+    "    metadata=inference_metadata,\n",
+    "    transforms=transforms,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "During inference, we can use `TopItemsCallback`. Such callback allows you to get scores for each user throughout the entire catalog and get recommendations in the form of ids of items with the highest score values.\n",
+    "\n",
+    "\n",
+    "Recommendations can be fetched in four formats: PySpark DataFrame, Pandas DataFrame, Polars DataFrame or raw PyTorch tensors. Each of the types corresponds a callback. In this example, we'll be using the `PandasTopItemsCallback`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 25,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "💡 Tip: For seamless cloud uploads and versioning, try installing [litmodels](https://pypi.org/project/litmodels/) to enable LitModelCheckpoint, which syncs automatically with the Lightning model registry.\n",
+      "GPU available: True (cuda), used: True\n",
+      "TPU available: False, using: 0 TPU cores\n",
+      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "9b94b2f1745140e5ae3bc1d63bdbbe4e",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Predicting: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "from replay.nn.lightning.callback import PandasTopItemsCallback\n",
+    "\n",
+    "csv_logger = CSVLogger(save_dir=\"sasrec/logs/test\", name=\"SasRec-example\")\n",
+    "\n",
+    "TOPK = [1, 5, 10, 20]\n",
+    "\n",
+    "pandas_prediction_callback = PandasTopItemsCallback(\n",
+    "    top_k=max(TOPK),\n",
+    "    query_column=\"user_id\",\n",
+    "    item_column=\"item_id\",\n",
+    "    rating_column=\"score\",\n",
+    ")\n",
+    "\n",
+    "trainer = L.Trainer(callbacks=[pandas_prediction_callback], logger=csv_logger, inference_mode=True)\n",
+    "trainer.predict(best_model, datamodule=parquet_module, return_predictions=False)\n",
+    "\n",
+    "pandas_res = pandas_prediction_callback.get_result()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 26,
+   "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>486</td>\n",
+       "      <td>6.457348</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>0</td>\n",
+       "      <td>224</td>\n",
+       "      <td>6.311622</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>0</td>\n",
+       "      <td>210</td>\n",
+       "      <td>6.299765</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>0</td>\n",
+       "      <td>101</td>\n",
+       "      <td>6.193144</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>0</td>\n",
+       "      <td>2142</td>\n",
+       "      <td>6.179489</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>...</th>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6037</th>\n",
+       "      <td>6039</td>\n",
+       "      <td>84</td>\n",
+       "      <td>10.029899</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6037</th>\n",
+       "      <td>6039</td>\n",
+       "      <td>2634</td>\n",
+       "      <td>9.860862</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6037</th>\n",
+       "      <td>6039</td>\n",
+       "      <td>1490</td>\n",
+       "      <td>9.738086</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6037</th>\n",
+       "      <td>6039</td>\n",
+       "      <td>2633</td>\n",
+       "      <td>9.715384</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6037</th>\n",
+       "      <td>6039</td>\n",
+       "      <td>2497</td>\n",
+       "      <td>9.684793</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "<p>120760 rows × 3 columns</p>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "      user_id item_id      score\n",
+       "0           0     486   6.457348\n",
+       "0           0     224   6.311622\n",
+       "0           0     210   6.299765\n",
+       "0           0     101   6.193144\n",
+       "0           0    2142   6.179489\n",
+       "...       ...     ...        ...\n",
+       "6037     6039      84  10.029899\n",
+       "6037     6039    2634   9.860862\n",
+       "6037     6039    1490   9.738086\n",
+       "6037     6039    2633   9.715384\n",
+       "6037     6039    2497   9.684793\n",
+       "\n",
+       "[120760 rows x 3 columns]"
+      ]
+     },
+     "execution_count": 26,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "pandas_res"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Calculating metrics\n",
+    "\n",
+    "*test_gt* is already encoded, so we can use it for computing metrics."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 27,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from replay.metrics import MAP, OfflineMetrics, Precision, Recall\n",
+    "from replay.metrics.torch_metrics_builder import metrics_to_df"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 28,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "result_metrics = OfflineMetrics(\n",
+    "    [Recall(TOPK), Precision(TOPK), MAP(TOPK)],\n",
+    "    query_column=\"user_id\",\n",
+    "    rating_column=\"score\",\n",
+    ")(pandas_res, test_gt.explode(\"item_id\"))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 29,
+   "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>5</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>MAP</th>\n",
+       "      <td>0.020371</td>\n",
+       "      <td>0.061189</td>\n",
+       "      <td>0.068575</td>\n",
+       "      <td>0.051101</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>Precision</th>\n",
+       "      <td>0.020371</td>\n",
+       "      <td>0.018731</td>\n",
+       "      <td>0.014790</td>\n",
+       "      <td>0.022093</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>Recall</th>\n",
+       "      <td>0.020371</td>\n",
+       "      <td>0.187314</td>\n",
+       "      <td>0.295793</td>\n",
+       "      <td>0.110467</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "k                 1        10        20         5\n",
+       "MAP        0.020371  0.061189  0.068575  0.051101\n",
+       "Precision  0.020371  0.018731  0.014790  0.022093\n",
+       "Recall     0.020371  0.187314  0.295793  0.110467"
+      ]
+     },
+     "execution_count": 29,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "metrics_to_df(result_metrics)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Let's call the `inverse_transform` encoder's function to get the final dataframe with recommendations"
+   ]
+  },
+  {
+   "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></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>2012</td>\n",
+       "      <td>500</td>\n",
+       "      <td>6.457348</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>2012</td>\n",
+       "      <td>231</td>\n",
+       "      <td>6.311622</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>2012</td>\n",
+       "      <td>216</td>\n",
+       "      <td>6.299765</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>2012</td>\n",
+       "      <td>104</td>\n",
+       "      <td>6.193144</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>2012</td>\n",
+       "      <td>2335</td>\n",
+       "      <td>6.179489</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>...</th>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6037</th>\n",
+       "      <td>5727</td>\n",
+       "      <td>86</td>\n",
+       "      <td>10.029899</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6037</th>\n",
+       "      <td>5727</td>\n",
+       "      <td>2841</td>\n",
+       "      <td>9.860862</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6037</th>\n",
+       "      <td>5727</td>\n",
+       "      <td>1623</td>\n",
+       "      <td>9.738086</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6037</th>\n",
+       "      <td>5727</td>\n",
+       "      <td>2840</td>\n",
+       "      <td>9.715384</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6037</th>\n",
+       "      <td>5727</td>\n",
+       "      <td>2702</td>\n",
+       "      <td>9.684793</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "<p>120760 rows × 3 columns</p>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "      user_id  item_id      score\n",
+       "0        2012      500   6.457348\n",
+       "0        2012      231   6.311622\n",
+       "0        2012      216   6.299765\n",
+       "0        2012      104   6.193144\n",
+       "0        2012     2335   6.179489\n",
+       "...       ...      ...        ...\n",
+       "6037     5727       86  10.029899\n",
+       "6037     5727     2841   9.860862\n",
+       "6037     5727     1623   9.738086\n",
+       "6037     5727     2840   9.715384\n",
+       "6037     5727     2702   9.684793\n",
+       "\n",
+       "[120760 rows x 3 columns]"
+      ]
+     },
+     "execution_count": 30,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "encoder.inverse_transform(pandas_res)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "new_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.12.3"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/examples/09_sasrec_example_sampled.ipynb b/examples/09_sasrec_example_sampled.ipynb
new file mode 100644
index 000000000..6672edd85
--- /dev/null
+++ b/examples/09_sasrec_example_sampled.ipynb
@@ -0,0 +1,4744 @@
+{
+ "cells": [
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Example of SasRec training/inference with Parquet Module"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Seed set to 42\n"
+     ]
+    }
+   ],
+   "source": [
+    "from typing import Optional\n",
+    "\n",
+    "import lightning as L\n",
+    "import pandas as pd\n",
+    "\n",
+    "L.seed_everything(42)\n",
+    "\n",
+    "import warnings\n",
+    "warnings.filterwarnings(\"ignore\")"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Preparing data\n",
+    "In this example, we will be using the MovieLens dataset, namely the 1m subset. It's demonstrated a simple case, so only item ids will be used as model input.\n",
+    "\n",
+    "---\n",
+    "**NOTE**\n",
+    "\n",
+    "Current implementation of SasRec is able to handle item and interactions features. It does not take into account user features. \n",
+    "\n",
+    "---"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "interactions = pd.read_csv(\"./data/ml1m_ratings.dat\", sep=\"\\t\", names=[\"user_id\", \"item_id\",\"rating\",\"timestamp\"])\n",
+    "interactions = interactions.drop(columns=[\"rating\"])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "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>timestamp</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>1000138</th>\n",
+       "      <td>6040</td>\n",
+       "      <td>858</td>\n",
+       "      <td>0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1000153</th>\n",
+       "      <td>6040</td>\n",
+       "      <td>2384</td>\n",
+       "      <td>1</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>999873</th>\n",
+       "      <td>6040</td>\n",
+       "      <td>593</td>\n",
+       "      <td>2</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1000007</th>\n",
+       "      <td>6040</td>\n",
+       "      <td>1961</td>\n",
+       "      <td>3</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1000192</th>\n",
+       "      <td>6040</td>\n",
+       "      <td>2019</td>\n",
+       "      <td>4</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>...</th>\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>446</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>825438</th>\n",
+       "      <td>4958</td>\n",
+       "      <td>1407</td>\n",
+       "      <td>447</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>825724</th>\n",
+       "      <td>4958</td>\n",
+       "      <td>3264</td>\n",
+       "      <td>448</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>825731</th>\n",
+       "      <td>4958</td>\n",
+       "      <td>2634</td>\n",
+       "      <td>449</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>825603</th>\n",
+       "      <td>4958</td>\n",
+       "      <td>1924</td>\n",
+       "      <td>450</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "<p>1000209 rows × 3 columns</p>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "         user_id  item_id  timestamp\n",
+       "1000138     6040      858          0\n",
+       "1000153     6040     2384          1\n",
+       "999873      6040      593          2\n",
+       "1000007     6040     1961          3\n",
+       "1000192     6040     2019          4\n",
+       "...          ...      ...        ...\n",
+       "825793      4958     2399        446\n",
+       "825438      4958     1407        447\n",
+       "825724      4958     3264        448\n",
+       "825731      4958     2634        449\n",
+       "825603      4958     1924        450\n",
+       "\n",
+       "[1000209 rows x 3 columns]"
+      ]
+     },
+     "execution_count": 3,
+     "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"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Encode catagorical data.\n",
+    "To ensure all categorical data is fit for training, it needs to be encoded using the `LabelEncoder` class. Create an instance of the encoder, providing a `LabelEncodingRule` for each categorcial column in the dataset that will be used in model. Note that ids of users and ids of items are always used."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "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>timestamp</th>\n",
+       "      <th>user_id</th>\n",
+       "      <th>item_id</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>12</td>\n",
+       "      <td>0</td>\n",
+       "      <td>0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>68</td>\n",
+       "      <td>1</td>\n",
+       "      <td>0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>67</td>\n",
+       "      <td>2</td>\n",
+       "      <td>0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>12</td>\n",
+       "      <td>3</td>\n",
+       "      <td>0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>140</td>\n",
+       "      <td>4</td>\n",
+       "      <td>0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>...</th>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1000204</th>\n",
+       "      <td>14</td>\n",
+       "      <td>4555</td>\n",
+       "      <td>3705</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1000205</th>\n",
+       "      <td>90</td>\n",
+       "      <td>2813</td>\n",
+       "      <td>3705</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1000206</th>\n",
+       "      <td>70</td>\n",
+       "      <td>2404</td>\n",
+       "      <td>3705</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1000207</th>\n",
+       "      <td>25</td>\n",
+       "      <td>5835</td>\n",
+       "      <td>3705</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1000208</th>\n",
+       "      <td>380</td>\n",
+       "      <td>979</td>\n",
+       "      <td>3705</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "<p>1000209 rows × 3 columns</p>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "         timestamp  user_id  item_id\n",
+       "0               12        0        0\n",
+       "1               68        1        0\n",
+       "2               67        2        0\n",
+       "3               12        3        0\n",
+       "4              140        4        0\n",
+       "...            ...      ...      ...\n",
+       "1000204         14     4555     3705\n",
+       "1000205         90     2813     3705\n",
+       "1000206         70     2404     3705\n",
+       "1000207         25     5835     3705\n",
+       "1000208        380      979     3705\n",
+       "\n",
+       "[1000209 rows x 3 columns]"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from replay.preprocessing.label_encoder import LabelEncoder, LabelEncodingRule\n",
+    "\n",
+    "encoder = LabelEncoder(\n",
+    "    [\n",
+    "        LabelEncodingRule(\"user_id\", default_value=\"last\"),\n",
+    "        LabelEncodingRule(\"item_id\", default_value=\"last\"),\n",
+    "    ]\n",
+    ")\n",
+    "interactions = interactions.sort_values(by=\"item_id\", ascending=True)\n",
+    "encoded_interactions = encoder.fit_transform(interactions)\n",
+    "encoded_interactions"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Split interactions into the train, validation and test datasets using LastNSplitter\n",
+    "We use widespread splitting strategy Last-One-Out. We filter out cold items and users for simplicity."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from replay.splitters import LastNSplitter\n",
+    "\n",
+    "splitter = LastNSplitter(\n",
+    "    N=1,\n",
+    "    divide_column=\"user_id\",\n",
+    "    query_column=\"user_id\",\n",
+    "    strategy=\"interactions\",\n",
+    "    drop_cold_users=True,\n",
+    "    drop_cold_items=True\n",
+    ")\n",
+    "\n",
+    "test_events, test_gt = splitter.split(encoded_interactions)\n",
+    "validation_events, validation_gt = splitter.split(test_events)\n",
+    "train_events = validation_events"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Dataset preprocessing (\"baking\")\n",
+    "SasRec expects each user in the batch to provide their events in form of a sequence. For this reason, the event splits must be properly processed using the `groupby_sequences` function provided by RePlay."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from replay.data.nn.utils import groupby_sequences\n",
+    "\n",
+    "\n",
+    "def bake_data(full_data):\n",
+    "    grouped_interactions = groupby_sequences(events=full_data, groupby_col=\"user_id\", sort_col=\"timestamp\")\n",
+    "    return grouped_interactions"
+   ]
+  },
+  {
+   "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>timestamp</th>\n",
+       "      <th>item_id</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>0</td>\n",
+       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
+       "      <td>[2426, 822, 2733, 2587, 2937, 3618, 2943, 708,...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>1</td>\n",
+       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
+       "      <td>[3272, 3026, 2760, 851, 346, 3393, 1107, 515, ...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>2</td>\n",
+       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
+       "      <td>[579, 1140, 1154, 2426, 1524, 1260, 2160, 2621...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>3</td>\n",
+       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
+       "      <td>[1781, 2940, 2468, 890, 948, 106, 593, 309, 49...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>4</td>\n",
+       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
+       "      <td>[1108, 2229, 21, 2435, 2142, 106, 1167, 593, 1...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>...</th>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6035</th>\n",
+       "      <td>6035</td>\n",
+       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
+       "      <td>[2426, 1279, 3151, 3321, 1178, 3301, 2501, 278...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6036</th>\n",
+       "      <td>6036</td>\n",
+       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
+       "      <td>[1592, 2302, 1633, 1813, 2879, 1482, 2651, 250...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6037</th>\n",
+       "      <td>6037</td>\n",
+       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
+       "      <td>[1971, 3500, 2077, 1666, 1399, 2651, 2748, 283...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6038</th>\n",
+       "      <td>6038</td>\n",
+       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
+       "      <td>[1486, 1485, 3384, 3512, 3302, 3126, 3650, 330...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6039</th>\n",
+       "      <td>6039</td>\n",
+       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
+       "      <td>[2432, 2960, 1848, 2114, 2142, 3091, 3248, 317...</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "<p>6040 rows × 3 columns</p>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "      user_id                                          timestamp  \\\n",
+       "0           0  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
+       "1           1  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
+       "2           2  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
+       "3           3  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
+       "4           4  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
+       "...       ...                                                ...   \n",
+       "6035     6035  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
+       "6036     6036  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
+       "6037     6037  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
+       "6038     6038  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
+       "6039     6039  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
+       "\n",
+       "                                                item_id  \n",
+       "0     [2426, 822, 2733, 2587, 2937, 3618, 2943, 708,...  \n",
+       "1     [3272, 3026, 2760, 851, 346, 3393, 1107, 515, ...  \n",
+       "2     [579, 1140, 1154, 2426, 1524, 1260, 2160, 2621...  \n",
+       "3     [1781, 2940, 2468, 890, 948, 106, 593, 309, 49...  \n",
+       "4     [1108, 2229, 21, 2435, 2142, 106, 1167, 593, 1...  \n",
+       "...                                                 ...  \n",
+       "6035  [2426, 1279, 3151, 3321, 1178, 3301, 2501, 278...  \n",
+       "6036  [1592, 2302, 1633, 1813, 2879, 1482, 2651, 250...  \n",
+       "6037  [1971, 3500, 2077, 1666, 1399, 2651, 2748, 283...  \n",
+       "6038  [1486, 1485, 3384, 3512, 3302, 3126, 3650, 330...  \n",
+       "6039  [2432, 2960, 1848, 2114, 2142, 3091, 3248, 317...  \n",
+       "\n",
+       "[6040 rows x 3 columns]"
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "train_events = bake_data(train_events)\n",
+    "\n",
+    "validation_events = bake_data(validation_events)\n",
+    "validation_gt = bake_data(validation_gt)\n",
+    "\n",
+    "test_events = bake_data(test_events)\n",
+    "test_gt = bake_data(test_gt)\n",
+    "\n",
+    "train_events"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "To ensure we don't have unknown users in ground truth, we join validation events and validation ground truth (also join test events and test ground truth correspondingly) by user ids to leave only the common ones.  "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def add_gt_to_events(events_df, gt_df):\n",
+    "    gt_to_join = gt_df[[\"user_id\", \"item_id\"]].rename(columns={\"item_id\": \"ground_truth\"})\n",
+    "\n",
+    "    events_df = events_df.merge(gt_to_join, on=\"user_id\", how=\"inner\")\n",
+    "    return events_df\n",
+    "\n",
+    "validation_events = add_gt_to_events(validation_events, validation_gt)\n",
+    "test_events = add_gt_to_events(test_events, test_gt)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pathlib import Path\n",
+    "\n",
+    "data_dir = Path(\"temp/data/\")\n",
+    "data_dir.mkdir(parents=True, exist_ok=True)\n",
+    "\n",
+    "TRAIN_PATH = data_dir / \"train.parquet\"\n",
+    "VAL_PATH = data_dir / \"val.parquet\"\n",
+    "PREDICT_PATH = data_dir / \"test.parquet\"\n",
+    "\n",
+    "ENCODER_PATH = data_dir / \"encoder\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_events.to_parquet(TRAIN_PATH)\n",
+    "validation_events.to_parquet(VAL_PATH)\n",
+    "test_events.to_parquet(PREDICT_PATH)\n",
+    "\n",
+    "encoder.save(ENCODER_PATH)"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Prepare to model training\n",
+    "### 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. It is required by the SasRec model to correctly create embeddings for every source column. Note that user_id does not required in `TensorSchema`.\n",
+    "\n",
+    "Note that the **padding value** is the next value (item_id) after the last one. **Cardinality** is the number of unique values ​​given the padding value."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from replay.data import FeatureHint, FeatureType\n",
+    "from replay.data.nn import TensorFeatureInfo, TensorSchema\n",
+    "\n",
+    "\n",
+    "EMBEDDING_DIM = 64\n",
+    "\n",
+    "encoder = encoder.load(ENCODER_PATH)\n",
+    "NUM_UNIQUE_ITEMS = len(encoder.mapping[\"item_id\"])\n",
+    "\n",
+    "tensor_schema = TensorSchema(\n",
+    "    [\n",
+    "        TensorFeatureInfo(\n",
+    "            name=\"item_id\",\n",
+    "            is_seq=True,\n",
+    "            padding_value=NUM_UNIQUE_ITEMS,\n",
+    "            cardinality=NUM_UNIQUE_ITEMS + 1,  # taking into account padding\n",
+    "            embedding_dim=EMBEDDING_DIM,\n",
+    "            feature_type=FeatureType.CATEGORICAL,\n",
+    "            feature_hint=FeatureHint.ITEM_ID,\n",
+    "        )\n",
+    "    ]\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Configure ParquetModule and transformation pipelines\n",
+    "\n",
+    "The `ParquetModule` class enables training of models on large datasets by reading data in batch-wise way. This class initialized with **paths to every data split, a metadata dict containing information about shape and padding value of every column and a dict of transforms**. `ParquetModule`'s  \"transform pipelines\" are stage-specific modules implementing additional preprocessing to be performed on batch level right before the forward pass.  \n",
+    "\n",
+    "For SasRec model, RePlay provides a function that generates a sequence of appropriate transforms for each data split named **make_default_sasrec_transforms**.\n",
+    "\n",
+    "Internally this function creates the following transforms:\n",
+    "1) Training:\n",
+    "    1. Create a target, which contains the shifted item sequence that represents the next item in the sequence (for the next item prediction task).\n",
+    "    2. Rename features to match it with expected format by the model during training.\n",
+    "    3. Unsqueeze target (*positive_labels*) and it's padding mask (*target_padding_mask*) for getting required shape of this tensors for loss computation.\n",
+    "    4. Group input features to be embed in expected format.\n",
+    "\n",
+    "2) Validation/Inference:\n",
+    "    1. Rename/group features to match it with expected format by the model during valdiation/inference.\n",
+    "\n",
+    "If a different set of transforms is required, you can create them yourself and submit them to the ParquetModule in the form of a dictionary where the key is the name of the split, and the value is the list of transforms. Available transforms are in the replay/nn/transforms/.\n",
+    "\n",
+    "**Note:** One of the transforms for the training data prepares the initial sequence for the task of Next Item Prediction so it shifts the sequence of items. For the final sequence length to be correct, you need to set shape of item_id in metadata as **model sequence length + shift**. Default shift value is 1."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import copy\n",
+    "\n",
+    "import torch\n",
+    "\n",
+    "from replay.data.nn import TensorSchema\n",
+    "from replay.nn.transform import GroupTransform, NextTokenTransform, RenameTransform, UnsqueezeTransform, UniformNegativeSamplingTransform"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def make_sasrec_transforms(\n",
+    "    tensor_schema: TensorSchema, query_column: str = \"query_id\", num_negative_samples: int = 128,\n",
+    ") -> dict[str, list[torch.nn.Module]]:\n",
+    "    item_column = tensor_schema.item_id_feature_name\n",
+    "    vocab_size = tensor_schema[item_column].cardinality\n",
+    "    train_transforms = [\n",
+    "        UniformNegativeSamplingTransform(vocab_size, num_negative_samples),\n",
+    "        NextTokenTransform(label_field=item_column, query_features=query_column, shift=1),\n",
+    "        RenameTransform(\n",
+    "            {\n",
+    "                query_column: \"query_id\",\n",
+    "                f\"{item_column}_mask\": \"padding_mask\",\n",
+    "                \"positive_labels_mask\": \"target_padding_mask\",\n",
+    "            }\n",
+    "        ),\n",
+    "        UnsqueezeTransform(\"target_padding_mask\", -1),\n",
+    "        UnsqueezeTransform(\"positive_labels\", -1),\n",
+    "        GroupTransform({\"feature_tensors\": [item_column]}),\n",
+    "    ]\n",
+    "\n",
+    "    val_transforms = [\n",
+    "        RenameTransform({query_column: \"query_id\", f\"{item_column}_mask\": \"padding_mask\"}),\n",
+    "        GroupTransform({\"feature_tensors\": [item_column]}),\n",
+    "    ]\n",
+    "    test_transforms = copy.deepcopy(val_transforms)\n",
+    "\n",
+    "    predict_transforms = copy.deepcopy(val_transforms)\n",
+    "\n",
+    "    transforms = {\n",
+    "        \"train\": train_transforms,\n",
+    "        \"validate\": val_transforms,\n",
+    "        \"test\": test_transforms,\n",
+    "        \"predict\": predict_transforms,\n",
+    "    }\n",
+    "\n",
+    "    return transforms\n",
+    "\n",
+    "transforms = make_sasrec_transforms(tensor_schema, query_column=\"user_id\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "MAX_SEQ_LEN = 50\n",
+    "\n",
+    "def create_meta(shape: int, gt_shape: Optional[int] = None):\n",
+    "    meta = {\n",
+    "        \"user_id\": {},\n",
+    "        \"item_id\": {\"shape\": shape, \"padding\": tensor_schema[\"item_id\"].padding_value},\n",
+    "    }\n",
+    "    if gt_shape is not None:\n",
+    "        meta.update({\"ground_truth\": {\"shape\": gt_shape, \"padding\": -1}})\n",
+    "\n",
+    "    return meta\n",
+    "\n",
+    "train_metadata = {\n",
+    "    \"train\": create_meta(shape=MAX_SEQ_LEN+1),\n",
+    "    \"validate\": create_meta(shape=MAX_SEQ_LEN, gt_shape=1),\n",
+    "}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from replay.data.nn import ParquetModule\n",
+    "\n",
+    "BATCH_SIZE = 32\n",
+    "\n",
+    "parquet_module = ParquetModule(\n",
+    "    train_path=TRAIN_PATH,\n",
+    "    validate_path=VAL_PATH,\n",
+    "    batch_size=BATCH_SIZE,\n",
+    "    metadata=train_metadata,\n",
+    "    transforms=transforms,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Train model\n",
+    "### Create SasRec model instance and run the training stage using lightning\n",
+    "We may now train the model using the Lightning trainer class. \n",
+    "\n",
+    "RePlay's implementation of SasRec is designed in a modular, **block-based approach**. Instead of passing configuration parameters to the constructor, SasRec is now built by providing fully initialized components that makes the model more flexible and easier to extend.\n",
+    "\n",
+    "#### Default Configuration\n",
+    "\n",
+    "Default SasRec model may be created quickly via method **from_params**. Default model instance has CE loss, original SasRec transformer layes, and embeddings are aggregated via sum."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 23,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from replay.nn.sequential import SasRec\n",
+    "from typing import Literal\n",
+    "def make_sasrec(\n",
+    "    schema: TensorSchema,\n",
+    "    embedding_dim: int = 192,\n",
+    "    num_heads: int = 4,\n",
+    "    num_blocks: int = 2,\n",
+    "    max_sequence_length: int = 50,\n",
+    "    dropout: float = 0.3,\n",
+    "    excluded_features: Optional[list[str]] = None,\n",
+    "    categorical_list_feature_aggregation_method: Literal[\"sum\", \"mean\", \"max\"] = \"sum\",\n",
+    ") -> SasRec:\n",
+    "    from replay.nn.sequential.sasrec import SasRecBody, SasRecTransformerLayer\n",
+    "    from replay.nn.agg import SumAggregator\n",
+    "    from replay.nn.embedding import SequenceEmbedding\n",
+    "    from replay.nn.loss import CE, CESampled\n",
+    "    from replay.nn.mask import DefaultAttentionMask\n",
+    "    from replay.nn.sequential.sasrec.agg import PositionAwareAggregator\n",
+    "    from replay.nn.sequential.sasrec.transformer import SasRecTransformerLayer\n",
+    "    excluded_features = [\n",
+    "        schema.query_id_feature_name,\n",
+    "        schema.timestamp_feature_name,\n",
+    "        *(excluded_features or []),\n",
+    "    ]\n",
+    "    excluded_features = list(set(excluded_features))\n",
+    "    body = SasRecBody(\n",
+    "        embedder=SequenceEmbedding(\n",
+    "            schema=schema,\n",
+    "            categorical_list_feature_aggregation_method=categorical_list_feature_aggregation_method,\n",
+    "            excluded_features=excluded_features,\n",
+    "        ),\n",
+    "        embedding_aggregator=PositionAwareAggregator(\n",
+    "            embedding_aggregator=SumAggregator(embedding_dim=embedding_dim),\n",
+    "            max_sequence_length=max_sequence_length,\n",
+    "            dropout=dropout,\n",
+    "        ),\n",
+    "        attn_mask_builder=DefaultAttentionMask(\n",
+    "            reference_feature_name=schema.item_id_feature_name,\n",
+    "            num_heads=num_heads,\n",
+    "        ),\n",
+    "        encoder=SasRecTransformerLayer(\n",
+    "            embedding_dim=embedding_dim,\n",
+    "            num_heads=num_heads,\n",
+    "            num_blocks=num_blocks,\n",
+    "            dropout=dropout,\n",
+    "            activation=\"relu\",\n",
+    "        ),\n",
+    "        output_normalization=torch.nn.LayerNorm(embedding_dim),\n",
+    "    )\n",
+    "    padding_idx = schema.item_id_features.item().padding_value\n",
+    "    return SasRec(\n",
+    "        body=body,\n",
+    "        loss=CESampled(padding_idx=padding_idx),\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 24,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "NUM_BLOCKS = 2\n",
+    "NUM_HEADS = 2\n",
+    "DROPOUT = 0.3\n",
+    "\n",
+    "sasrec = make_sasrec(\n",
+    "    schema=tensor_schema,\n",
+    "    embedding_dim=EMBEDDING_DIM,\n",
+    "    max_sequence_length=MAX_SEQ_LEN,\n",
+    "    num_heads=NUM_HEADS,\n",
+    "    num_blocks=NUM_BLOCKS,\n",
+    "    dropout=DROPOUT,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "A universal PyTorch Lightning module is provided. It can work with any NN model."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 25,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from replay.nn.lightning.optimizer import OptimizerFactory\n",
+    "from replay.nn.lightning.scheduler import LRSchedulerFactory\n",
+    "from replay.nn.lightning import LightningModule\n",
+    "\n",
+    "model = LightningModule(\n",
+    "    sasrec,\n",
+    "    optimizer_factory=OptimizerFactory(),\n",
+    "    lr_scheduler_factory=LRSchedulerFactory(),\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "To facilitate training, we add the following callbacks:\n",
+    "1) `ModelCheckpoint` - to save the best trained model based on its Recall metric. It's a default Lightning Callback.\n",
+    "1) `ComputeMetricsCallback` - to display a detailed validation metric matrix after each epoch. It's a custom RePlay callback for computing recsys metrics on validation and test stages.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 26,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "GPU available: True (cuda), used: True\n",
+      "TPU available: False, using: 0 TPU cores\n",
+      "You are using a CUDA device ('NVIDIA GeForce RTX 3060 Laptop GPU') 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 | Mode  | FLOPs\n",
+      "-------------------------------------------------\n",
+      "0 | model | SasRec | 291 K  | train | 0    \n",
+      "-------------------------------------------------\n",
+      "291 K     Trainable params\n",
+      "0         Non-trainable params\n",
+      "291 K     Total params\n",
+      "1.164     Total estimated model params size (MB)\n",
+      "39        Modules in train mode\n",
+      "0         Modules in eval mode\n",
+      "0         Total Flops\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "bdbe6004cd2b40f09355b6ed047313af",
+       "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": "5f67652c4a59456087990b5eb2a048c6",
+       "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": "cde930c138ce478fb113c82a14740877",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 0, global step 189: 'recall@10' reached 0.01507 (best 0.01507), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=0-step=189-v3.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.001821  0.005005  0.005654  0.004220\n",
+      "ndcg    0.001821  0.007341  0.009693  0.005424\n",
+      "recall  0.001821  0.015069  0.024342  0.009107\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "e2c4de9f0d544a3c9fe936be40e767ea",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 1, global step 378: 'recall@10' reached 0.02302 (best 0.02302), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=1-step=378-v2.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.002981  0.007374  0.008710  0.006044\n",
+      "ndcg    0.002981  0.010976  0.016031  0.007759\n",
+      "recall  0.002981  0.023017  0.043385  0.013082\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "ebeeac2ff096409c9d92c8f500f5a4fa",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 2, global step 567: 'recall@10' reached 0.02981 (best 0.02981), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=2-step=567.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.003809  0.009334  0.011016  0.007628\n",
+      "ndcg    0.003809  0.014037  0.020154  0.009877\n",
+      "recall  0.003809  0.029806  0.053982  0.016890\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "af33cdb8efe947fb96e680d605c57ad8",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 3, global step 756: 'recall@10' reached 0.04554 (best 0.04554), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=3-step=756-v3.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.004637  0.013212  0.015986  0.010153\n",
+      "ndcg    0.004637  0.020583  0.030981  0.013059\n",
+      "recall  0.004637  0.045537  0.087266  0.022024\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "a80a1830e850491ba4bb54053894886b",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 4, global step 945: 'recall@10' reached 0.06143 (best 0.06143), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=4-step=945-v4.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.008611  0.020454  0.024102  0.016998\n",
+      "ndcg    0.008611  0.029885  0.043547  0.021429\n",
+      "recall  0.008611  0.061434  0.116244  0.035105\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "2e60d216a807444693f4880a5d3026d9",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 5, global step 1134: 'recall@10' reached 0.08147 (best 0.08147), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=5-step=1134-v1.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.006789  0.023326  0.027358  0.018168\n",
+      "ndcg    0.006789  0.036654  0.051677  0.023948\n",
+      "recall  0.006789  0.081470  0.141580  0.041729\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "9a55138f358d4264b3f326e2b087c7f7",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 6, global step 1323: 'recall@10' reached 0.09141 (best 0.09141), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=6-step=1323-v2.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.008776  0.027117  0.031929  0.021949\n",
+      "ndcg    0.008776  0.041864  0.059872  0.028990\n",
+      "recall  0.008776  0.091406  0.163603  0.050836\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "3be95e5408e844238aea4dad092d3f6c",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 7, global step 1512: 'recall@10' reached 0.10383 (best 0.10383), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=7-step=1512.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k             1        10        20         5\n",
+      "map     0.01126  0.032126  0.037192  0.026006\n",
+      "ndcg    0.01126  0.048599  0.067473  0.033582\n",
+      "recall  0.01126  0.103825  0.179334  0.056963\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "9fac6ccdf1cf496fa5a226221d9e35ee",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 8, global step 1701: 'recall@10' reached 0.11244 (best 0.11244), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=8-step=1701.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.012088  0.033863  0.039145  0.026379\n",
+      "ndcg    0.012088  0.051835  0.071279  0.033566\n",
+      "recall  0.012088  0.112436  0.189767  0.055638\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "200aaa74132545ebb133abb0f7132f38",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 9, global step 1890: 'recall@10' reached 0.12038 (best 0.12038), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=9-step=1890-v2.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.014075  0.037768  0.043720  0.030035\n",
+      "ndcg    0.014075  0.056689  0.078432  0.037731\n",
+      "recall  0.014075  0.120384  0.206491  0.061268\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "d95e0c88f6664255a963b4dfd25fb420",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 10, global step 2079: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.014572  0.038940  0.045264  0.032039\n",
+      "ndcg    0.014572  0.057550  0.080690  0.040461\n",
+      "recall  0.014572  0.119887  0.211624  0.066236\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "d0e225b1fda2442894cbded20f774fbe",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 11, global step 2268: 'recall@10' reached 0.13313 (best 0.13313), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=11-step=2268.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.014572  0.041118  0.047172  0.032944\n",
+      "ndcg    0.014572  0.062227  0.084644  0.042165\n",
+      "recall  0.014572  0.133135  0.222553  0.070541\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "289a63e4e8d248e9830917c74921b213",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 12, global step 2457: 'recall@10' reached 0.13578 (best 0.13578), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=12-step=2457-v1.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.012916  0.041002  0.048065  0.032950\n",
+      "ndcg    0.012916  0.062827  0.088817  0.043172\n",
+      "recall  0.012916  0.135784  0.239112  0.074681\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "da5831ed697c40a5be21af8ed077c0af",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 13, global step 2646: 'recall@10' reached 0.13761 (best 0.13761), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=13-step=2646.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016393  0.043745  0.049922  0.035814\n",
+      "ndcg    0.016393  0.065313  0.088050  0.045814\n",
+      "recall  0.016393  0.137606  0.228018  0.076668\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "df8c7feb16c040749320c57ac1fb6929",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 14, global step 2835: 'recall@10' reached 0.14572 (best 0.14572), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=14-step=2835.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.013744  0.044572  0.051397  0.036082\n",
+      "ndcg    0.013744  0.067861  0.093062  0.046887\n",
+      "recall  0.013744  0.145719  0.246067  0.079980\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "24ca4ca06969457cbadeb87d9b786a1b",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 15, global step 3024: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.014572  0.043037  0.049347  0.033872\n",
+      "ndcg    0.014572  0.066414  0.089676  0.043945\n",
+      "recall  0.014572  0.145057  0.237622  0.075012\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "9a5ee338856b4205ba96c417dc1cd897",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 16, global step 3213: 'recall@10' reached 0.15301 (best 0.15301), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=16-step=3213.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.020533  0.050175  0.056921  0.041307\n",
+      "ndcg    0.020533  0.073800  0.098705  0.052068\n",
+      "recall  0.020533  0.153005  0.252194  0.085279\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "88de3cabb37446c4bd994c385fb71254",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 17, global step 3402: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016228  0.044922  0.052359  0.035591\n",
+      "ndcg    0.016228  0.068935  0.096304  0.045990\n",
+      "recall  0.016228  0.149859  0.258652  0.078159\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "1424cea483354fcd960f0b2fa686d6ff",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 18, global step 3591: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017056  0.048151  0.055104  0.039250\n",
+      "ndcg    0.017056  0.072077  0.097874  0.050346\n",
+      "recall  0.017056  0.152012  0.255009  0.084451\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "3ec2ba11a2da4485965830c2ef84fef7",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 19, global step 3780: 'recall@10' reached 0.16443 (best 0.16443), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=19-step=3780.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017553  0.050482  0.057555  0.040688\n",
+      "ndcg    0.017553  0.076663  0.102664  0.052659\n",
+      "recall  0.017553  0.164431  0.267760  0.089584\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "019723cee5394e998e598c7a50892aa4",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 20, global step 3969: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.015731  0.048458  0.055842  0.039204\n",
+      "ndcg    0.015731  0.074221  0.101518  0.051256\n",
+      "recall  0.015731  0.160623  0.269415  0.088425\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "6041f647e1434cb3892b2a2f0f3dcecd",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 21, global step 4158: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.018049  0.049611  0.056758  0.040399\n",
+      "ndcg    0.018049  0.074154  0.100495  0.051515\n",
+      "recall  0.018049  0.156317  0.261136  0.085610\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "6ee1a62f7622477bb8648edcd1b92588",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 22, global step 4347: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016393  0.050516  0.057873  0.040887\n",
+      "ndcg    0.016393  0.076624  0.103727  0.053190\n",
+      "recall  0.016393  0.163603  0.271403  0.090909\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "3f8eab96a86443af85c2f62d01a9cf5a",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 23, global step 4536: 'recall@10' reached 0.16559 (best 0.16559), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=23-step=4536.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.014903  0.050339  0.057592  0.040376\n",
+      "ndcg    0.014903  0.076891  0.103476  0.052462\n",
+      "recall  0.014903  0.165590  0.271071  0.089419\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "9c8ded0ade1b4879986fd73a0ebebc2c",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 24, global step 4725: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k             1        10        20         5\n",
+      "map     0.01689  0.049176  0.056484  0.040109\n",
+      "ndcg    0.01689  0.074093  0.100858  0.052004\n",
+      "recall  0.01689  0.157145  0.263289  0.088591\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "c979af1f9bbc4295a23eb8b2472fb96d",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 25, global step 4914: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.018049  0.051505  0.058838  0.042118\n",
+      "ndcg    0.018049  0.077604  0.104511  0.054477\n",
+      "recall  0.018049  0.164928  0.271734  0.092565\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "3d0e8715d4a249f1966939aecd397b8a",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 26, global step 5103: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017718  0.051380  0.059348  0.042295\n",
+      "ndcg    0.017718  0.076988  0.106090  0.054788\n",
+      "recall  0.017718  0.162279  0.277529  0.093227\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "73181d4a5c4e4013b7a8eeebf742ce9e",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 27, global step 5292: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017221  0.048906  0.057236  0.039648\n",
+      "ndcg    0.017221  0.074121  0.104692  0.051467\n",
+      "recall  0.017221  0.158470  0.279848  0.087928\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "1634e7c2003d44d697030ab44288f0b1",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 28, global step 5481: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k             1        10        20         5\n",
+      "map     0.01954  0.051349  0.059358  0.041872\n",
+      "ndcg    0.01954  0.076578  0.106111  0.053326\n",
+      "recall  0.01954  0.161119  0.278689  0.088591\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "908c27a408b84ef7b6c91700b34f1f0d",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 29, global step 5670: 'recall@10' reached 0.17221 (best 0.17221), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=29-step=5670.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.021527  0.055773  0.063681  0.046164\n",
+      "ndcg    0.021527  0.082563  0.111508  0.058804\n",
+      "recall  0.021527  0.172214  0.286968  0.097698\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "b51365e160ab4865a27cf655e17707a5",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 30, global step 5859: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.018877  0.053691  0.061561  0.044171\n",
+      "ndcg    0.018877  0.080600  0.109718  0.057056\n",
+      "recall  0.018877  0.170558  0.286637  0.096705\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "8547295f7ccb4fd790366aab9fe9de78",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 31, global step 6048: 'recall@10' reached 0.17404 (best 0.17404), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=31-step=6048.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.014903  0.050917  0.058470  0.040520\n",
+      "ndcg    0.014903  0.079214  0.107110  0.053566\n",
+      "recall  0.014903  0.174035  0.285147  0.093724\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "71f02a600023460aa36238e118f89453",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 32, global step 6237: 'recall@10' reached 0.17553 (best 0.17553), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=32-step=6237.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.018049  0.053705  0.061464  0.043180\n",
+      "ndcg    0.018049  0.081750  0.110451  0.056098\n",
+      "recall  0.018049  0.175526  0.289949  0.095877\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "a94c40bad33b4fef92fbe3b7ca2dd7a9",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 33, global step 6426: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.019208  0.053874  0.061629  0.043597\n",
+      "ndcg    0.019208  0.081296  0.109911  0.056049\n",
+      "recall  0.019208  0.173207  0.287134  0.094386\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "93e62f688cae4954826ce5f65bc33ea6",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 34, global step 6615: 'recall@10' reached 0.18016 (best 0.18016), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=34-step=6615.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017553  0.054073  0.061606  0.043559\n",
+      "ndcg    0.017553  0.083073  0.111055  0.057211\n",
+      "recall  0.017553  0.180162  0.291936  0.099354\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "4bf8a5210e174568bd46212f9ea3fa7c",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 35, global step 6804: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.018215  0.055078  0.062641  0.044364\n",
+      "ndcg    0.018215  0.082983  0.110672  0.056984\n",
+      "recall  0.018215  0.176023  0.285809  0.095546\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "c43ebbf0977a431d8a961330f45570c5",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 36, global step 6993: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017221  0.052581  0.060386  0.042071\n",
+      "ndcg    0.017221  0.080544  0.109221  0.054559\n",
+      "recall  0.017221  0.174367  0.288293  0.092896\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "2d4241ed36f44607be6c6153931ac151",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 37, global step 7182: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017884  0.053871  0.061073  0.043766\n",
+      "ndcg    0.017884  0.082064  0.108563  0.057241\n",
+      "recall  0.017884  0.176354  0.281669  0.098857\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "aba0b7085a074368884b0651c8a2a0b3",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 38, global step 7371: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.018546  0.054212  0.061996  0.044508\n",
+      "ndcg    0.018546  0.081737  0.110419  0.057725\n",
+      "recall  0.018546  0.173704  0.287796  0.098361\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "feba795a523446eb8e5e6cd96bad39ba",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 39, global step 7560: 'recall@10' reached 0.18082 (best 0.18082), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=39-step=7560.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k             1        10        20         5\n",
+      "map     0.02103  0.057426  0.064922  0.046829\n",
+      "ndcg    0.02103  0.085795  0.113317  0.059729\n",
+      "recall  0.02103  0.180825  0.290114  0.099354\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "83bf44f90b024f2aa4f5b69bbc18b023",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 40, global step 7749: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.020699  0.055658  0.063230  0.044966\n",
+      "ndcg    0.020699  0.083080  0.111028  0.057106\n",
+      "recall  0.020699  0.174863  0.286140  0.094386\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "4abdd8633d3e453d93a2af9a1d6a9ccc",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 41, global step 7938: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016725  0.052379  0.060527  0.042126\n",
+      "ndcg    0.016725  0.079931  0.110088  0.054617\n",
+      "recall  0.016725  0.172214  0.292433  0.092896\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "b20c77f4d192498d81c2ed6baa4c811d",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 42, global step 8127: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.021692  0.058423  0.066561  0.048173\n",
+      "ndcg    0.021692  0.086242  0.116025  0.061053\n",
+      "recall  0.021692  0.179169  0.297235  0.100513\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "c4929ab9384a4a6ba9bf29a9c939b1ab",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 43, global step 8316: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.019705  0.055943  0.063840  0.045391\n",
+      "ndcg    0.019705  0.084424  0.113300  0.058431\n",
+      "recall  0.019705  0.179831  0.294254  0.098526\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "9ffbd31684ad4a6aa515233f9e8ef5aa",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 44, global step 8505: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.015069  0.053308  0.061374  0.043291\n",
+      "ndcg    0.015069  0.081830  0.111383  0.057192\n",
+      "recall  0.015069  0.176850  0.294088  0.099851\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "e6f025667c594ba0aa1029fc126858f2",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 45, global step 8694: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.018049  0.055185  0.062849  0.044538\n",
+      "ndcg    0.018049  0.083015  0.111301  0.057084\n",
+      "recall  0.018049  0.175857  0.288458  0.095380\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "ab51906b826a401396d92dd47c51028d",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 46, global step 8883: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017387  0.054915  0.062717  0.044329\n",
+      "ndcg    0.017387  0.083518  0.112327  0.057648\n",
+      "recall  0.017387  0.179003  0.293757  0.098526\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "e68d1d8754fa43958a9f7eb8b86ddfab",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 47, global step 9072: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017553  0.055809  0.064044  0.045537\n",
+      "ndcg    0.017553  0.084333  0.114317  0.059376\n",
+      "recall  0.017553  0.179169  0.297731  0.101838\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "dd0ec52c482f4fc789b0458eb0c2da5d",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 48, global step 9261: 'recall@10' reached 0.18596 (best 0.18596), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=48-step=9261.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.022355  0.060448  0.068135  0.049738\n",
+      "ndcg    0.022355  0.089385  0.117671  0.063264\n",
+      "recall  0.022355  0.185958  0.298394  0.104819\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "1f2fd19e798845c8959bd278a1e47c99",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 49, global step 9450: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.020533  0.058365  0.065847  0.047737\n",
+      "ndcg    0.020533  0.087662  0.115249  0.061533\n",
+      "recall  0.020533  0.185627  0.295413  0.103991\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "51c229a5a2ac4551a01b109eee0f72c9",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 50, global step 9639: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.019705  0.057252  0.065311  0.046550\n",
+      "ndcg    0.019705  0.086161  0.115841  0.059957\n",
+      "recall  0.019705  0.182812  0.300878  0.101176\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "c1bf323886a147aab5eb81a49910d99c",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 51, global step 9828: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.019871  0.056373  0.064782  0.046465\n",
+      "ndcg    0.019871  0.084487  0.115268  0.060018\n",
+      "recall  0.019871  0.178341  0.300381  0.101672\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "5a0a40c3297446219f782ccaddb48f36",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 52, global step 10017: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017056  0.055635  0.063238  0.044657\n",
+      "ndcg    0.017056  0.085148  0.113217  0.058444\n",
+      "recall  0.017056  0.183640  0.295413  0.100845\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "66e75dd4c62d419e8b41a635baeea9b6",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 53, global step 10206: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016725  0.056399  0.064312  0.046448\n",
+      "ndcg    0.016725  0.085609  0.114614  0.061160\n",
+      "recall  0.016725  0.182646  0.297731  0.106309\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "7b552a2b6ae343ab914f022626199df5",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 54, global step 10395: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.020864  0.056358  0.064559  0.045976\n",
+      "ndcg    0.020864  0.084551  0.114658  0.059012\n",
+      "recall  0.020864  0.179003  0.298559  0.099189\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "7d7d5e67c4264a8d936a886955a1e468",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 55, global step 10584: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.015897  0.052936  0.060713  0.042005\n",
+      "ndcg    0.015897  0.082446  0.111144  0.055636\n",
+      "recall  0.015897  0.181321  0.295579  0.097698\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "567aa7273f024f0496d7b105c0e5745f",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 56, global step 10773: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.018215  0.055970  0.064171  0.045659\n",
+      "ndcg    0.018215  0.084831  0.115129  0.059707\n",
+      "recall  0.018215  0.180990  0.301706  0.102997\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "effe5af26f79489f9fe7ced8ee84dbd6",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 57, global step 10962: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.019705  0.057038  0.065075  0.046796\n",
+      "ndcg    0.019705  0.085973  0.115791  0.060577\n",
+      "recall  0.019705  0.182812  0.301871  0.102997\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "5c3359ce59cd45e883e49fd6b07b642d",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 58, global step 11151: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.015731  0.054501  0.062389  0.043956\n",
+      "ndcg    0.015731  0.084609  0.113613  0.058502\n",
+      "recall  0.015731  0.185296  0.300546  0.103328\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "77112b36b0154d3cb233771cac236769",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 59, global step 11340: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.019705  0.057285  0.065306  0.046415\n",
+      "ndcg    0.019705  0.086774  0.116052  0.060147\n",
+      "recall  0.019705  0.185461  0.301374  0.102500\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "32f38bfe60b24b90b0a6d006a0babd31",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 60, global step 11529: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.018049  0.056947  0.064405  0.046498\n",
+      "ndcg    0.018049  0.086052  0.113448  0.060260\n",
+      "recall  0.018049  0.183143  0.291936  0.102335\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "f094d09d3e7b453480b5b485f25692f3",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 61, global step 11718: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017553  0.057223  0.064880  0.046870\n",
+      "ndcg    0.017553  0.086840  0.115219  0.061361\n",
+      "recall  0.017553  0.185461  0.298725  0.105812\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "f4948db841b44c168157b0b11e79efa4",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 62, global step 11907: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017718  0.056284  0.064414  0.045866\n",
+      "ndcg    0.017718  0.085691  0.115490  0.060021\n",
+      "recall  0.017718  0.183805  0.302037  0.103494\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "5fc0767f2e6a4b4ea45162db8c71ad1c",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 63, global step 12096: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.014075  0.053323  0.061450  0.042706\n",
+      "ndcg    0.014075  0.083044  0.113033  0.057016\n",
+      "recall  0.014075  0.182149  0.301540  0.101010\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "ffa03ec64fe842c2a11b1d4b2336c1b4",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 64, global step 12285: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016559  0.055950  0.063635  0.045606\n",
+      "ndcg    0.016559  0.085525  0.113860  0.060119\n",
+      "recall  0.016559  0.183971  0.296738  0.104653\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "1da052c1711a48eb9179718b1e8075cf",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 65, global step 12474: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017718  0.056836  0.065167  0.046782\n",
+      "ndcg    0.017718  0.085591  0.116301  0.060821\n",
+      "recall  0.017718  0.181321  0.303527  0.103825\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "b3993c1d275a4924b6731e4499bde5c6",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 66, global step 12663: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.015897  0.054146  0.063006  0.043884\n",
+      "ndcg    0.015897  0.083296  0.116044  0.058099\n",
+      "recall  0.015897  0.180493  0.310979  0.101838\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "832191caf90d44c7b8e8223afb823b76",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 67, global step 12852: 'recall@10' reached 0.19159 (best 0.19159), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=67-step=12852.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017056  0.057459  0.065661  0.045971\n",
+      "ndcg    0.017056  0.088389  0.118391  0.060327\n",
+      "recall  0.017056  0.191588  0.310482  0.104322\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "12bcca54c067420494528199338a1e56",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 68, global step 13041: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017718  0.058559  0.066859  0.048270\n",
+      "ndcg    0.017718  0.088819  0.119630  0.063497\n",
+      "recall  0.017718  0.189435  0.312469  0.110283\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "361a005333cf4228bd88446f5d3c7efd",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 69, global step 13230: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017387  0.057117  0.065430  0.046705\n",
+      "ndcg    0.017387  0.087054  0.117606  0.061578\n",
+      "recall  0.017387  0.186620  0.307998  0.107303\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "816b4845818f4e1ab83abd1cef17d94d",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 70, global step 13419: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.019043  0.057870  0.066302  0.047017\n",
+      "ndcg    0.019043  0.087867  0.118650  0.061382\n",
+      "recall  0.019043  0.187945  0.309820  0.105647\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "a7c54580cab949ed9162e5e2163f5144",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 71, global step 13608: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.015731  0.054928  0.063489  0.044635\n",
+      "ndcg    0.015731  0.084751  0.116133  0.059363\n",
+      "recall  0.015731  0.184136  0.308660  0.104653\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "cd4ae17c40d34d28a0d326139782a376",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 72, global step 13797: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.018381  0.056898  0.065084  0.046131\n",
+      "ndcg    0.018381  0.086764  0.116853  0.060242\n",
+      "recall  0.018381  0.186620  0.306177  0.103660\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "cb7d5821d4b24d928d834ac8f92c37b6",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 73, global step 13986: 'recall@10' reached 0.19308 (best 0.19308), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=73-step=13986.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k             1        10        20         5\n",
+      "map     0.01954  0.060165  0.068179  0.048888\n",
+      "ndcg    0.01954  0.090827  0.120453  0.063258\n",
+      "recall  0.01954  0.193078  0.311144  0.107303\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "8c82b678f2374af7859f13db93cc66a1",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 74, global step 14175: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.020202  0.059898  0.068066  0.049208\n",
+      "ndcg    0.020202  0.090335  0.120156  0.063797\n",
+      "recall  0.020202  0.192085  0.310151  0.108627\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "bff6682c9771473eba533531e1d63df0",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 75, global step 14364: 'recall@10' reached 0.19540 (best 0.19540), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=75-step=14364.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017553  0.059239  0.067168  0.047867\n",
+      "ndcg    0.017553  0.090614  0.119642  0.062577\n",
+      "recall  0.017553  0.195397  0.310482  0.107634\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "56513ed0de4f488e8a67ddab39bff555",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 76, global step 14553: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017553  0.057571  0.066031  0.046285\n",
+      "ndcg    0.017553  0.088607  0.119752  0.060844\n",
+      "recall  0.017553  0.192416  0.316278  0.105647\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "6c21be0f8bf64ff5a9a0433140db4afb",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 77, global step 14742: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k             1        10        20         5\n",
+      "map     0.02252  0.059148  0.067676  0.048570\n",
+      "ndcg    0.02252  0.087980  0.119156  0.062026\n",
+      "recall  0.02252  0.184468  0.307998  0.103494\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "67f774849e614a5e951e72ddc8847b0d",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 78, global step 14931: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.020864  0.059902  0.068102  0.048949\n",
+      "ndcg    0.020864  0.089627  0.119711  0.062921\n",
+      "recall  0.020864  0.188773  0.308164  0.105812\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "489fccb9e3774e5e9924cb0ae484229a",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 79, global step 15120: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.018049  0.057458  0.065510  0.046583\n",
+      "ndcg    0.018049  0.087580  0.117393  0.060762\n",
+      "recall  0.018049  0.188276  0.307170  0.104322\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "752a04f3c8ce4e36a5675c8a76c95503",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 80, global step 15309: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.018546  0.057481  0.065625  0.046603\n",
+      "ndcg    0.018546  0.087740  0.117695  0.061046\n",
+      "recall  0.018546  0.188773  0.307832  0.105481\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "9e682711164c445ba54cfe4a84262d8c",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 81, global step 15498: 'recall@10' reached 0.19672 (best 0.19672), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=81-step=15498.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.019705  0.060155  0.067919  0.048319\n",
+      "ndcg    0.019705  0.091545  0.120256  0.062412\n",
+      "recall  0.019705  0.196721  0.311144  0.105647\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "f0da4eba218f4fffb4deb95efa8df791",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 82, global step 15687: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.018381  0.057389  0.065948  0.046868\n",
+      "ndcg    0.018381  0.087341  0.118833  0.061256\n",
+      "recall  0.018381  0.187448  0.312635  0.105481\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "1c626b061743416287ba6fc1aa65e8c6",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 83, global step 15876: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017387  0.057521  0.065933  0.046556\n",
+      "ndcg    0.017387  0.088367  0.119428  0.061458\n",
+      "recall  0.017387  0.191257  0.314953  0.107303\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "fd43c2909b574172a1726662394f32f7",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 84, global step 16065: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.019374  0.058125  0.066462  0.047210\n",
+      "ndcg    0.019374  0.087957  0.118482  0.061159\n",
+      "recall  0.019374  0.187614  0.308660  0.103991\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "33d3e2b831274e0b9783039482899584",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 85, global step 16254: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016725  0.055607  0.064210  0.044947\n",
+      "ndcg    0.016725  0.085456  0.117166  0.059262\n",
+      "recall  0.016725  0.185130  0.311310  0.103328\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "45f8fd0d5bd34318afd47a5b5b700d20",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 86, global step 16443: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016559  0.057457  0.065991  0.046661\n",
+      "ndcg    0.016559  0.087773  0.119155  0.061410\n",
+      "recall  0.016559  0.188607  0.313297  0.106640\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "bc750dde2e4d42eb8cf7655a39a465df",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 87, global step 16632: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.015731  0.057037  0.065072  0.045678\n",
+      "ndcg    0.015731  0.088516  0.118281  0.060867\n",
+      "recall  0.015731  0.193410  0.312138  0.107634\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "9b25c05436d24e3d80a56010edcff25b",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 88, global step 16821: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.018712  0.059290  0.067419  0.048380\n",
+      "ndcg    0.018712  0.090193  0.120127  0.063372\n",
+      "recall  0.018712  0.193244  0.312303  0.109455\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "89bed0fc3db14129badb8168c668c450",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 89, global step 17010: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.019208  0.060027  0.067680  0.048468\n",
+      "ndcg    0.019208  0.091291  0.119452  0.062846\n",
+      "recall  0.019208  0.195893  0.307832  0.106971\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "a66ad86b8b664ca6a4defd3c7cf97be1",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 90, global step 17199: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017553  0.059004  0.066841  0.047811\n",
+      "ndcg    0.017553  0.089908  0.118707  0.062456\n",
+      "recall  0.017553  0.192913  0.307336  0.107303\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "e6f035c7845b41f28da5d971ad4fc92f",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 91, global step 17388: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.019871  0.059618  0.068323  0.048620\n",
+      "ndcg    0.019871  0.089395  0.121484  0.062390\n",
+      "recall  0.019871  0.188773  0.316443  0.104488\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "782c0bc4a498460ca213a3930447fdc1",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 92, global step 17577: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016393  0.056838  0.065549  0.046489\n",
+      "ndcg    0.016393  0.087088  0.118991  0.061682\n",
+      "recall  0.016393  0.187779  0.314290  0.108462\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "448e00539d8c4ebcba2e51d0b5dbbb9e",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 93, global step 17766: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.020864  0.060046  0.068195  0.049431\n",
+      "ndcg    0.020864  0.090098  0.120281  0.063916\n",
+      "recall  0.020864  0.190429  0.310813  0.108462\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "2fea31da94254d8d8f7b87f88743a754",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 94, global step 17955: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.018215  0.059519  0.067797  0.048775\n",
+      "ndcg    0.018215  0.090521  0.120919  0.064276\n",
+      "recall  0.018215  0.193575  0.314290  0.111939\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "310abc5fa09f4f4597dba48212181900",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 95, global step 18144: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016393  0.057420  0.065561  0.045932\n",
+      "ndcg    0.016393  0.088958  0.118821  0.060803\n",
+      "recall  0.016393  0.194237  0.312800  0.106475\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "b3c83c7aa7c74abb8db94651d34e8a98",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 96, global step 18333: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016725  0.056762  0.065223  0.045708\n",
+      "ndcg    0.016725  0.087184  0.118141  0.060252\n",
+      "recall  0.016725  0.188607  0.311310  0.104984\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "54598b8e0f5b4bd3a2ed8a7c2ccbc8ae",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 97, global step 18522: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016228  0.056993  0.065326  0.045982\n",
+      "ndcg    0.016228  0.087804  0.118434  0.060850\n",
+      "recall  0.016228  0.190429  0.312138  0.106475\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "836e2a6ef8564b2bbaa9311a90e8817d",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 98, global step 18711: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016393  0.056134  0.064101  0.044497\n",
+      "ndcg    0.016393  0.087227  0.116437  0.058923\n",
+      "recall  0.016393  0.191091  0.307004  0.103328\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "e6338f1e7e1f409b9bb55d350f538a20",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 99, global step 18900: 'recall@10' was not in top 1\n",
+      "`Trainer.fit` stopped: `max_epochs=100` reached.\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017884  0.058572  0.066894  0.048223\n",
+      "ndcg    0.017884  0.088838  0.119484  0.063391\n",
+      "recall  0.017884  0.189435  0.311310  0.109952\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "from lightning.pytorch.callbacks import ModelCheckpoint\n",
+    "from lightning.pytorch.loggers import CSVLogger\n",
+    "\n",
+    "from replay.nn.lightning.callback import ComputeMetricsCallback\n",
+    "\n",
+    "\n",
+    "checkpoint_callback = ModelCheckpoint(\n",
+    "    dirpath=\"sasrec/checkpoints\",\n",
+    "    save_top_k=1,\n",
+    "    verbose=True,\n",
+    "    monitor=\"recall@10\",\n",
+    "    mode=\"max\",\n",
+    ")\n",
+    "\n",
+    "validation_metrics_callback = ComputeMetricsCallback(\n",
+    "    metrics=[\"map\", \"ndcg\", \"recall\"],\n",
+    "    ks=[1, 5, 10, 20],\n",
+    "    item_count=NUM_UNIQUE_ITEMS,\n",
+    ")\n",
+    "\n",
+    "csv_logger = CSVLogger(save_dir=\"sasrec/logs/train\", name=\"SasRec-example\")\n",
+    "\n",
+    "trainer = L.Trainer(\n",
+    "    max_epochs=100,\n",
+    "    callbacks=[checkpoint_callback, validation_metrics_callback],\n",
+    "    logger=csv_logger,\n",
+    ")\n",
+    "\n",
+    "trainer.fit(model, datamodule=parquet_module)"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Now we can get the best model path stored in the checkpoint callback."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 27,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "'/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=81-step=15498.ckpt'"
+      ]
+     },
+     "execution_count": 27,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "best_model_path = checkpoint_callback.best_model_path\n",
+    "best_model_path"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Inference\n",
+    "\n",
+    "To obtain model scores, we will load the weights from the best checkpoint. To do this, we use the `LightningModule`, provide there the path to the checkpoint and the model instance."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 33,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import replay\n",
+    "torch.serialization.add_safe_globals([\n",
+    "    replay.nn.lightning.optimizer.OptimizerFactory,\n",
+    "    replay.nn.lightning.scheduler.LRSchedulerFactory\n",
+    "])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 34,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sasrec = make_sasrec(\n",
+    "    schema=tensor_schema,\n",
+    "    embedding_dim=EMBEDDING_DIM,\n",
+    "    max_sequence_length=MAX_SEQ_LEN,\n",
+    "    num_heads=NUM_HEADS,\n",
+    "    num_blocks=NUM_BLOCKS,\n",
+    "    dropout=DROPOUT,\n",
+    ")\n",
+    "\n",
+    "best_model = LightningModule.load_from_checkpoint(best_model_path, model=sasrec)\n",
+    "best_model.eval();"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Configure `ParquetModule` for inference"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 44,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "inference_metadata = {\"predict\": create_meta(shape=MAX_SEQ_LEN)}\n",
+    "\n",
+    "parquet_module = ParquetModule(\n",
+    "    predict_path=PREDICT_PATH,\n",
+    "    batch_size=BATCH_SIZE,\n",
+    "    metadata=inference_metadata,\n",
+    "    transforms=transforms,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "During inference, we can use `TopItemsCallback`. Such callback allows you to get scores for each user throughout the entire catalog and get recommendations in the form of ids of items with the highest score values.\n",
+    "\n",
+    "\n",
+    "Recommendations can be fetched in four formats: PySpark DataFrame, Pandas DataFrame, Polars DataFrame or raw PyTorch tensors. Each of the types corresponds a callback. In this example, we'll be using the `PandasTopItemsCallback`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 46,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "💡 Tip: For seamless cloud uploads and versioning, try installing [litmodels](https://pypi.org/project/litmodels/) to enable LitModelCheckpoint, which syncs automatically with the Lightning model registry.\n",
+      "GPU available: True (cuda), used: True\n",
+      "TPU available: False, using: 0 TPU cores\n",
+      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "cd1373dad90640fe8cbc523662ccb563",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Predicting: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "from replay.nn.lightning.callback import PandasTopItemsCallback\n",
+    "\n",
+    "csv_logger = CSVLogger(save_dir=\"sasrec/logs/test\", name=\"SasRec-example\")\n",
+    "\n",
+    "TOPK = [1, 5, 10, 20]\n",
+    "\n",
+    "pandas_prediction_callback = PandasTopItemsCallback(\n",
+    "    top_k=max(TOPK),\n",
+    "    query_column=\"user_id\",\n",
+    "    item_column=\"item_id\",\n",
+    "    rating_column=\"score\",\n",
+    ")\n",
+    "\n",
+    "trainer = L.Trainer(callbacks=[pandas_prediction_callback], logger=csv_logger, inference_mode=True)\n",
+    "trainer.predict(best_model, datamodule=parquet_module, return_predictions=False)\n",
+    "\n",
+    "pandas_res = pandas_prediction_callback.get_result()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 47,
+   "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>224</td>\n",
+       "      <td>7.242342</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>0</td>\n",
+       "      <td>572</td>\n",
+       "      <td>6.818249</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>0</td>\n",
+       "      <td>486</td>\n",
+       "      <td>6.81148</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>0</td>\n",
+       "      <td>1371</td>\n",
+       "      <td>6.534966</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>0</td>\n",
+       "      <td>210</td>\n",
+       "      <td>6.52649</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>...</th>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6037</th>\n",
+       "      <td>6039</td>\n",
+       "      <td>2497</td>\n",
+       "      <td>10.457304</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6037</th>\n",
+       "      <td>6039</td>\n",
+       "      <td>3503</td>\n",
+       "      <td>10.305973</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6037</th>\n",
+       "      <td>6039</td>\n",
+       "      <td>2601</td>\n",
+       "      <td>10.280416</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6037</th>\n",
+       "      <td>6039</td>\n",
+       "      <td>2750</td>\n",
+       "      <td>10.01198</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6037</th>\n",
+       "      <td>6039</td>\n",
+       "      <td>1490</td>\n",
+       "      <td>9.916577</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "<p>120760 rows × 3 columns</p>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "      user_id item_id      score\n",
+       "0           0     224   7.242342\n",
+       "0           0     572   6.818249\n",
+       "0           0     486    6.81148\n",
+       "0           0    1371   6.534966\n",
+       "0           0     210    6.52649\n",
+       "...       ...     ...        ...\n",
+       "6037     6039    2497  10.457304\n",
+       "6037     6039    3503  10.305973\n",
+       "6037     6039    2601  10.280416\n",
+       "6037     6039    2750   10.01198\n",
+       "6037     6039    1490   9.916577\n",
+       "\n",
+       "[120760 rows x 3 columns]"
+      ]
+     },
+     "execution_count": 47,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "pandas_res"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Calculating metrics\n",
+    "\n",
+    "*test_gt* is already encoded, so we can use it for computing metrics."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 48,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from replay.metrics import MAP, OfflineMetrics, Precision, Recall\n",
+    "from replay.metrics.torch_metrics_builder import metrics_to_df"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 49,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "result_metrics = OfflineMetrics(\n",
+    "    [Recall(TOPK), Precision(TOPK), MAP(TOPK)],\n",
+    "    query_column=\"user_id\",\n",
+    "    rating_column=\"score\",\n",
+    ")(pandas_res, test_gt.explode(\"item_id\"))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 50,
+   "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>5</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>MAP</th>\n",
+       "      <td>0.016065</td>\n",
+       "      <td>0.054039</td>\n",
+       "      <td>0.061749</td>\n",
+       "      <td>0.043969</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>Precision</th>\n",
+       "      <td>0.016065</td>\n",
+       "      <td>0.017655</td>\n",
+       "      <td>0.014458</td>\n",
+       "      <td>0.020073</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>Recall</th>\n",
+       "      <td>0.016065</td>\n",
+       "      <td>0.176549</td>\n",
+       "      <td>0.289169</td>\n",
+       "      <td>0.100364</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "k                 1        10        20         5\n",
+       "MAP        0.016065  0.054039  0.061749  0.043969\n",
+       "Precision  0.016065  0.017655  0.014458  0.020073\n",
+       "Recall     0.016065  0.176549  0.289169  0.100364"
+      ]
+     },
+     "execution_count": 50,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "metrics_to_df(result_metrics)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Let's call the `inverse_transform` encoder's function to get the final dataframe with recommendations"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 51,
+   "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>2012</td>\n",
+       "      <td>231</td>\n",
+       "      <td>7.242342</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>2012</td>\n",
+       "      <td>586</td>\n",
+       "      <td>6.818249</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>2012</td>\n",
+       "      <td>500</td>\n",
+       "      <td>6.81148</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>2012</td>\n",
+       "      <td>1485</td>\n",
+       "      <td>6.534966</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>2012</td>\n",
+       "      <td>216</td>\n",
+       "      <td>6.52649</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>...</th>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6037</th>\n",
+       "      <td>5727</td>\n",
+       "      <td>2702</td>\n",
+       "      <td>10.457304</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6037</th>\n",
+       "      <td>5727</td>\n",
+       "      <td>3745</td>\n",
+       "      <td>10.305973</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6037</th>\n",
+       "      <td>5727</td>\n",
+       "      <td>2806</td>\n",
+       "      <td>10.280416</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6037</th>\n",
+       "      <td>5727</td>\n",
+       "      <td>2961</td>\n",
+       "      <td>10.01198</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6037</th>\n",
+       "      <td>5727</td>\n",
+       "      <td>1623</td>\n",
+       "      <td>9.916577</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "<p>120760 rows × 3 columns</p>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "      user_id  item_id      score\n",
+       "0        2012      231   7.242342\n",
+       "0        2012      586   6.818249\n",
+       "0        2012      500    6.81148\n",
+       "0        2012     1485   6.534966\n",
+       "0        2012      216    6.52649\n",
+       "...       ...      ...        ...\n",
+       "6037     5727     2702  10.457304\n",
+       "6037     5727     3745  10.305973\n",
+       "6037     5727     2806  10.280416\n",
+       "6037     5727     2961   10.01198\n",
+       "6037     5727     1623   9.916577\n",
+       "\n",
+       "[120760 rows x 3 columns]"
+      ]
+     },
+     "execution_count": 51,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "encoder.inverse_transform(pandas_res)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "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.11.14"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/examples/09_sasrec_example_thr_sampled.ipynb b/examples/09_sasrec_example_thr_sampled.ipynb
new file mode 100644
index 000000000..b27ec3f8d
--- /dev/null
+++ b/examples/09_sasrec_example_thr_sampled.ipynb
@@ -0,0 +1,4768 @@
+{
+ "cells": [
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Example of SasRec training/inference with Parquet Module"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Seed set to 42\n"
+     ]
+    }
+   ],
+   "source": [
+    "from typing import Optional\n",
+    "\n",
+    "import lightning as L\n",
+    "import pandas as pd\n",
+    "\n",
+    "L.seed_everything(42)\n",
+    "\n",
+    "import warnings\n",
+    "warnings.filterwarnings(\"ignore\")"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Preparing data\n",
+    "In this example, we will be using the MovieLens dataset, namely the 1m subset. It's demonstrated a simple case, so only item ids will be used as model input.\n",
+    "\n",
+    "---\n",
+    "**NOTE**\n",
+    "\n",
+    "Current implementation of SasRec is able to handle item and interactions features. It does not take into account user features. \n",
+    "\n",
+    "---"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "interactions = pd.read_csv(\"./data/ml1m_ratings.dat\", sep=\"\\t\", names=[\"user_id\", \"item_id\",\"rating\",\"timestamp\"])\n",
+    "interactions = interactions.drop(columns=[\"rating\"])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "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>timestamp</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>1000138</th>\n",
+       "      <td>6040</td>\n",
+       "      <td>858</td>\n",
+       "      <td>0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1000153</th>\n",
+       "      <td>6040</td>\n",
+       "      <td>2384</td>\n",
+       "      <td>1</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>999873</th>\n",
+       "      <td>6040</td>\n",
+       "      <td>593</td>\n",
+       "      <td>2</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1000007</th>\n",
+       "      <td>6040</td>\n",
+       "      <td>1961</td>\n",
+       "      <td>3</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1000192</th>\n",
+       "      <td>6040</td>\n",
+       "      <td>2019</td>\n",
+       "      <td>4</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>...</th>\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>446</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>825438</th>\n",
+       "      <td>4958</td>\n",
+       "      <td>1407</td>\n",
+       "      <td>447</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>825724</th>\n",
+       "      <td>4958</td>\n",
+       "      <td>3264</td>\n",
+       "      <td>448</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>825731</th>\n",
+       "      <td>4958</td>\n",
+       "      <td>2634</td>\n",
+       "      <td>449</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>825603</th>\n",
+       "      <td>4958</td>\n",
+       "      <td>1924</td>\n",
+       "      <td>450</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "<p>1000209 rows × 3 columns</p>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "         user_id  item_id  timestamp\n",
+       "1000138     6040      858          0\n",
+       "1000153     6040     2384          1\n",
+       "999873      6040      593          2\n",
+       "1000007     6040     1961          3\n",
+       "1000192     6040     2019          4\n",
+       "...          ...      ...        ...\n",
+       "825793      4958     2399        446\n",
+       "825438      4958     1407        447\n",
+       "825724      4958     3264        448\n",
+       "825731      4958     2634        449\n",
+       "825603      4958     1924        450\n",
+       "\n",
+       "[1000209 rows x 3 columns]"
+      ]
+     },
+     "execution_count": 3,
+     "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"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Encode catagorical data.\n",
+    "To ensure all categorical data is fit for training, it needs to be encoded using the `LabelEncoder` class. Create an instance of the encoder, providing a `LabelEncodingRule` for each categorcial column in the dataset that will be used in model. Note that ids of users and ids of items are always used."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "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>timestamp</th>\n",
+       "      <th>user_id</th>\n",
+       "      <th>item_id</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>12</td>\n",
+       "      <td>0</td>\n",
+       "      <td>0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>68</td>\n",
+       "      <td>1</td>\n",
+       "      <td>0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>67</td>\n",
+       "      <td>2</td>\n",
+       "      <td>0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>12</td>\n",
+       "      <td>3</td>\n",
+       "      <td>0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>140</td>\n",
+       "      <td>4</td>\n",
+       "      <td>0</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>...</th>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1000204</th>\n",
+       "      <td>14</td>\n",
+       "      <td>4555</td>\n",
+       "      <td>3705</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1000205</th>\n",
+       "      <td>90</td>\n",
+       "      <td>2813</td>\n",
+       "      <td>3705</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1000206</th>\n",
+       "      <td>70</td>\n",
+       "      <td>2404</td>\n",
+       "      <td>3705</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1000207</th>\n",
+       "      <td>25</td>\n",
+       "      <td>5835</td>\n",
+       "      <td>3705</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1000208</th>\n",
+       "      <td>380</td>\n",
+       "      <td>979</td>\n",
+       "      <td>3705</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "<p>1000209 rows × 3 columns</p>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "         timestamp  user_id  item_id\n",
+       "0               12        0        0\n",
+       "1               68        1        0\n",
+       "2               67        2        0\n",
+       "3               12        3        0\n",
+       "4              140        4        0\n",
+       "...            ...      ...      ...\n",
+       "1000204         14     4555     3705\n",
+       "1000205         90     2813     3705\n",
+       "1000206         70     2404     3705\n",
+       "1000207         25     5835     3705\n",
+       "1000208        380      979     3705\n",
+       "\n",
+       "[1000209 rows x 3 columns]"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from replay.preprocessing.label_encoder import LabelEncoder, LabelEncodingRule\n",
+    "\n",
+    "encoder = LabelEncoder(\n",
+    "    [\n",
+    "        LabelEncodingRule(\"user_id\", default_value=\"last\"),\n",
+    "        LabelEncodingRule(\"item_id\", default_value=\"last\"),\n",
+    "    ]\n",
+    ")\n",
+    "interactions = interactions.sort_values(by=\"item_id\", ascending=True)\n",
+    "encoded_interactions = encoder.fit_transform(interactions)\n",
+    "encoded_interactions"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Split interactions into the train, validation and test datasets using LastNSplitter\n",
+    "We use widespread splitting strategy Last-One-Out. We filter out cold items and users for simplicity."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from replay.splitters import LastNSplitter\n",
+    "\n",
+    "splitter = LastNSplitter(\n",
+    "    N=1,\n",
+    "    divide_column=\"user_id\",\n",
+    "    query_column=\"user_id\",\n",
+    "    strategy=\"interactions\",\n",
+    "    drop_cold_users=True,\n",
+    "    drop_cold_items=True\n",
+    ")\n",
+    "\n",
+    "test_events, test_gt = splitter.split(encoded_interactions)\n",
+    "validation_events, validation_gt = splitter.split(test_events)\n",
+    "train_events = validation_events"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Dataset preprocessing (\"baking\")\n",
+    "SasRec expects each user in the batch to provide their events in form of a sequence. For this reason, the event splits must be properly processed using the `groupby_sequences` function provided by RePlay."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from replay.data.nn.utils import groupby_sequences\n",
+    "\n",
+    "\n",
+    "def bake_data(full_data):\n",
+    "    grouped_interactions = groupby_sequences(events=full_data, groupby_col=\"user_id\", sort_col=\"timestamp\")\n",
+    "    return grouped_interactions"
+   ]
+  },
+  {
+   "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>timestamp</th>\n",
+       "      <th>item_id</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>0</td>\n",
+       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
+       "      <td>[2426, 822, 2733, 2587, 2937, 3618, 2943, 708,...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>1</td>\n",
+       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
+       "      <td>[3272, 3026, 2760, 851, 346, 3393, 1107, 515, ...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>2</td>\n",
+       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
+       "      <td>[579, 1140, 1154, 2426, 1524, 1260, 2160, 2621...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>3</td>\n",
+       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
+       "      <td>[1781, 2940, 2468, 890, 948, 106, 593, 309, 49...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>4</td>\n",
+       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
+       "      <td>[1108, 2229, 21, 2435, 2142, 106, 1167, 593, 1...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>...</th>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6035</th>\n",
+       "      <td>6035</td>\n",
+       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
+       "      <td>[2426, 1279, 3151, 3321, 1178, 3301, 2501, 278...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6036</th>\n",
+       "      <td>6036</td>\n",
+       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
+       "      <td>[1592, 2302, 1633, 1813, 2879, 1482, 2651, 250...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6037</th>\n",
+       "      <td>6037</td>\n",
+       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
+       "      <td>[1971, 3500, 2077, 1666, 1399, 2651, 2748, 283...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6038</th>\n",
+       "      <td>6038</td>\n",
+       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
+       "      <td>[1486, 1485, 3384, 3512, 3302, 3126, 3650, 330...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6039</th>\n",
+       "      <td>6039</td>\n",
+       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
+       "      <td>[2432, 2960, 1848, 2114, 2142, 3091, 3248, 317...</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "<p>6040 rows × 3 columns</p>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "      user_id                                          timestamp  \\\n",
+       "0           0  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
+       "1           1  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
+       "2           2  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
+       "3           3  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
+       "4           4  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
+       "...       ...                                                ...   \n",
+       "6035     6035  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
+       "6036     6036  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
+       "6037     6037  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
+       "6038     6038  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
+       "6039     6039  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
+       "\n",
+       "                                                item_id  \n",
+       "0     [2426, 822, 2733, 2587, 2937, 3618, 2943, 708,...  \n",
+       "1     [3272, 3026, 2760, 851, 346, 3393, 1107, 515, ...  \n",
+       "2     [579, 1140, 1154, 2426, 1524, 1260, 2160, 2621...  \n",
+       "3     [1781, 2940, 2468, 890, 948, 106, 593, 309, 49...  \n",
+       "4     [1108, 2229, 21, 2435, 2142, 106, 1167, 593, 1...  \n",
+       "...                                                 ...  \n",
+       "6035  [2426, 1279, 3151, 3321, 1178, 3301, 2501, 278...  \n",
+       "6036  [1592, 2302, 1633, 1813, 2879, 1482, 2651, 250...  \n",
+       "6037  [1971, 3500, 2077, 1666, 1399, 2651, 2748, 283...  \n",
+       "6038  [1486, 1485, 3384, 3512, 3302, 3126, 3650, 330...  \n",
+       "6039  [2432, 2960, 1848, 2114, 2142, 3091, 3248, 317...  \n",
+       "\n",
+       "[6040 rows x 3 columns]"
+      ]
+     },
+     "execution_count": 7,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "train_events = bake_data(train_events)\n",
+    "\n",
+    "validation_events = bake_data(validation_events)\n",
+    "validation_gt = bake_data(validation_gt)\n",
+    "\n",
+    "test_events = bake_data(test_events)\n",
+    "test_gt = bake_data(test_gt)\n",
+    "\n",
+    "train_events"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "To ensure we don't have unknown users in ground truth, we join validation events and validation ground truth (also join test events and test ground truth correspondingly) by user ids to leave only the common ones.  "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def add_gt_to_events(events_df, gt_df):\n",
+    "    gt_to_join = gt_df[[\"user_id\", \"item_id\"]].rename(columns={\"item_id\": \"ground_truth\"})\n",
+    "\n",
+    "    events_df = events_df.merge(gt_to_join, on=\"user_id\", how=\"inner\")\n",
+    "    return events_df\n",
+    "\n",
+    "validation_events = add_gt_to_events(validation_events, validation_gt)\n",
+    "test_events = add_gt_to_events(test_events, test_gt)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pathlib import Path\n",
+    "\n",
+    "data_dir = Path(\"temp/data/\")\n",
+    "data_dir.mkdir(parents=True, exist_ok=True)\n",
+    "\n",
+    "TRAIN_PATH = data_dir / \"train.parquet\"\n",
+    "VAL_PATH = data_dir / \"val.parquet\"\n",
+    "PREDICT_PATH = data_dir / \"test.parquet\"\n",
+    "\n",
+    "ENCODER_PATH = data_dir / \"encoder\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "train_events.to_parquet(TRAIN_PATH)\n",
+    "validation_events.to_parquet(VAL_PATH)\n",
+    "test_events.to_parquet(PREDICT_PATH)\n",
+    "\n",
+    "encoder.save(ENCODER_PATH)"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Prepare to model training\n",
+    "### 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. It is required by the SasRec model to correctly create embeddings for every source column. Note that user_id does not required in `TensorSchema`.\n",
+    "\n",
+    "Note that the **padding value** is the next value (item_id) after the last one. **Cardinality** is the number of unique values ​​given the padding value."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from replay.data import FeatureHint, FeatureType\n",
+    "from replay.data.nn import TensorFeatureInfo, TensorSchema\n",
+    "\n",
+    "\n",
+    "EMBEDDING_DIM = 64\n",
+    "\n",
+    "encoder = encoder.load(ENCODER_PATH)\n",
+    "NUM_UNIQUE_ITEMS = len(encoder.mapping[\"item_id\"])\n",
+    "\n",
+    "tensor_schema = TensorSchema(\n",
+    "    [\n",
+    "        TensorFeatureInfo(\n",
+    "            name=\"item_id\",\n",
+    "            is_seq=True,\n",
+    "            padding_value=NUM_UNIQUE_ITEMS,\n",
+    "            cardinality=NUM_UNIQUE_ITEMS + 1,  # taking into account padding\n",
+    "            embedding_dim=EMBEDDING_DIM,\n",
+    "            feature_type=FeatureType.CATEGORICAL,\n",
+    "            feature_hint=FeatureHint.ITEM_ID,\n",
+    "        )\n",
+    "    ]\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Configure ParquetModule and transformation pipelines\n",
+    "\n",
+    "The `ParquetModule` class enables training of models on large datasets by reading data in batch-wise way. This class initialized with **paths to every data split, a metadata dict containing information about shape and padding value of every column and a dict of transforms**. `ParquetModule`'s  \"transform pipelines\" are stage-specific modules implementing additional preprocessing to be performed on batch level right before the forward pass.  \n",
+    "\n",
+    "For SasRec model, RePlay provides a function that generates a sequence of appropriate transforms for each data split named **make_default_sasrec_transforms**.\n",
+    "\n",
+    "Internally this function creates the following transforms:\n",
+    "1) Training:\n",
+    "    1. Create a target, which contains the shifted item sequence that represents the next item in the sequence (for the next item prediction task).\n",
+    "    2. Rename features to match it with expected format by the model during training.\n",
+    "    3. Unsqueeze target (*positive_labels*) and it's padding mask (*target_padding_mask*) for getting required shape of this tensors for loss computation.\n",
+    "    4. Group input features to be embed in expected format.\n",
+    "\n",
+    "2) Validation/Inference:\n",
+    "    1. Rename/group features to match it with expected format by the model during valdiation/inference.\n",
+    "\n",
+    "If a different set of transforms is required, you can create them yourself and submit them to the ParquetModule in the form of a dictionary where the key is the name of the split, and the value is the list of transforms. Available transforms are in the replay/nn/transforms/.\n",
+    "\n",
+    "**Note:** One of the transforms for the training data prepares the initial sequence for the task of Next Item Prediction so it shifts the sequence of items. For the final sequence length to be correct, you need to set shape of item_id in metadata as **model sequence length + shift**. Default shift value is 1."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import copy\n",
+    "\n",
+    "import torch\n",
+    "\n",
+    "from replay.data.nn import TensorSchema\n",
+    "from replay.nn.transform import GroupTransform, NextTokenTransform, RenameTransform, UnsqueezeTransform, ThresholdNegativeSamplingTransform"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def make_sasrec_transforms(\n",
+    "    tensor_schema: TensorSchema, query_column: str = \"query_id\", num_negative_samples: int = 128,\n",
+    ") -> dict[str, list[torch.nn.Module]]:\n",
+    "    item_column = tensor_schema.item_id_feature_name\n",
+    "    vocab_size = tensor_schema[item_column].cardinality\n",
+    "    train_transforms = [\n",
+    "        ThresholdNegativeSamplingTransform(vocab_size, num_negative_samples),\n",
+    "        NextTokenTransform(label_field=item_column, query_features=query_column, shift=1),\n",
+    "        RenameTransform(\n",
+    "            {\n",
+    "                query_column: \"query_id\",\n",
+    "                f\"{item_column}_mask\": \"padding_mask\",\n",
+    "                \"positive_labels_mask\": \"target_padding_mask\",\n",
+    "            }\n",
+    "        ),\n",
+    "        UnsqueezeTransform(\"target_padding_mask\", -1),\n",
+    "        UnsqueezeTransform(\"positive_labels\", -1),\n",
+    "        GroupTransform({\"feature_tensors\": [item_column]}),\n",
+    "    ]\n",
+    "\n",
+    "    val_transforms = [\n",
+    "        RenameTransform({query_column: \"query_id\", f\"{item_column}_mask\": \"padding_mask\"}),\n",
+    "        GroupTransform({\"feature_tensors\": [item_column]}),\n",
+    "    ]\n",
+    "    test_transforms = copy.deepcopy(val_transforms)\n",
+    "\n",
+    "    predict_transforms = copy.deepcopy(val_transforms)\n",
+    "\n",
+    "    transforms = {\n",
+    "        \"train\": train_transforms,\n",
+    "        \"validate\": val_transforms,\n",
+    "        \"test\": test_transforms,\n",
+    "        \"predict\": predict_transforms,\n",
+    "    }\n",
+    "\n",
+    "    return transforms\n",
+    "\n",
+    "transforms = make_sasrec_transforms(tensor_schema, query_column=\"user_id\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "MAX_SEQ_LEN = 50\n",
+    "\n",
+    "def create_meta(shape: int, gt_shape: Optional[int] = None):\n",
+    "    meta = {\n",
+    "        \"user_id\": {},\n",
+    "        \"item_id\": {\"shape\": shape, \"padding\": tensor_schema[\"item_id\"].padding_value},\n",
+    "    }\n",
+    "    if gt_shape is not None:\n",
+    "        meta.update({\"ground_truth\": {\"shape\": gt_shape, \"padding\": -1}})\n",
+    "\n",
+    "    return meta\n",
+    "\n",
+    "train_metadata = {\n",
+    "    \"train\": create_meta(shape=MAX_SEQ_LEN+1),\n",
+    "    \"validate\": create_meta(shape=MAX_SEQ_LEN, gt_shape=1),\n",
+    "}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from replay.data.nn import ParquetModule\n",
+    "\n",
+    "BATCH_SIZE = 32\n",
+    "\n",
+    "parquet_module = ParquetModule(\n",
+    "    train_path=TRAIN_PATH,\n",
+    "    validate_path=VAL_PATH,\n",
+    "    batch_size=BATCH_SIZE,\n",
+    "    metadata=train_metadata,\n",
+    "    transforms=transforms,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Train model\n",
+    "### Create SasRec model instance and run the training stage using lightning\n",
+    "We may now train the model using the Lightning trainer class. \n",
+    "\n",
+    "RePlay's implementation of SasRec is designed in a modular, **block-based approach**. Instead of passing configuration parameters to the constructor, SasRec is now built by providing fully initialized components that makes the model more flexible and easier to extend.\n",
+    "\n",
+    "#### Default Configuration\n",
+    "\n",
+    "Default SasRec model may be created quickly via method **from_params**. Default model instance has CE loss, original SasRec transformer layes, and embeddings are aggregated via sum."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from replay.nn.sequential import SasRec\n",
+    "from typing import Literal\n",
+    "def make_sasrec(\n",
+    "    schema: TensorSchema,\n",
+    "    embedding_dim: int = 192,\n",
+    "    num_heads: int = 4,\n",
+    "    num_blocks: int = 2,\n",
+    "    max_sequence_length: int = 50,\n",
+    "    dropout: float = 0.3,\n",
+    "    excluded_features: Optional[list[str]] = None,\n",
+    "    categorical_list_feature_aggregation_method: Literal[\"sum\", \"mean\", \"max\"] = \"sum\",\n",
+    ") -> SasRec:\n",
+    "    from replay.nn.sequential.sasrec import SasRecBody, SasRecTransformerLayer\n",
+    "    from replay.nn.agg import SumAggregator\n",
+    "    from replay.nn.embedding import SequenceEmbedding\n",
+    "    from replay.nn.loss import CE, CESampled\n",
+    "    from replay.nn.mask import DefaultAttentionMask\n",
+    "    from replay.nn.sequential.sasrec.agg import PositionAwareAggregator\n",
+    "    from replay.nn.sequential.sasrec.transformer import SasRecTransformerLayer\n",
+    "    excluded_features = [\n",
+    "        schema.query_id_feature_name,\n",
+    "        schema.timestamp_feature_name,\n",
+    "        *(excluded_features or []),\n",
+    "    ]\n",
+    "    excluded_features = list(set(excluded_features))\n",
+    "    body = SasRecBody(\n",
+    "        embedder=SequenceEmbedding(\n",
+    "            schema=schema,\n",
+    "            categorical_list_feature_aggregation_method=categorical_list_feature_aggregation_method,\n",
+    "            excluded_features=excluded_features,\n",
+    "        ),\n",
+    "        embedding_aggregator=PositionAwareAggregator(\n",
+    "            embedding_aggregator=SumAggregator(embedding_dim=embedding_dim),\n",
+    "            max_sequence_length=max_sequence_length,\n",
+    "            dropout=dropout,\n",
+    "        ),\n",
+    "        attn_mask_builder=DefaultAttentionMask(\n",
+    "            reference_feature_name=schema.item_id_feature_name,\n",
+    "            num_heads=num_heads,\n",
+    "        ),\n",
+    "        encoder=SasRecTransformerLayer(\n",
+    "            embedding_dim=embedding_dim,\n",
+    "            num_heads=num_heads,\n",
+    "            num_blocks=num_blocks,\n",
+    "            dropout=dropout,\n",
+    "            activation=\"relu\",\n",
+    "        ),\n",
+    "        output_normalization=torch.nn.LayerNorm(embedding_dim),\n",
+    "    )\n",
+    "    padding_idx = schema.item_id_features.item().padding_value\n",
+    "    return SasRec(\n",
+    "        body=body,\n",
+    "        loss=CESampled(padding_idx=padding_idx),\n",
+    "    )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "NUM_BLOCKS = 2\n",
+    "NUM_HEADS = 2\n",
+    "DROPOUT = 0.3\n",
+    "\n",
+    "sasrec = make_sasrec(\n",
+    "    schema=tensor_schema,\n",
+    "    embedding_dim=EMBEDDING_DIM,\n",
+    "    max_sequence_length=MAX_SEQ_LEN,\n",
+    "    num_heads=NUM_HEADS,\n",
+    "    num_blocks=NUM_BLOCKS,\n",
+    "    dropout=DROPOUT,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "A universal PyTorch Lightning module is provided. It can work with any NN model."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from replay.nn.lightning.optimizer import OptimizerFactory\n",
+    "from replay.nn.lightning.scheduler import LRSchedulerFactory\n",
+    "from replay.nn.lightning import LightningModule\n",
+    "\n",
+    "model = LightningModule(\n",
+    "    sasrec,\n",
+    "    optimizer_factory=OptimizerFactory(),\n",
+    "    lr_scheduler_factory=LRSchedulerFactory(),\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "To facilitate training, we add the following callbacks:\n",
+    "1) `ModelCheckpoint` - to save the best trained model based on its Recall metric. It's a default Lightning Callback.\n",
+    "1) `ComputeMetricsCallback` - to display a detailed validation metric matrix after each epoch. It's a custom RePlay callback for computing recsys metrics on validation and test stages.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "GPU available: True (cuda), used: True\n",
+      "TPU available: False, using: 0 TPU cores\n",
+      "You are using a CUDA device ('NVIDIA GeForce RTX 3060 Laptop GPU') 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 | Mode  | FLOPs\n",
+      "-------------------------------------------------\n",
+      "0 | model | SasRec | 291 K  | train | 0    \n",
+      "-------------------------------------------------\n",
+      "291 K     Trainable params\n",
+      "0         Non-trainable params\n",
+      "291 K     Total params\n",
+      "1.164     Total estimated model params size (MB)\n",
+      "39        Modules in train mode\n",
+      "0         Modules in eval mode\n",
+      "0         Total Flops\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "21ca5ea332d9418fb2481e0296b399a0",
+       "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": "dfbb5a69743640d282090979774e93a2",
+       "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": "b6b1a774cbc84c73b23e546285b0a116",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 0, global step 189: 'recall@10' reached 0.02517 (best 0.02517), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=0-step=189-v3.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.003477  0.008391  0.009940  0.006889\n",
+      "ndcg    0.003477  0.012236  0.018080  0.008533\n",
+      "recall  0.003477  0.025170  0.048684  0.013578\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "f47ec635a7c143f281864d439cbbcc8e",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 1, global step 378: 'recall@10' reached 0.04090 (best 0.04090), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=1-step=378-v2.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.005133  0.013343  0.015591  0.011023\n",
+      "ndcg    0.005133  0.019693  0.028004  0.014007\n",
+      "recall  0.005133  0.040901  0.074019  0.023183\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "bbdfceaef5df4df98b29883914a7c174",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 2, global step 567: 'recall@10' reached 0.06773 (best 0.06773), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=2-step=567.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.008776  0.021569  0.024780  0.017903\n",
+      "ndcg    0.008776  0.032154  0.044198  0.022983\n",
+      "recall  0.008776  0.067726  0.116079  0.038748\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "7f43be0bcd8c49c7b29e3689b24ad818",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 3, global step 756: 'recall@10' reached 0.08197 (best 0.08197), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=3-step=756-v3.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k             1        10        20         5\n",
+      "map     0.00977  0.024920  0.029111  0.019702\n",
+      "ndcg    0.00977  0.037940  0.053492  0.025112\n",
+      "recall  0.00977  0.081967  0.144064  0.041894\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "0df4dc7b7f964a4089205386735e2968",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 4, global step 945: 'recall@10' reached 0.10051 (best 0.10051), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=4-step=945-v4.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.011757  0.032053  0.036894  0.025981\n",
+      "ndcg    0.011757  0.047823  0.065510  0.033102\n",
+      "recall  0.011757  0.100513  0.170558  0.054976\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "92b4eac8339b483684bc947c3e9f035b",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 5, global step 1134: 'recall@10' reached 0.11707 (best 0.11707), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=5-step=1134-v1.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k            1        10        20         5\n",
+      "map     0.0154  0.038303  0.043634  0.031517\n",
+      "ndcg    0.0154  0.056389  0.076198  0.039699\n",
+      "recall  0.0154  0.117072  0.196225  0.064911\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "2dde8a65268b46ce85312ed5529729a1",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 6, global step 1323: 'recall@10' reached 0.13214 (best 0.13214), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=6-step=1323-v2.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.015069  0.042212  0.047618  0.034832\n",
+      "ndcg    0.015069  0.062962  0.082865  0.044922\n",
+      "recall  0.015069  0.132141  0.211293  0.076006\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "d452da2b13ca4953b58fffaecb5143dc",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 7, global step 1512: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.015731  0.041482  0.047466  0.033985\n",
+      "ndcg    0.015731  0.061963  0.083911  0.043462\n",
+      "recall  0.015731  0.130651  0.217751  0.072694\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "0dc27c98e71f493da3ef87e551bd0c60",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 8, global step 1701: 'recall@10' reached 0.14406 (best 0.14406), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=8-step=1701-v1.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.015897  0.045616  0.051492  0.037233\n",
+      "ndcg    0.015897  0.068287  0.089852  0.047670\n",
+      "recall  0.015897  0.144064  0.229674  0.079649\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "b143a23c75a347139ec26ef54abad1b6",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 9, global step 1890: 'recall@10' reached 0.15102 (best 0.15102), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=9-step=1890-v2.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017387  0.048506  0.054854  0.040073\n",
+      "ndcg    0.017387  0.072139  0.095364  0.051380\n",
+      "recall  0.017387  0.151018  0.243087  0.086107\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "23cf891fa04440408cf156e3bb1d3245",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 10, global step 2079: 'recall@10' reached 0.15201 (best 0.15201), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=10-step=2079.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k             1        10        20         5\n",
+      "map     0.01689  0.048957  0.055559  0.040876\n",
+      "ndcg    0.01689  0.072775  0.096923  0.052810\n",
+      "recall  0.01689  0.152012  0.247723  0.089419\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "b9eee01e7c704ed9a0559f130d4ca137",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 11, global step 2268: 'recall@10' reached 0.16128 (best 0.16128), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=11-step=2268.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016228  0.049785  0.056337  0.040514\n",
+      "ndcg    0.016228  0.075524  0.099693  0.052985\n",
+      "recall  0.016228  0.161285  0.257493  0.091406\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "060953ee21ef46d49a34a75df75b57bd",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 12, global step 2457: 'recall@10' reached 0.16327 (best 0.16327), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=12-step=2457-v1.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.014738  0.048876  0.055841  0.039215\n",
+      "ndcg    0.014738  0.075207  0.100830  0.051456\n",
+      "recall  0.014738  0.163272  0.265110  0.089088\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "9db57c2e20e046e3a7635fec4ee98680",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 13, global step 2646: 'recall@10' reached 0.16393 (best 0.16393), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=13-step=2646.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016725  0.051229  0.058088  0.041919\n",
+      "ndcg    0.016725  0.077239  0.102477  0.054437\n",
+      "recall  0.016725  0.163934  0.264282  0.092896\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "4de755f88a624acd90ee5787f42d0eec",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 14, global step 2835: 'recall@10' reached 0.16791 (best 0.16791), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=14-step=2835.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.018049  0.052408  0.059556  0.042344\n",
+      "ndcg    0.018049  0.079003  0.105266  0.054474\n",
+      "recall  0.018049  0.167909  0.272231  0.091737\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "294c0a5c4b5a43d1850d722063b45faa",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 15, global step 3024: 'recall@10' reached 0.16907 (best 0.16907), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=15-step=3024.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017718  0.052563  0.059905  0.042822\n",
+      "ndcg    0.017718  0.079424  0.106455  0.055627\n",
+      "recall  0.017718  0.169068  0.276536  0.095049\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "5cec32b151e34d558aadb0a54a9fe645",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 16, global step 3213: 'recall@10' reached 0.17255 (best 0.17255), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=16-step=3213.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016725  0.053275  0.060719  0.043368\n",
+      "ndcg    0.016725  0.080840  0.108407  0.056778\n",
+      "recall  0.016725  0.172545  0.282497  0.098029\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "964972bd33324e2e88184dfa18ac414f",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 17, global step 3402: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.015069  0.052175  0.059507  0.042386\n",
+      "ndcg    0.015069  0.079773  0.106812  0.055673\n",
+      "recall  0.015069  0.171717  0.279351  0.096374\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "749e9689bfc842daa9ed09f725f4323d",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 18, global step 3591: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017387  0.052979  0.060514  0.042849\n",
+      "ndcg    0.017387  0.079924  0.107494  0.055138\n",
+      "recall  0.017387  0.169896  0.279185  0.092731\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "1226764060764a59a3f2cfb1a822f8b2",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 19, global step 3780: 'recall@10' reached 0.18066 (best 0.18066), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=19-step=3780.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016559  0.053610  0.060994  0.042965\n",
+      "ndcg    0.016559  0.082830  0.110151  0.056616\n",
+      "recall  0.016559  0.180659  0.289617  0.098692\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "3073e689c7d340b59b6989506c74c700",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 20, global step 3969: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.015069  0.052912  0.060238  0.042455\n",
+      "ndcg    0.015069  0.081482  0.108433  0.056028\n",
+      "recall  0.015069  0.176685  0.283822  0.097698\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "f86fc39ef2104c69b1d1f89818719d08",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 21, global step 4158: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.015731  0.053080  0.060646  0.043241\n",
+      "ndcg    0.015731  0.080854  0.108792  0.056898\n",
+      "recall  0.015731  0.173207  0.284484  0.098857\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "356d76af5ca94c5ca0efa94d57a7f71c",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 22, global step 4347: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017056  0.055227  0.062985  0.044949\n",
+      "ndcg    0.017056  0.083916  0.112344  0.058574\n",
+      "recall  0.017056  0.179666  0.292433  0.100348\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "79f90e00872f44b7b33c0f8c77ff011a",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 23, global step 4536: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.015069  0.052468  0.060395  0.042041\n",
+      "ndcg    0.015069  0.080925  0.110088  0.055468\n",
+      "recall  0.015069  0.175857  0.291770  0.096705\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "8d04fd5e94fe4711aee0ad54b356d9eb",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 24, global step 4725: 'recall@10' reached 0.18397 (best 0.18397), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=24-step=4725.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017718  0.055757  0.063565  0.044925\n",
+      "ndcg    0.017718  0.085273  0.114202  0.058594\n",
+      "recall  0.017718  0.183971  0.299387  0.100513\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "cfd12e9d0ad847c09dcf07d9532add8f",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 25, global step 4914: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017387  0.054950  0.062547  0.044381\n",
+      "ndcg    0.017387  0.083432  0.111543  0.057612\n",
+      "recall  0.017387  0.178506  0.290611  0.098195\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "363e70c34eff435597795909950744bd",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 26, global step 5103: 'recall@10' reached 0.18463 (best 0.18463), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=26-step=5103.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.019043  0.057219  0.064977  0.046506\n",
+      "ndcg    0.019043  0.086570  0.114830  0.060294\n",
+      "recall  0.019043  0.184633  0.296407  0.102666\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "48eb130d2ce549d4a70d48979d153fd2",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 27, global step 5292: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.019208  0.057371  0.065001  0.046650\n",
+      "ndcg    0.019208  0.085884  0.114015  0.059848\n",
+      "recall  0.019208  0.180825  0.292764  0.100182\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "4f715157ede74adba3455a0ef7e5131e",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 28, global step 5481: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.018381  0.056056  0.063808  0.045615\n",
+      "ndcg    0.018381  0.085206  0.113792  0.059325\n",
+      "recall  0.018381  0.182812  0.296572  0.101507\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "32ee6a09dcc248cfa8af8e6f803c9970",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 29, global step 5670: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k             1        10        20         5\n",
+      "map     0.01689  0.054551  0.062793  0.044119\n",
+      "ndcg    0.01689  0.083030  0.113385  0.057502\n",
+      "recall  0.01689  0.178010  0.298725  0.098526\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "8b0022b7e1334ae198b407aa045798f0",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 30, global step 5859: 'recall@10' reached 0.18828 (best 0.18828), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=30-step=5859.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017387  0.057006  0.064480  0.046299\n",
+      "ndcg    0.017387  0.087281  0.114612  0.060979\n",
+      "recall  0.017387  0.188276  0.296572  0.106143\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "7e91bd3715e74bb88a5186a9bd2018b0",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 31, global step 6048: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016559  0.055557  0.063455  0.045518\n",
+      "ndcg    0.016559  0.084519  0.113808  0.059953\n",
+      "recall  0.016559  0.180825  0.297731  0.104322\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "6518c10046494dceb7539431f4342d40",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 32, global step 6237: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.015234  0.055037  0.063161  0.044312\n",
+      "ndcg    0.015234  0.085493  0.115172  0.059173\n",
+      "recall  0.015234  0.187117  0.304686  0.104984\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "f6ed4206a4514b399847c7f23dc9ded5",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 33, global step 6426: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.015897  0.055037  0.063057  0.043967\n",
+      "ndcg    0.015897  0.085474  0.114917  0.058397\n",
+      "recall  0.015897  0.187117  0.304024  0.102832\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "07380901a70f4cf4bccda988256f4c63",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 34, global step 6615: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.014738  0.053996  0.061706  0.042709\n",
+      "ndcg    0.014738  0.083953  0.112215  0.056473\n",
+      "recall  0.014738  0.183971  0.296076  0.098692\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "c2f3ecd2dccc4e188f7a1e1098e35644",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 35, global step 6804: 'recall@10' reached 0.18844 (best 0.18844), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=35-step=6804-v1.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.015069  0.054859  0.063045  0.043283\n",
+      "ndcg    0.015069  0.085649  0.116068  0.057529\n",
+      "recall  0.015069  0.188442  0.309985  0.101341\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "f668d2ae168b4a18aab4b55f571d65b5",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 36, global step 6993: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016725  0.057387  0.065605  0.046467\n",
+      "ndcg    0.016725  0.087272  0.117565  0.060911\n",
+      "recall  0.016725  0.186455  0.307004  0.105150\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "e99e9ff82afd4989910381ef0a5203b5",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 37, global step 7182: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.014903  0.055713  0.063883  0.044455\n",
+      "ndcg    0.014903  0.086266  0.116276  0.058719\n",
+      "recall  0.014903  0.188111  0.307336  0.102335\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "6db4c2801b96444ea9c1242071243bea",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 38, global step 7371: 'recall@10' reached 0.19208 (best 0.19208), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=38-step=7371.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.015897  0.056195  0.063902  0.044337\n",
+      "ndcg    0.015897  0.087493  0.115857  0.058652\n",
+      "recall  0.015897  0.192085  0.304852  0.102666\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "815fab87a5fe4603adf4c1db9881772e",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 39, global step 7560: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016559  0.056413  0.064644  0.045493\n",
+      "ndcg    0.016559  0.086769  0.117049  0.059937\n",
+      "recall  0.016559  0.188111  0.308495  0.104322\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "8759b410188f4ac795f88e785b27542f",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 40, global step 7749: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.014406  0.055492  0.063463  0.044392\n",
+      "ndcg    0.014406  0.086624  0.116225  0.059581\n",
+      "recall  0.014406  0.190263  0.308495  0.106309\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "7ec1302b87b3481db013eb0b41f690a8",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 41, global step 7938: 'recall@10' reached 0.19225 (best 0.19225), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=41-step=7938.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017221  0.058234  0.066382  0.047257\n",
+      "ndcg    0.017221  0.089155  0.118997  0.062111\n",
+      "recall  0.017221  0.192250  0.310647  0.107634\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "fef8f8d3fb0a4a6f8e6d3fcd7166421c",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 42, global step 8127: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.018381  0.057681  0.065612  0.046263\n",
+      "ndcg    0.018381  0.088403  0.117568  0.060508\n",
+      "recall  0.018381  0.191091  0.307004  0.104322\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "75e8ff0777b545d28653647466ff5fd3",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 43, global step 8316: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.015731  0.057188  0.065122  0.045949\n",
+      "ndcg    0.015731  0.088375  0.117464  0.060972\n",
+      "recall  0.015731  0.192250  0.307667  0.107137\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "b0dec430572c4d49a50a461c816aafe7",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 44, global step 8505: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.014903  0.055369  0.063525  0.044646\n",
+      "ndcg    0.014903  0.086007  0.115973  0.059935\n",
+      "recall  0.014903  0.187779  0.306839  0.106971\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "dfa485d5d42e49a8a9070ee85aef434f",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 45, global step 8694: 'recall@10' reached 0.19374 (best 0.19374), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=45-step=8694.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k             1        10        20         5\n",
+      "map     0.01391  0.055919  0.063806  0.044141\n",
+      "ndcg    0.01391  0.087688  0.116770  0.058898\n",
+      "recall  0.01391  0.193741  0.309488  0.104156\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "a38a593a5f2345caaf79560dc66e8d9f",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 46, global step 8883: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.018546  0.057837  0.065681  0.046448\n",
+      "ndcg    0.018546  0.088962  0.118050  0.061033\n",
+      "recall  0.018546  0.193078  0.309157  0.105978\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "6ce539a3aabb4dd695acc6c07ad72ac9",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 47, global step 9072: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016559  0.057167  0.065194  0.045987\n",
+      "ndcg    0.016559  0.088154  0.117904  0.060651\n",
+      "recall  0.016559  0.191588  0.310316  0.105647\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "30f737adc35f4b4c90a419732e8e1e91",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 48, global step 9261: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k             1        10        20         5\n",
+      "map     0.01689  0.058332  0.066214  0.047022\n",
+      "ndcg    0.01689  0.089275  0.118570  0.061644\n",
+      "recall  0.01689  0.192416  0.309488  0.106475\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "bf5a91c9655947cc91918c15b2a9462b",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 49, global step 9450: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016393  0.057926  0.066244  0.046906\n",
+      "ndcg    0.016393  0.088820  0.119375  0.062076\n",
+      "recall  0.016393  0.191422  0.312800  0.108627\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "f21c5ac4b4734d74a23756a1fdf040f2",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 50, global step 9639: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017718  0.058476  0.066364  0.046785\n",
+      "ndcg    0.017718  0.089596  0.118568  0.061171\n",
+      "recall  0.017718  0.193410  0.308495  0.105315\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "7269482228054872bb3fd02ba7ed438b",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 51, global step 9828: 'recall@10' reached 0.19540 (best 0.19540), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=51-step=9828-v1.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.015731  0.057566  0.065119  0.045990\n",
+      "ndcg    0.015731  0.089326  0.117305  0.060894\n",
+      "recall  0.015731  0.195397  0.307004  0.106640\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "f2e119321d274785b3001a5f44f941a4",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 52, global step 10017: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k             1        10        20         5\n",
+      "map     0.01954  0.058868  0.067041  0.048454\n",
+      "ndcg    0.01954  0.088730  0.118800  0.063159\n",
+      "recall  0.01954  0.188276  0.307832  0.108462\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "712a7c0cf2a048679ce1e68b12de03b6",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 53, global step 10206: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017718  0.058614  0.066505  0.047370\n",
+      "ndcg    0.017718  0.089564  0.118682  0.062093\n",
+      "recall  0.017718  0.192747  0.308660  0.107303\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "dfb0a6436d2147adb94f4a9c8ea16aac",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 54, global step 10395: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016725  0.057676  0.065385  0.046299\n",
+      "ndcg    0.016725  0.089061  0.117540  0.061333\n",
+      "recall  0.016725  0.193741  0.307170  0.107634\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "f890babbe4e64f739ba391bce80360e9",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 55, global step 10584: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016062  0.056661  0.064644  0.045372\n",
+      "ndcg    0.016062  0.087800  0.117280  0.060313\n",
+      "recall  0.016062  0.191588  0.308992  0.106309\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "f4c2e6243e484477bba7d4de22a19810",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 56, global step 10773: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.018049  0.058768  0.066948  0.047558\n",
+      "ndcg    0.018049  0.090098  0.120345  0.062504\n",
+      "recall  0.018049  0.194734  0.315284  0.108462\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "87dba884d40f4568acd8ca8bbbd64f6b",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 57, global step 10962: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.014572  0.053518  0.062111  0.042355\n",
+      "ndcg    0.014572  0.083976  0.115510  0.056682\n",
+      "recall  0.014572  0.185792  0.310979  0.100845\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "58699cb7550047feb645c6af4f054fcc",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 58, global step 11151: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.018049  0.058147  0.066599  0.047524\n",
+      "ndcg    0.018049  0.088199  0.119005  0.062165\n",
+      "recall  0.018049  0.188276  0.310151  0.107137\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "086f8398ee09499398aa7d3b70c79aa8",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 59, global step 11340: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.014406  0.055658  0.063896  0.044665\n",
+      "ndcg    0.014406  0.086842  0.117136  0.059875\n",
+      "recall  0.014406  0.190760  0.311144  0.106640\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "eb274ba8390345fc95f0e0f78a1f68e4",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 60, global step 11529: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.015731  0.056472  0.064719  0.045739\n",
+      "ndcg    0.015731  0.087657  0.117903  0.061421\n",
+      "recall  0.015731  0.191422  0.311475  0.109786\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "b4c4b8b17b194484b484a34ecf6c211b",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 61, global step 11718: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k             1        10        20         5\n",
+      "map     0.01689  0.057109  0.065266  0.045570\n",
+      "ndcg    0.01689  0.088058  0.118054  0.060104\n",
+      "recall  0.01689  0.191257  0.310482  0.104819\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "8869b85b84e141069b219175f8a5fd05",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 62, global step 11907: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.014738  0.056182  0.064093  0.044226\n",
+      "ndcg    0.014738  0.087940  0.116915  0.058911\n",
+      "recall  0.014738  0.193906  0.308826  0.103991\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "284a8281380c4b51bc3e4d1fa821ce55",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 63, global step 12096: 'recall@10' reached 0.19556 (best 0.19556), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=63-step=12096.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016228  0.058258  0.066323  0.047014\n",
+      "ndcg    0.016228  0.089951  0.119723  0.062421\n",
+      "recall  0.016228  0.195562  0.314125  0.109786\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "a172b3e4752e4ac8b1df1f435e3d0401",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 64, global step 12285: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016725  0.058308  0.066315  0.047469\n",
+      "ndcg    0.016725  0.089611  0.119061  0.063139\n",
+      "recall  0.016725  0.193741  0.310813  0.111442\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "3481bce9e0ee43c58b478c79937e831b",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 65, global step 12474: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.015069  0.055754  0.063692  0.043901\n",
+      "ndcg    0.015069  0.087118  0.116341  0.058438\n",
+      "recall  0.015069  0.191754  0.307998  0.103163\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "929ab3208337487998be56f8cc63c14a",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 66, global step 12663: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.015234  0.056930  0.065395  0.045352\n",
+      "ndcg    0.015234  0.088645  0.119658  0.060367\n",
+      "recall  0.015234  0.194403  0.317437  0.106475\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "2672359bbd564354a429be48c4ac375b",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 67, global step 12852: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.014241  0.055383  0.063480  0.044845\n",
+      "ndcg    0.014241  0.086363  0.116422  0.060699\n",
+      "recall  0.014241  0.189270  0.309323  0.109621\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "0403c1ea57694538b5bcdf72a28d3e25",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 68, global step 13041: 'recall@10' reached 0.19622 (best 0.19622), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=68-step=13041.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.018712  0.058867  0.067064  0.047028\n",
+      "ndcg    0.018712  0.090484  0.120756  0.061726\n",
+      "recall  0.018712  0.196225  0.316774  0.107137\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "a697daf59428459886d557106ec3cc86",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 69, global step 13230: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.018712  0.059296  0.067623  0.047822\n",
+      "ndcg    0.018712  0.090358  0.121063  0.062465\n",
+      "recall  0.018712  0.193906  0.316112  0.107468\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "0d87a6c21b0e4db48b4a6387e1482296",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 70, global step 13419: 'recall@10' reached 0.19672 (best 0.19672), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=70-step=13419.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017718  0.058932  0.066943  0.047301\n",
+      "ndcg    0.017718  0.090675  0.119968  0.062253\n",
+      "recall  0.017718  0.196721  0.312800  0.108296\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "bbb281775c9545eaa92d24f00b763b73",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 71, global step 13608: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017056  0.057241  0.065559  0.045708\n",
+      "ndcg    0.017056  0.088863  0.119476  0.060634\n",
+      "recall  0.017056  0.194569  0.316278  0.106640\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "285af00d69c3444a8d58e6349850a98f",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 72, global step 13797: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.015069  0.056120  0.064450  0.044969\n",
+      "ndcg    0.015069  0.087871  0.118417  0.060586\n",
+      "recall  0.015069  0.193741  0.314953  0.108793\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "021ea8662cb5411eb86c8f63337dceeb",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 73, global step 13986: 'recall@10' reached 0.19854 (best 0.19854), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=73-step=13986.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016559  0.057389  0.065353  0.045104\n",
+      "ndcg    0.016559  0.089853  0.119195  0.059893\n",
+      "recall  0.016559  0.198543  0.315284  0.105481\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "f45164cbd400421dada62034fa3d49eb",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 74, global step 14175: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.015565  0.056737  0.065008  0.045043\n",
+      "ndcg    0.015565  0.088335  0.118772  0.059753\n",
+      "recall  0.015565  0.193906  0.314953  0.104984\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "7990b92ef50c48daa04193a0749c46e1",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 75, global step 14364: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k            1        10        20         5\n",
+      "map     0.0154  0.056695  0.064954  0.044974\n",
+      "ndcg    0.0154  0.088831  0.119369  0.059974\n",
+      "recall  0.0154  0.196390  0.318099  0.106143\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "5246f04cb4324f55b47ecb33070d932c",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 76, global step 14553: 'recall@10' reached 0.19904 (best 0.19904), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=76-step=14553.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.014406  0.056638  0.064821  0.044654\n",
+      "ndcg    0.014406  0.089447  0.119451  0.060199\n",
+      "recall  0.014406  0.199040  0.318099  0.108130\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "786ee80f518d4518ae0f92245a103b52",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 77, global step 14742: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k             1        10        20         5\n",
+      "map     0.01689  0.057254  0.065636  0.045783\n",
+      "ndcg    0.01689  0.088655  0.119626  0.060673\n",
+      "recall  0.01689  0.193575  0.316940  0.106640\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "bb577bcfe6fb48b2acecd29f761a841c",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 78, global step 14931: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.013413  0.053570  0.062183  0.042154\n",
+      "ndcg    0.013413  0.084824  0.116217  0.056864\n",
+      "recall  0.013413  0.189270  0.313462  0.102169\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "56ce55e76d164531a232448502b5e89e",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 79, global step 15120: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.013744  0.056911  0.065346  0.045518\n",
+      "ndcg    0.013744  0.089081  0.119783  0.061084\n",
+      "recall  0.013744  0.196225  0.317602  0.108793\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "216cc04625a84a04ae600fcf7a58760f",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 80, global step 15309: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017056  0.059605  0.067669  0.048805\n",
+      "ndcg    0.017056  0.091420  0.121109  0.064925\n",
+      "recall  0.017056  0.197218  0.315284  0.114589\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "cd55a7a84fa444398684901f380306c7",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 81, global step 15498: 'recall@10' reached 0.19970 (best 0.19970), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=81-step=15498-v1.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.019208  0.059915  0.067934  0.047475\n",
+      "ndcg    0.019208  0.092051  0.121520  0.061684\n",
+      "recall  0.019208  0.199702  0.316774  0.105315\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "1b8397e043464992abbf5f67594004eb",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 82, global step 15687: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017056  0.059126  0.066993  0.047193\n",
+      "ndcg    0.017056  0.091040  0.119902  0.061589\n",
+      "recall  0.017056  0.197880  0.312469  0.105647\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "7d663efa0c9b418da3409d032859661c",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 83, global step 15876: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.016393  0.058910  0.067300  0.047367\n",
+      "ndcg    0.016393  0.090708  0.121604  0.062527\n",
+      "recall  0.016393  0.196556  0.319424  0.108958\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "6a583740ecc14ad9a0c10b8f102de5d9",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 84, global step 16065: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k            1        10        20         5\n",
+      "map     0.0154  0.057330  0.065707  0.046324\n",
+      "ndcg    0.0154  0.088907  0.119815  0.061892\n",
+      "recall  0.0154  0.194072  0.317105  0.109786\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "0a7154f282df47c58f709f68c7bea385",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 85, global step 16254: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017884  0.058865  0.067077  0.046901\n",
+      "ndcg    0.017884  0.090545  0.120613  0.061410\n",
+      "recall  0.017884  0.196390  0.315615  0.105978\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "638857741499485686c1b4b6733f2713",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 86, global step 16443: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.013578  0.056565  0.064611  0.044949\n",
+      "ndcg    0.013578  0.088694  0.118386  0.060223\n",
+      "recall  0.013578  0.195728  0.313959  0.106971\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "c6604bb63cc448b994d1cbc8076c0e58",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 87, global step 16632: 'recall@10' reached 0.20003 (best 0.20003), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=87-step=16632.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.014738  0.057765  0.065777  0.045747\n",
+      "ndcg    0.014738  0.090573  0.120003  0.061197\n",
+      "recall  0.014738  0.200033  0.316940  0.108627\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "62a8bd2f89314bb3a6a6a3a45e6c0877",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 88, global step 16821: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.014406  0.056543  0.064898  0.045063\n",
+      "ndcg    0.014406  0.088396  0.118998  0.060280\n",
+      "recall  0.014406  0.194569  0.315946  0.106971\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "1e01bdb063ae49369e758c0a4da17182",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 89, global step 17010: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.015731  0.057330  0.065518  0.045824\n",
+      "ndcg    0.015731  0.088940  0.118949  0.060987\n",
+      "recall  0.015731  0.194237  0.313297  0.107634\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "bbc483c2ccad4f0da0c4311a56f1a4b5",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 90, global step 17199: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.015897  0.058550  0.066817  0.046854\n",
+      "ndcg    0.015897  0.090655  0.121206  0.062134\n",
+      "recall  0.015897  0.197549  0.319258  0.108958\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "64e9ddfeb5424717a97741b58dbd6f0e",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 91, global step 17388: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.014075  0.057416  0.065911  0.046128\n",
+      "ndcg    0.014075  0.089700  0.120977  0.062278\n",
+      "recall  0.014075  0.196887  0.321245  0.111939\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "e3d7f5856e1f4fedb35c0f77a956fab0",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 92, global step 17577: 'recall@10' reached 0.20036 (best 0.20036), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=92-step=17577-v2.ckpt' as top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.014241  0.057366  0.065310  0.044878\n",
+      "ndcg    0.014241  0.090287  0.119605  0.059684\n",
+      "recall  0.014241  0.200364  0.317105  0.104984\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "c6e17c5b3e6342c79419d1b83e848f30",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 93, global step 17766: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017221  0.058259  0.066380  0.047455\n",
+      "ndcg    0.017221  0.089409  0.119171  0.063038\n",
+      "recall  0.017221  0.193078  0.311144  0.111111\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "1dd6b48245004c5d994afe4da809e807",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 94, global step 17955: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.015897  0.058542  0.066742  0.046857\n",
+      "ndcg    0.015897  0.090490  0.120543  0.062166\n",
+      "recall  0.015897  0.196721  0.315946  0.109124\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "626ec34631fd432f82b6ab0bcb51a20c",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 95, global step 18144: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017221  0.059014  0.067464  0.047632\n",
+      "ndcg    0.017221  0.090545  0.121542  0.062936\n",
+      "recall  0.017221  0.195397  0.318430  0.109952\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "e1ad3bc2d94c407b995ac6696468895f",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 96, global step 18333: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.017056  0.059395  0.067978  0.047806\n",
+      "ndcg    0.017056  0.091108  0.122518  0.062947\n",
+      "recall  0.017056  0.196721  0.321245  0.109455\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "d4b13e0d6c38444fa1533a6ca3c51133",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 97, global step 18522: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.015731  0.057509  0.065688  0.046131\n",
+      "ndcg    0.015731  0.089842  0.119812  0.062111\n",
+      "recall  0.015731  0.197549  0.316443  0.111442\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "271f58b7cce546aba843acb2cf87ebe5",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 98, global step 18711: 'recall@10' was not in top 1\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k              1        10        20         5\n",
+      "map     0.014241  0.057645  0.065858  0.045747\n",
+      "ndcg    0.014241  0.090083  0.120261  0.061015\n",
+      "recall  0.014241  0.198212  0.318099  0.107799\n",
+      "\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "33fc04054c0942dba61679e0f8bda4b2",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Validation: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Epoch 99, global step 18900: 'recall@10' reached 0.20136 (best 0.20136), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=99-step=18900.ckpt' as top 1\n",
+      "`Trainer.fit` stopped: `max_epochs=100` reached.\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "k             1        10        20         5\n",
+      "map     0.01391  0.056347  0.064440  0.044130\n",
+      "ndcg    0.01391  0.089694  0.119503  0.059647\n",
+      "recall  0.01391  0.201358  0.319921  0.107468\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "from lightning.pytorch.callbacks import ModelCheckpoint\n",
+    "from lightning.pytorch.loggers import CSVLogger\n",
+    "\n",
+    "from replay.nn.lightning.callback import ComputeMetricsCallback\n",
+    "\n",
+    "\n",
+    "checkpoint_callback = ModelCheckpoint(\n",
+    "    dirpath=\"sasrec/checkpoints\",\n",
+    "    save_top_k=1,\n",
+    "    verbose=True,\n",
+    "    monitor=\"recall@10\",\n",
+    "    mode=\"max\",\n",
+    ")\n",
+    "\n",
+    "validation_metrics_callback = ComputeMetricsCallback(\n",
+    "    metrics=[\"map\", \"ndcg\", \"recall\"],\n",
+    "    ks=[1, 5, 10, 20],\n",
+    "    item_count=NUM_UNIQUE_ITEMS,\n",
+    ")\n",
+    "\n",
+    "csv_logger = CSVLogger(save_dir=\"sasrec/logs/train\", name=\"SasRec-example\")\n",
+    "\n",
+    "trainer = L.Trainer(\n",
+    "    max_epochs=100,\n",
+    "    callbacks=[checkpoint_callback, validation_metrics_callback],\n",
+    "    logger=csv_logger,\n",
+    ")\n",
+    "\n",
+    "trainer.fit(model, datamodule=parquet_module)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAiEAAAGnCAYAAABhFoffAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjgsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvwVt1zgAAAAlwSFlzAAASdAAAEnQB3mYfeAAAKj5JREFUeJzt3XtU1XW+//EXqNwFJUVQseyoo4MameDIcSRLPUzemJI0x6bSdI5aap3K0OUq82QWWmtaxtHmdNLRpouXdHRKU0edTpk7b53OiYZDjgoIomOGG4FUPr8//LGPtEEufvGziedjLRbD57Y/3/es5fe19vezd37GGCMAAIDrzN/2BgAAQPNECAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWtLS9AV9w9uxZ7dmzR7GxsQoMDLS9HQAAmozy8nLl5uYqOTlZbdq0qddcQoikPXv2KDU11fY2AABosjZu3KgxY8bUaw4hRFJsbKykywXs1q2bI2u63W65XC4lJiYqLCzMkTWbM+rpPGrqPGrqLOrpvMaoaU5OjlJTUz330voghEieRzDdunVTXFycI2sWFxersLBQvXr1Unh4uCNrNmfU03nU1HnU1FnU03mNWdOGHGfgYCoAALCCEAIAAKwghAAAACsIIQAAwApCCAAAsIIQAgAArCCEAAAAKwghAADACkIIAACwghACAACsIIQAAAArCCEAAMAKQggAALCCEAIAAKxoaXsDAADAGTc9/aer9kcHG6XHX5+91AXvhAAAACsIIQAAwApCCAAAsIIQAgAArCCEAAAAKwghAADACkIIAACwghACAACsIIQAAAArCCEAAMAKQggAALCCEAIAAKwghAAAACsIIQAAwApCCAAAsIIQAgAArCCEAAAAKwghAADACkIIAACwghACAACsIIQAAAArCCEAAMAKQggAALCCEAIAAKwghAAAACsIIQAAwApCCAAAsIIQAgAArCCEAAAAKwghAADACkIIAACwokEh5PPPP9cjjzyiuLg4hYaGqkuXLrr33nuVnZ3tNTYrK0spKSkKCwtTZGSk7r//fp06dcprXEVFhV566SV17dpVQUFB6tu3r95+++1qX7+uawIAAN/VsiGTXnzxRX3yySdKS0tT3759VVhYqGXLlqlfv3767LPP1Lt3b0lSXl6eBg8erIiICC1atEhut1tLlizRl19+KZfLpYCAAM+a8+bN0+LFizVlyhQlJCRo06ZNmjBhgvz8/DR+/HjPuPqsCQAAfFeDQsjjjz+uP/zhD1Vu+OPGjVOfPn20ePFirVmzRpK0aNEilZSU6MCBA+rSpYskKTExUcOGDdPKlSs1depUSVJ+fr6WLl2qGTNmaNmyZZKkhx9+WMnJyXryySeVlpamFi1a1GtNAADg2xr0OCYpKcnrHYfu3bsrLi5OWVlZnrb169dr5MiRnrAgSUOHDlWPHj303nvvedo2bdqkCxcuaPr06Z42Pz8/TZs2TXl5edq7d2+91wQAAL6tQe+EVMcYo5MnTyouLk7S5Xc3ioqK1L9/f6+xiYmJ+uCDDzx/Hzp0SKGhoerVq5fXuMr+QYMG1WvNmhQVFXmdH8nJyZEkud1uFRcX17pGXZSUlFT5jWtDPZ1HTZ1HTZ1FPesvOthctb9d0OV+J2vqdrsbPNexEPLWW28pPz9fzz33nCSpoKBAkhQTE+M1NiYmRmfOnFF5ebkCAwNVUFCgDh06yM/Pz2ucJJ04caLea9YkMzNTCxYsqLbP5XKpsLCwtkutF5fL5eh6zR31dB41dR41dRb1rLv0+LqNc7Kmx48fb/BcR0LI119/rRkzZmjgwIF64IEHJEmlpaWSVG0gCAoK8owJDAz0/L7auPquWZPp06crLS2tSltOTo5SU1OVmJjo9W5MQ5WUlMjlcikxMVGhoaGOrNmcUU/nUVPnUVNnUc/6G/7KX67a3y7IaErPCkdreuUxjPq65hBSWFioESNGKCIiQuvWrfMcIA0ODpYklZeXe80pKyurMiY4OLjO4+q6Zk2ioqIUFRVVbV9YWJjCw8OvOr++QkNDHV+zOaOezqOmzqOmzqKedVdY6lf7IDlb07CwsAbPvaYvK/vuu+/0i1/8QmfPntXWrVvVsWNHT1/lI5PKRyhXKigoUGRkpOcdi5iYGBUWFsoY4zVOkmfd+qwJAAB8W4NDSFlZmUaNGqXs7Gxt2bJFP/3pT6v0d+rUSe3bt9f+/fu95rpcLsXHx3v+jo+P1/nz573e0tm3b5+nv75rAgAA39agEHLp0iWNGzdOe/fu1dq1azVw4MBqx91zzz3asmWLcnNzPW07d+5UdnZ2lXMZY8aMUatWrZSZmelpM8Zo+fLl6tSpk5KSkuq9JgAA8G0NOhPyL//yL/rjH/+oUaNG6cyZM54vJ6s0ceJESdLcuXO1du1aDRkyRLNmzZLb7VZGRob69Omjhx56yDO+c+fOmj17tjIyMnThwgUlJCRo48aN+vjjj/XWW295zpnUZ00AAODbGhRCDh8+LEnavHmzNm/e7NVfGUJiY2O1Z88ePf7443r66acVEBCgESNGaOnSpV5nNxYvXqy2bdtqxYoVWrlypbp37641a9ZowoQJVcbVZ00AAOC7GhRCdu/eXeexcXFx2rZtW63j/P39lZ6ervT0dMfWBAAAvuuaPh0DAADQUIQQAABgBSEEAABYQQgBAABWEEIAAIAVhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWEEIAAIAVhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWEEIAAIAVhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWEEIAAIAVhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWEEIAAIAVhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWNDiEuN1uPfPMM0pJSVFkZKT8/Py0cuVKr3EPPvig/Pz8vH569uzpNbaiokIvvfSSunbtqqCgIPXt21dvv/12ta+flZWllJQUhYWFKTIyUvfff79OnTrV0MsBAADXWcuGTjx9+rSee+45denSRbfccot2795d49jAwED9+7//e5W2iIgIr3Hz5s3T4sWLNWXKFCUkJGjTpk2aMGGC/Pz8NH78eM+4vLw8DR48WBEREVq0aJHcbreWLFmiL7/8Ui6XSwEBAQ29LAAAcJ00OITExMSooKBA0dHR2r9/vxISEmp+kZYtNXHixKuul5+fr6VLl2rGjBlatmyZJOnhhx9WcnKynnzySaWlpalFixaSpEWLFqmkpEQHDhxQly5dJEmJiYkaNmyYVq5cqalTpzb0sgAAwHXS4McxgYGBio6OrvP4S5cuqbi4uMb+TZs26cKFC5o+fbqnzc/PT9OmTVNeXp727t3raV+/fr1GjhzpCSCSNHToUPXo0UPvvfdePa8EAADYcF0Opp4/f17h4eGKiIhQZGSkZsyYIbfbXWXMoUOHFBoaql69elVpT0xM9PRLl98xKSoqUv/+/b1eJzEx0TMOAAD4tgY/jqmrmJgYPfXUU+rXr58qKiq0detWZWZm6osvvtDu3bvVsuXlLRQUFKhDhw7y8/Pzmi9JJ06c8Iy7sv2HY8+cOaPy8nIFBgZWu5+ioiKvA6w5OTmSLh+2vdq7NfVRUlJS5TeuDfV0HjV1HjV1FvWsv+hgc9X+dkGX+52s6Q/fVKiPRg8hL7zwQpW/x48frx49emjevHlat26d58BpaWlptcEhKCjI03/l79rG1hRCMjMztWDBgmr7XC6XCgsL63JZdeZyuRxdr7mjns6jps6jps6innWXHl+3cU7W9Pjx4w2e2+ghpDqPPfaY5s+frx07dnhCSHBwsMrLy73GlpWVefqv/F2XsdWZPn260tLSqrTl5OQoNTVViYmJXo+DGqqkpEQul0uJiYkKDQ11ZM3mjHo6j5o6j5o6i3rW3/BX/nLV/nZBRlN6Vjha06ysrAbPtRJCgoODdcMNN+jMmTOetpiYGO3atUvGmCqPZCofv3Ts2NEz7sr2KxUUFCgyMrLGd0EkKSoqSlFRUdX2hYWFKTw8vP4XdBWhoaGOr9mcUU/nUVPnUVNnUc+6Kyz1q32QnK1pWFhYg+da+cbUc+fO6fTp02rfvr2nLT4+XufPn/dKVPv27fP0S1KnTp3Uvn177d+/32tdl8vlGQcAAHxbo4aQsrIynTt3zqt94cKFMsYoJSXF0zZmzBi1atVKmZmZnjZjjJYvX65OnTopKSnJ037PPfdoy5Ytys3N9bTt3LlT2dnZXo9aAACAb7qmxzHLli3T2bNnPZ9c2bx5s/Ly8iRJjz76qL799lvdeuutuu+++zxf075t2zZ98MEHSklJ0ZgxYzxrde7cWbNnz1ZGRoYuXLighIQEbdy4UR9//LHeeustzxeVSdLcuXO1du1aDRkyRLNmzZLb7VZGRob69Omjhx566FouyXHDX/lLjW+PHV084jrvBgAA33FNIWTJkiU6duyY5+8NGzZow4YNkqSJEyeqTZs2GjlypLZv365Vq1bp0qVL6tatmxYtWqQnnnhC/v5V34hZvHix2rZtqxUrVmjlypXq3r271qxZowkTJlQZFxsbqz179ujxxx/X008/rYCAAI0YMUJLly696nkQAADgO64phBw9erTWMatXr67zev7+/kpPT1d6enqtY+Pi4rRt27Y6rw0AAHyLlYOpAAAAhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWEEIAAIAVhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWEEIAAIAVhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWEEIAAIAVhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWEEIAAIAVhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWEEIAAIAVhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBUNDiFut1vPPPOMUlJSFBkZKT8/P61cubLasVlZWUpJSVFYWJgiIyN1//3369SpU17jKioq9NJLL6lr164KCgpS37599fbbb1/TmgAAwDe1bOjE06dP67nnnlOXLl10yy23aPfu3dWOy8vL0+DBgxUREaFFixbJ7XZryZIl+vLLL+VyuRQQEOAZO2/ePC1evFhTpkxRQkKCNm3apAkTJsjPz0/jx49v0JoAAMA3NTiExMTEqKCgQNHR0dq/f78SEhKqHbdo0SKVlJTowIED6tKliyQpMTFRw4YN08qVKzV16lRJUn5+vpYuXaoZM2Zo2bJlkqSHH35YycnJevLJJ5WWlqYWLVrUa00AAOC7Gvw4JjAwUNHR0bWOW79+vUaOHOkJC5I0dOhQ9ejRQ++9956nbdOmTbpw4YKmT5/uafPz89O0adOUl5envXv31ntNAADguxr1YGp+fr6KiorUv39/r77ExEQdOnTI8/ehQ4cUGhqqXr16eY2r7K/vmgAAwHc1+HFMXRQUFEi6/Ojmh2JiYnTmzBmVl5crMDBQBQUF6tChg/z8/LzGSdKJEyfqvWZ1ioqKvA6w5uTkSLp82La4uLg+l1ijkpISSVK7IFPjGKdeqzmorGflb1w7auo8auos6ll/0cE133Ok/7snOVlTt9vd4LmNGkJKS0slqdpAEBQU5BkTGBjo+X21cfVdszqZmZlasGBBtX0ul0uFhYVXvab6mtKzosa+Xbt2OfpazYHL5bK9hR8dauo8auos6ll36fF1G+dkTY8fP97guY0aQoKDgyVJ5eXlXn1lZWVVxgQHB9d5XF3XrM706dOVlpZWpS0nJ0epqalKTEz0ehzUUCUlJXK5XPrd1/46XeZX7ZiPHhvsyGs1B5X1TExMVGhoqO3t/ChQU+dRU2dRz/ob/spfrtrfLshoSs8KR2ualZXV4LmNGkIqH5lUPkK5UkFBgSIjIz3vWMTExGjXrl0yxlR5JFM5t2PHjvVeszpRUVGKioqqti8sLEzh4eF1ubQ6O13mp8LS6kOI06/VHISGhlI3h1FT51FTZ1HPuqvpfvNDTtY0LCyswXMb9WBqp06d1L59e+3fv9+rz+VyKT4+3vN3fHy8zp8/75Wo9u3b5+mv75oAAMB3NfrXtt9zzz3asmWLcnNzPW07d+5UdnZ2lcciY8aMUatWrZSZmelpM8Zo+fLl6tSpk5KSkuq9JgAA8F3X9Dhm2bJlOnv2rOeTK5s3b1ZeXp4k6dFHH1VERITmzp2rtWvXasiQIZo1a5bcbrcyMjLUp08fPfTQQ561OnfurNmzZysjI0MXLlxQQkKCNm7cqI8//lhvvfWW54vKJNV5TQAA4LuuKYQsWbJEx44d8/y9YcMGbdiwQZI0ceJERUREKDY2Vnv27NHjjz+up59+WgEBARoxYoSWLl3qdXZj8eLFatu2rVasWKGVK1eqe/fuWrNmjSZMmFBlXH3WBAAAvumaQsjRo0frNC4uLk7btm2rdZy/v7/S09OVnp7u2JoAAMA3NfqZEAAAgOoQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWEEIAAIAVhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWEEIAAIAVhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWEEIAAIAVhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWEEIAAIAVhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWEEIAAIAVhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVjR6CNm9e7f8/Pyq/fnss8+qjP300081aNAghYSEKDo6WjNnzpTb7fZas7y8XHPmzFHHjh0VHBysAQMGaPv27Y19KQAAwEEtr9cLzZw5UwkJCVXaunXr5vnfhw8f1p133qlevXrp5ZdfVl5enpYsWaL//d//1Ycfflhl3oMPPqh169Zp9uzZ6t69u1auXKm77rpLu3bt0qBBg67L9QAAgGtz3ULIz3/+c40dO7bG/rlz56pt27bavXu3wsPDJUk33XSTpkyZoo8++kjDhw+XJLlcLr3zzjvKyMjQE088IUn69a9/rd69e+upp57Sp59+2vgXAwAArtl1PRNy7tw5Xbx40au9uLhY27dv18SJEz0BRLocLsLCwvTee+952tatW6cWLVpo6tSpnragoCBNnjxZe/fuVW5ubuNeBAAAcMR1CyEPPfSQwsPDFRQUpCFDhmj//v2evi+//FIXL15U//79q8wJCAhQfHy8Dh065Gk7dOiQevToUSWsSFJiYqKky491AACA72v0xzEBAQG65557dNddd6ldu3b66quvtGTJEv385z/Xp59+qltvvVUFBQWSpJiYGK/5MTEx+vjjjz1/FxQU1DhOkk6cOHHV/RQVFenUqVNV2nJyciRJbrdbxcXF9bvAGpSUlEiS2gWZGsc49VrNQWU9K3/j2lFT51FTZ1HP+osOrvmeI/3fPcnJmlb3AZK6avQQkpSUpKSkJM/fo0eP1tixY9W3b1+lp6dr69atKi0tlSQFBgZ6zQ8KCvL0S1JpaWmN4yr7ryYzM1MLFiyots/lcqmwsLD2i6qHKT0rauzbtWuXo6/VHLhcLttb+NGhps6jps6innWXHl+3cU7W9Pjx4w2ee90Opl6pW7duGjNmjDZs2KBLly4pODhY0uWP3v5QWVmZp1+SgoODaxxX2X8106dPV1paWpW2nJwcpaamKjExUb169ar39VSnpKRELpdLv/vaX6fL/Kod89Fjgx15reagsp6JiYkKDQ21vZ0fBWrqPGrqLOpZf8Nf+ctV+9sFGU3pWeFoTbOysho810oIkaTY2Fh9//33Kikp8TxKqXwsc6WCggJ17NjR83dMTIzy8/OrHSepytjqREVFKSoqqtq+sLAwr7Mm1+p0mZ8KS6sPIU6/VnMQGhpK3RxGTZ1HTZ1FPeuupvvNDzlZ07CwsAbPtfaNqUeOHFFQUJDCwsLUu3dvtWzZssphVUn6/vvvdfjwYcXHx3va4uPjlZ2d7XWeYt++fZ5+AADg+xo9hPzwEKgkffHFF/rjH/+o4cOHy9/fXxERERo6dKjWrFmjc+fOecatXr1abre7yuOTsWPH6tKlS3r99dc9beXl5XrzzTc1YMAAxcbGNu4FAQAARzT645hx48YpODhYSUlJioqK0ldffaXXX39dISEhWrx4sWfc888/r6SkJCUnJ2vq1KnKy8vT0qVLNXz4cKWkpHjGDRgwQGlpaUpPT1dRUZG6deumVatW6ejRo3rjjTca+3IAAIBDGv2dkNTUVJ0+fVovv/yypk+frnfffVd333239u/fX+UQaL9+/bRjxw4FBwfrscce0+uvv67Jkydr3bp1Xmv+/ve/1+zZs7V69WrNnDlTFy5c0JYtWzR4MAc9AQBoKhr9nZCZM2dq5syZdRo7aNAgffLJJ7WOCwoKUkZGhjIyMq51ewAAwBJrB1MBAEDzRggBAABWEEIAAIAVhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWEEIAAIAVhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWEEIAAIAVhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWEEIAAIAVhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWEEIAAIAVhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWEEIAAIAVhBAAAGAFIQQAAFjRZENIeXm55syZo44dOyo4OFgDBgzQ9u3bbW8LAADUUZMNIQ8++KBefvll/epXv9Jvf/tbtWjRQnfddZf+8z//0/bWAABAHbS0vYGGcLlceuedd5SRkaEnnnhCkvTrX/9avXv31lNPPaVPP/3U8g4BAEBtmuQ7IevWrVOLFi00depUT1tQUJAmT56svXv3Kjc31+LuAABAXTTJd0IOHTqkHj16KDw8vEp7YmKiJOnw4cOKjY2tdm5RUZFOnTpVpe2rr76SJP3Xf/2X3G63I3ssLS3V8ePHFfidv8LLqx+zb98+R16rOais58GDBxUcHGx7Oz8K1NR51NRZ1LP+ws8dvWp/4PfS8eMVjtb0yJEjki6f1ayvJhlCCgoKFBMT49Ve2XbixIka52ZmZmrBggXV9k2YMMGZDdbRzzKv68sBAKC/NNK6ubm56tevX73mNMkQUlpaqsDAQK/2oKAgT39Npk+frrS0tCptxcXFys7OVp8+fapdtyFycnKUmpqqjRs3qlu3bo6s2ZxRT+dRU+dRU2dRT+c1Rk3Ly8uVm5ur5OTkes9tkiEkODi42rd9ysrKPP01iYqKUlRUlFf7wIEDndvgFbp166a4uLhGWbs5op7Oo6bOo6bOop7Oc7qm9X0HpFKTPJgaExOjgoICr/bKto4dO17vLQEAgHpqkiEkPj5e2dnZKi4urtJeedAzPj7ewq4AAEB9NMkQMnbsWF26dEmvv/66p628vFxvvvmmBgwYUOMnYwAAgO9okmdCBgwYoLS0NKWnp6uoqEjdunXTqlWrdPToUb3xxhu2tydJat++vZ555hm1b9/e9lZ+FKin86ip86ips6in83ytpn7GGGN7Ew1RVlam+fPna82aNfr222/Vt29fLVy4UP/0T/9ke2sAAKAOmmwIAQAATVuTPBMCAACaPkIIAACwghACAACsIIQAAAArCCH/38GDBzV69GhFRkYqJCREvXv31quvvurpv/322+Xn5+f1k5KSUmWdzz//XI888oji4uIUGhqqLl266N5771V2dnad93L27FlNnTpV7du3V2hoqIYMGaKDBw86dq3Xi6/UdOfOnZo0aZJ69OihkJAQ3XzzzXr44Yer/dZdX+Yr9fyhKVOmyM/PTyNHjrym67PB12q6Y8cO3XHHHYqIiFDr1q1122236d1333XkWq8XX6rpgQMHNHLkSEVHRyssLEx9+/bVq6++qkuXLjl2vY3NqXr+z//8j9LS0nTzzTcrJCRE7dq10+DBg7V58+Y676Ux7k1N8ntCnPbRRx9p1KhRuvXWWzV//nyFhYXpm2++UV5eXpVxnTt31gsvvFCl7YdfEf/iiy/qk08+UVpamvr27avCwkItW7ZM/fr102effabevXtfdS8VFRUaMWKEvvjiCz355JNq166dMjMzdfvtt+vAgQPq3r27MxfdyHyppnPmzNGZM2eUlpam7t2768iRI1q2bJm2bNmiw4cPKzo62pmLbkS+VM8r7d+/XytXrvT8xyObEl+r6ZtvvqnJkydr2LBhWrRokVq0aKG//vWvys3NvfaLvU58qaYHDhxQUlKSunfvrjlz5igkJEQffvihZs2apW+++Ua//e1vnbnoRuRkPY8dO6Zz587pgQceUMeOHXX+/HmtX79eo0eP1ooVKzR16tSr7qXR7k2mmfvuu+9Mhw4dzC9/+Utz6dKlGsclJyebuLi4Wtf75JNPTHl5eZW27OxsExgYaH71q1/VOv/dd981kszatWs9bUVFRaZNmzbmvvvuq3W+L/C1mu7Zs8drH3v27DGSzLx582qdb5uv1bNSRUWFGThwoJk0aZK58cYbzYgRI+o81zZfq+nf/vY3ExwcbGbOnFn75n2Ur9V0ypQpJiAgwPz973+v0j548GATHh5e63zbnK5ndS5evGhuueUW85Of/KTWsY11b2r2j2P+8Ic/6OTJk3r++efl7++vkpISVVRU1Dj+4sWLcrvdNfYnJSUpICCgSlv37t0VFxenrKysWvezbt06dejQQXfffbenrX379rr33nu1adOmav/rwb7G12o6ePBg+fv7e7VFRkbWab5tvlbPSqtXr9Z///d/6/nnn6/zHF/hazVdvny5Ll26pOeee06S5Ha7ZZrYVzj5Wk2Li4sVFBSkNm3aVGmPiYm56n9p3Vc4Xc/qtGjRQrGxsTp79mytYxvr3tTsQ8iOHTsUHh6u/Px8/eQnP1FYWJjCw8M1bdo0lZWVVRmbnZ2t0NBQtW7dWtHR0Zo/f74uXLhQ62sYY3Ty5Em1a9eu1rGHDh1Sv379vG6aiYmJOn/+fIOf219PvlbT6rjdbrnd7gbPv558sZ7nzp3TnDlzNHfu3CbxOOuHfK2mO3bsUM+ePfXBBx+oc+fOat26tW644QbNnz//qjceX+JrNb399ttVXFys3/zmN8rKytKxY8e0fPlybdiwQenp6Q2+zuulsepZUlKi06dP65tvvtErr7yiDz/8UHfeeWet+2m0e1OD30P5kejbt68JCQkxISEh5tFHHzXr1683jz76qJFkxo8f7xk3adIk8+yzz5r169eb3//+92b06NFGkrn33ntrfY3Vq1cbSeaNN96odWxoaKiZNGmSV/uf/vQnI8ls3bq1fhdoga/VtDoLFy40kszOnTsbNP968sV6PvHEE6Zr166mrKzMGGOa3OMYX6tpeHi4adu2rQkMDDTz588369atMxMmTDCSzNNPP31N13q9+FpNL168aB555BHTqlUrI8lIMi1atDD/9m//dk3Xeb00Vj1/85vfeOrh7+9vxo4da86cOVPrfhrr3tTsQ8jNN99sJJl//ud/rtJe+X9UdnZ2jXOnTJliJJm9e/fWOCYrK8uEh4ebgQMHmosXL9a6H39/fzNt2jSv9p07dxpJ5v333691Ddt8raY/tGfPHtOyZcs6/aPnC3ytnn/9619Nq1atzLp16zxtTS2E+FpN/f39jSSzePHiKu0pKSkmODjYFBcX17qGbb5WU2OMeeWVV8zIkSPNqlWrzLvvvmtSU1NNy5Ytm/W/o1lZWWb79u1m1apVZsSIEeaXv/ylKSwsrHU/jXVvavYhJC4uzkgye/bsqdJeeXBx1apVNc79+uuvjSSzcOHCavsLCgrMzTffbGJjY01+fn6d9vNjeCfE12p6paysLBMZGWni4+ObxD/sxvhePVNSUkxycnKVtqYWQnytpqGhoUaSOXbsWJX2VatWVbtPX+RrNX3hhRdMdHS0OXfuXJX222+/3XTs2NFcuHChTuvY0pj1vNKwYcNMQkKCqaiouOq4xro3NfszIZUfY+rQoUOV9qioKEnSt99+W+Pc2NhYSdKZM2e8+r777jv94he/0NmzZ7V161avj0vVJCYmptrvr6hsq+s6NvlaTSvl5uZq+PDhioiI0AcffKDWrVvXa74tvlTPP//5z9q6datmzZqlo0ePen4uXryo0tJSHT16VMXFxXW+Nlt8qabXuh9f4Ws1zczM1B133KGwsLAq7aNHj9aJEyd09OjROq1jS2PV84fGjh2rzz//vNYzHY11b2r2IeS2226TJOXn51dpP3HihKTLp39rcuTIkWrHlJWVadSoUcrOztaWLVv005/+tM77iY+P18GDB70Oo+3bt08hISHq0aNHndeyxddqKkl///vfNXz4cJWXl2vbtm2KiYmp13ybfKmex48flyTdfffd6tq1q+cnPz9ff/7zn9W1a1f9x3/8R90uzCJfqum17sdX+FpNT548We2XklUe2Lx48WKd17KhMepZndLSUkmXw97VNNq9qUHvn/yIHDx40EgyEyZMqNJ+3333mZYtW5r8/Hzz3XffeQ7gVaqoqDDjxo0zksyBAwc87RcvXjSjR482LVu2NH/605+u+tonTpwwWVlZ5vvvv/e0vfPOO16fxT516pRp06aNGTdu3LVc6nXjazV1u90mMTHRtG7d2uzfv9+BK7y+fKmex44dM++//77XT/v27U3//v3N+++/b3Jychy68sbjSzU1xpj333/fSDJz5871tF26dMkMGjTIREZGeu3DF/laTXv37m0iIyPN6dOnq6x52223mdatW1cZ64ucrufJkye9XuP77783/fr1M8HBwVUeW13Pe1OzDyHGXD5drP9/mvi1114zaWlpRpJJT083xhiza9cuEx0dbR577DHz2muvmSVLlph//Md/NJLM1KlTq6w1a9YsI8mMGjXKrF692uvnSg888ICRZP72t7952i5evGh+9rOfmbCwMLNgwQLz2muvmbi4ONO6dWvz9ddfN3otnOJLNR0zZoyRZCZNmuQ1tykcUDPGt+pZnaZ2JsQY36ppRUWFufPOO42fn5+ZOnWqee2118ywYcOMJLNixYpGr4VTfKmma9asMZLMP/zDP5gXX3zRvPrqq2bgwIFGkvnXf/3XRq+FE5ysZ2pqqrnjjjvMs88+a373u9+ZhQsXmp49expJZunSpVXGXs97EyHEXE6Dzz77rLnxxhtNq1atTLdu3cwrr7zi6T9y5IhJS0szN910kwkKCjIhISHmtttuM8uXL/c6zJOcnOz5+FN1P1eq6R/4M2fOmMmTJ5sbbrjBhISEmOTkZPP555831uU3Cl+q6Y033ljj3BtvvLERq+AcX6pndZpiCPG1mp47d87MmjXLREdHm4CAANOnTx+zZs2axrr8RuFrNd26datJTk427dq189R0+fLljXX5jnOynm+//bYZOnSo6dChg2nZsqVp27atGTp0qNm0aZPX617Pe5OfMU3sa/kAAMCPQrM/mAoAAOwghAAAACsIIQAAwApCCAAAsIIQAgAArCCEAAAAKwghAADACkIIAACwghACAACsIIQAAAArCCEAAMAKQggAALCCEAIAAKwghAAAACsIIQAAwIr/B8ppMEdqEMoaAAAAAElFTkSuQmCC",
+      "text/plain": [
+       "<Figure size 600x480 with 1 Axes>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "import matplotlib.pyplot as plt\n",
+    "\n",
+    "plt.figure(figsize = (5, 4), dpi = 120)\n",
+    "plt.hist(transforms[\"train\"][0].frequencies.cpu().numpy(), bins = 50)\n",
+    "plt.grid()\n"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Now we can get the best model path stored in the checkpoint callback."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "'/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=99-step=18900.ckpt'"
+      ]
+     },
+     "execution_count": 22,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "best_model_path = checkpoint_callback.best_model_path\n",
+    "best_model_path"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Inference\n",
+    "\n",
+    "To obtain model scores, we will load the weights from the best checkpoint. To do this, we use the `LightningModule`, provide there the path to the checkpoint and the model instance."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 23,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import replay\n",
+    "torch.serialization.add_safe_globals([\n",
+    "    replay.nn.lightning.optimizer.OptimizerFactory,\n",
+    "    replay.nn.lightning.scheduler.LRSchedulerFactory\n",
+    "])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 24,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "sasrec = make_sasrec(\n",
+    "    schema=tensor_schema,\n",
+    "    embedding_dim=EMBEDDING_DIM,\n",
+    "    max_sequence_length=MAX_SEQ_LEN,\n",
+    "    num_heads=NUM_HEADS,\n",
+    "    num_blocks=NUM_BLOCKS,\n",
+    "    dropout=DROPOUT,\n",
+    ")\n",
+    "\n",
+    "best_model = LightningModule.load_from_checkpoint(best_model_path, model=sasrec)\n",
+    "best_model.eval();"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Configure `ParquetModule` for inference"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 25,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "inference_metadata = {\"predict\": create_meta(shape=MAX_SEQ_LEN)}\n",
+    "\n",
+    "parquet_module = ParquetModule(\n",
+    "    predict_path=PREDICT_PATH,\n",
+    "    batch_size=BATCH_SIZE,\n",
+    "    metadata=inference_metadata,\n",
+    "    transforms=transforms,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "During inference, we can use `TopItemsCallback`. Such callback allows you to get scores for each user throughout the entire catalog and get recommendations in the form of ids of items with the highest score values.\n",
+    "\n",
+    "\n",
+    "Recommendations can be fetched in four formats: PySpark DataFrame, Pandas DataFrame, Polars DataFrame or raw PyTorch tensors. Each of the types corresponds a callback. In this example, we'll be using the `PandasTopItemsCallback`."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 26,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "💡 Tip: For seamless cloud uploads and versioning, try installing [litmodels](https://pypi.org/project/litmodels/) to enable LitModelCheckpoint, which syncs automatically with the Lightning model registry.\n",
+      "GPU available: True (cuda), used: True\n",
+      "TPU available: False, using: 0 TPU cores\n",
+      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n"
+     ]
+    },
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "94a84cd4fca248e7b340909e2048450f",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "Predicting: |          | 0/? [00:00<?, ?it/s]"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    }
+   ],
+   "source": [
+    "from replay.nn.lightning.callback import PandasTopItemsCallback\n",
+    "\n",
+    "csv_logger = CSVLogger(save_dir=\"sasrec/logs/test\", name=\"SasRec-example\")\n",
+    "\n",
+    "TOPK = [1, 5, 10, 20]\n",
+    "\n",
+    "pandas_prediction_callback = PandasTopItemsCallback(\n",
+    "    top_k=max(TOPK),\n",
+    "    query_column=\"user_id\",\n",
+    "    item_column=\"item_id\",\n",
+    "    rating_column=\"score\",\n",
+    ")\n",
+    "\n",
+    "trainer = L.Trainer(callbacks=[pandas_prediction_callback], logger=csv_logger, inference_mode=True)\n",
+    "trainer.predict(best_model, datamodule=parquet_module, return_predictions=False)\n",
+    "\n",
+    "pandas_res = pandas_prediction_callback.get_result()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 27,
+   "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>486</td>\n",
+       "      <td>7.196251</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>0</td>\n",
+       "      <td>101</td>\n",
+       "      <td>7.086551</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>0</td>\n",
+       "      <td>210</td>\n",
+       "      <td>6.872698</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>0</td>\n",
+       "      <td>1613</td>\n",
+       "      <td>6.834656</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>0</td>\n",
+       "      <td>357</td>\n",
+       "      <td>6.555361</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>...</th>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6037</th>\n",
+       "      <td>6039</td>\n",
+       "      <td>84</td>\n",
+       "      <td>11.528885</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6037</th>\n",
+       "      <td>6039</td>\n",
+       "      <td>2700</td>\n",
+       "      <td>11.400557</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6037</th>\n",
+       "      <td>6039</td>\n",
+       "      <td>2557</td>\n",
+       "      <td>11.273298</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6037</th>\n",
+       "      <td>6039</td>\n",
+       "      <td>2634</td>\n",
+       "      <td>11.121119</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6037</th>\n",
+       "      <td>6039</td>\n",
+       "      <td>2633</td>\n",
+       "      <td>11.119842</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "<p>120760 rows × 3 columns</p>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "      user_id item_id      score\n",
+       "0           0     486   7.196251\n",
+       "0           0     101   7.086551\n",
+       "0           0     210   6.872698\n",
+       "0           0    1613   6.834656\n",
+       "0           0     357   6.555361\n",
+       "...       ...     ...        ...\n",
+       "6037     6039      84  11.528885\n",
+       "6037     6039    2700  11.400557\n",
+       "6037     6039    2557  11.273298\n",
+       "6037     6039    2634  11.121119\n",
+       "6037     6039    2633  11.119842\n",
+       "\n",
+       "[120760 rows x 3 columns]"
+      ]
+     },
+     "execution_count": 27,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "pandas_res"
+   ]
+  },
+  {
+   "attachments": {},
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Calculating metrics\n",
+    "\n",
+    "*test_gt* is already encoded, so we can use it for computing metrics."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 28,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from replay.metrics import MAP, OfflineMetrics, Precision, Recall\n",
+    "from replay.metrics.torch_metrics_builder import metrics_to_df"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 29,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "result_metrics = OfflineMetrics(\n",
+    "    [Recall(TOPK), Precision(TOPK), MAP(TOPK)],\n",
+    "    query_column=\"user_id\",\n",
+    "    rating_column=\"score\",\n",
+    ")(pandas_res, test_gt.explode(\"item_id\"))"
+   ]
+  },
+  {
+   "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>5</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>MAP</th>\n",
+       "      <td>0.016065</td>\n",
+       "      <td>0.056637</td>\n",
+       "      <td>0.064067</td>\n",
+       "      <td>0.045948</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>Precision</th>\n",
+       "      <td>0.016065</td>\n",
+       "      <td>0.018665</td>\n",
+       "      <td>0.014798</td>\n",
+       "      <td>0.021133</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>Recall</th>\n",
+       "      <td>0.016065</td>\n",
+       "      <td>0.186651</td>\n",
+       "      <td>0.295959</td>\n",
+       "      <td>0.105664</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "k                 1        10        20         5\n",
+       "MAP        0.016065  0.056637  0.064067  0.045948\n",
+       "Precision  0.016065  0.018665  0.014798  0.021133\n",
+       "Recall     0.016065  0.186651  0.295959  0.105664"
+      ]
+     },
+     "execution_count": 30,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "metrics_to_df(result_metrics)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Let's call the `inverse_transform` encoder's function to get the final dataframe with recommendations"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 31,
+   "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>2012</td>\n",
+       "      <td>500</td>\n",
+       "      <td>7.196251</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>2012</td>\n",
+       "      <td>104</td>\n",
+       "      <td>7.086551</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>2012</td>\n",
+       "      <td>216</td>\n",
+       "      <td>6.872698</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>2012</td>\n",
+       "      <td>1777</td>\n",
+       "      <td>6.834656</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>2012</td>\n",
+       "      <td>367</td>\n",
+       "      <td>6.555361</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>...</th>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "      <td>...</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6037</th>\n",
+       "      <td>5727</td>\n",
+       "      <td>86</td>\n",
+       "      <td>11.528885</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6037</th>\n",
+       "      <td>5727</td>\n",
+       "      <td>2907</td>\n",
+       "      <td>11.400557</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6037</th>\n",
+       "      <td>5727</td>\n",
+       "      <td>2762</td>\n",
+       "      <td>11.273298</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6037</th>\n",
+       "      <td>5727</td>\n",
+       "      <td>2841</td>\n",
+       "      <td>11.121119</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>6037</th>\n",
+       "      <td>5727</td>\n",
+       "      <td>2840</td>\n",
+       "      <td>11.119842</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "<p>120760 rows × 3 columns</p>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "      user_id  item_id      score\n",
+       "0        2012      500   7.196251\n",
+       "0        2012      104   7.086551\n",
+       "0        2012      216   6.872698\n",
+       "0        2012     1777   6.834656\n",
+       "0        2012      367   6.555361\n",
+       "...       ...      ...        ...\n",
+       "6037     5727       86  11.528885\n",
+       "6037     5727     2907  11.400557\n",
+       "6037     5727     2762  11.273298\n",
+       "6037     5727     2841  11.121119\n",
+       "6037     5727     2840  11.119842\n",
+       "\n",
+       "[120760 rows x 3 columns]"
+      ]
+     },
+     "execution_count": 31,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "encoder.inverse_transform(pandas_res)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "new_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.12.3"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/replay/nn/transform/__init__.py b/replay/nn/transform/__init__.py
index 355cf8351..40c2cdec1 100644
--- a/replay/nn/transform/__init__.py
+++ b/replay/nn/transform/__init__.py
@@ -1,6 +1,14 @@
 from .copy import CopyTransform
 from .grouping import GroupTransform
+<<<<<<< HEAD
 from .negative_sampling import MultiClassNegativeSamplingTransform, UniformNegativeSamplingTransform
+=======
+from .negative_sampling import (
+    UniformNegativeSamplingTransform, 
+    FrequencyNegativeSamplingTransform,
+    ThresholdNegativeSamplingTransform,
+)
+>>>>>>> 36b25601 (Gitignore updated to recent examples)
 from .next_token import NextTokenTransform
 from .rename import RenameTransform
 from .reshape import UnsqueezeTransform
@@ -18,5 +26,7 @@
     "TokenMaskTransform",
     "TrimTransform",
     "UniformNegativeSamplingTransform",
+    "FrequencyNegativeSamplingTransform",
+    "ThresholdNegativeSamplingTransform",
     "UnsqueezeTransform",
 ]
diff --git a/replay/nn/transform/negative_sampling.py b/replay/nn/transform/negative_sampling.py
index 4d1364eda..34debab8e 100644
--- a/replay/nn/transform/negative_sampling.py
+++ b/replay/nn/transform/negative_sampling.py
@@ -1,6 +1,9 @@
-from typing import Optional
+from typing import Optional, Literal, cast
 
 import torch
+import torch.nn.functional as func
+
+import warnings
 
 
 class UniformNegativeSamplingTransform(torch.nn.Module):
@@ -29,7 +32,7 @@ def __init__(
         cardinality: int,
         num_negative_samples: int,
         *,
-        out_feature_name: Optional[str] = "negative_labels",
+        out_feature_name: str = "negative_labels",
         sample_distribution: Optional[torch.Tensor] = None,
         generator: Optional[torch.Generator] = None,
     ) -> None:
@@ -43,12 +46,20 @@ def __init__(
         :param generator: Random number generator to be used for sampling
                 from the distribution. Default: ``None``.
         """
-        if sample_distribution is not None and sample_distribution.size(-1) != cardinality:
-            msg = (
-                "The sample_distribution parameter has an incorrect size. "
-                f"Got {sample_distribution.size(-1)}, expected {cardinality}."
-            )
-            raise ValueError(msg)
+        if sample_distribution is not None:
+            if sample_distribution.ndim != 1:
+                msg: str = (
+                    "The `sample_distribution` parameter must be 1D."
+                    f"Got {sample_distribution.ndim}, will be flattened."
+                )
+                warnings.warn(msg)
+                sample_distribution = sample_distribution.flatten()
+            if sample_distribution.size(-1) != cardinality:
+                msg: str = (
+                    "The sample_distribution parameter has an incorrect size. "
+                    f"Got {sample_distribution.size(-1)}, expected {cardinality}."
+                )
+                raise ValueError(msg)
 
         if num_negative_samples >= cardinality:
             msg = (
@@ -63,9 +74,13 @@ def __init__(
         self.num_negative_samples = num_negative_samples
         self.generator = generator
         if sample_distribution is not None:
-            self.sample_distribution = sample_distribution
+            sample_distribution = sample_distribution
         else:
-            self.sample_distribution = torch.ones(cardinality)
+            sample_distribution = torch.ones(cardinality)
+        
+        self.sample_distribution = torch.nn.Buffer(
+            cast(torch.Tensor, sample_distribution)
+        )
 
     def forward(self, batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
         output_batch = dict(batch.items())
@@ -77,7 +92,186 @@ def forward(self, batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
             generator=self.generator,
         )
 
-        output_batch[self.out_feature_name] = negatives.to(device=next(iter(output_batch.values())).device)
+        device = next(iter(output_batch.values())).device
+        output_batch[self.out_feature_name] = negatives.to(device)
+        return output_batch
+
+class FrequencyNegativeSamplingTransform(torch.nn.Module):
+    """
+    Transform for global negative sampling.
+
+    For every batch, transform generates a vector of size ``(num_negative_samples)``
+    consisting of random indices sampeled from a range of ``cardinality``. Unless a custom sample
+    distribution is provided, the indices are weighted equally.
+
+    Example:
+
+    .. code-block:: python
+
+        >>> _ = torch.manual_seed(0)
+        >>> input_batch = {"item_id": torch.LongTensor([[1, 0, 4]])}
+        >>> transform = UniformNegativeSamplingTransform(cardinality=4, num_negative_samples=2)
+        >>> output_batch = transform(input_batch)
+        >>> output_batch
+        {'item_id': tensor([[1, 0, 4]]), 'negative_labels': tensor([2, 1])}
+
+    """
+
+    def __init__(
+        self,
+        cardinality: int,
+        num_negative_samples: int,
+        *,
+        out_feature_name: str = "negative_labels",
+        generator: Optional[torch.Generator] = None,
+        mode: Literal["softmax", "softsum"] = "softmax",
+    ) -> None:
+        """
+        :param cardinality: The size of sample vocabulary.
+        :param num_negative_samples: The size of negatives vector to generate.
+        :param out_feature_name: The name of result feature in batch.
+        :param sample_distribution: The weighs of indices in the vocabulary. If specified, must
+                match the ``cardinality``. Default: ``None``.
+        :param generator: Random number generator to be used for sampling
+                from the distribution. Default: ``None``.
+        """
+        assert num_negative_samples < cardinality
+
+        super().__init__()
+
+        self.cardinality = cardinality
+        self.out_feature_name = out_feature_name
+        self.num_negative_samples = num_negative_samples
+        self.generator = generator
+        self.mode = mode
+
+        self.frequencies = torch.nn.Buffer(
+            torch.zeros(cardinality, dtype = torch.int64)
+        )
+
+    def get_probas(self) -> torch.Tensor:
+        raw: torch.Tensor = 1.0 / (1.0 + self.frequencies)
+        match self.mode:
+            case "softsum":
+                result: torch.Tensor = raw / torch.sum(raw)
+            case "softmax":
+                result: torch.Tensor = func.softmax(raw, dim = -1)
+            case _:
+                msg: str = f"Unsupported mode: {self.mode}."
+                raise TypeError(msg)
+        return result
+    
+    def update_probas(self, selected: torch.Tensor) -> None:
+        device = self.frequencies.device
+        one = torch.ones(1, dtype = torch.int64, device = device)
+        self.frequencies.index_add_(-1, selected, one.expand(selected.numel()))
+
+    def forward(self, batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+        output_batch = dict(batch.items())
+
+        negatives = torch.multinomial(
+            input = self.get_probas(),
+            num_samples=self.num_negative_samples,
+            replacement=False,
+            generator=self.generator,
+        )
+
+        self.update_probas(negatives)
+
+        device = next(iter(output_batch.values())).device
+        output_batch[self.out_feature_name] = negatives.to(device)
+        return output_batch
+
+class ThresholdNegativeSamplingTransform(torch.nn.Module):
+    """
+    Transform for global negative sampling.
+
+    For every batch, transform generates a vector of size ``(num_negative_samples)``
+    consisting of random indices sampeled from a range of ``cardinality``. Unless a custom sample
+    distribution is provided, the indices are weighted equally.
+
+    Example:
+
+    .. code-block:: python
+
+        >>> _ = torch.manual_seed(0)
+        >>> input_batch = {"item_id": torch.LongTensor([[1, 0, 4]])}
+        >>> transform = UniformNegativeSamplingTransform(cardinality=4, num_negative_samples=2)
+        >>> output_batch = transform(input_batch)
+        >>> output_batch
+        {'item_id': tensor([[1, 0, 4]]), 'negative_labels': tensor([2, 1])}
+
+    """
+
+    def __init__(
+        self,
+        cardinality: int,
+        num_negative_samples: int,
+        *,
+        out_feature_name: str = "negative_labels",
+        generator: Optional[torch.Generator] = None,
+        mode: Literal["softmax", "softsum"] = "softmax",
+    ) -> None:
+        """
+        :param cardinality: The size of sample vocabulary.
+        :param num_negative_samples: The size of negatives vector to generate.
+        :param out_feature_name: The name of result feature in batch.
+        :param sample_distribution: The weighs of indices in the vocabulary. If specified, must
+                match the ``cardinality``. Default: ``None``.
+        :param generator: Random number generator to be used for sampling
+                from the distribution. Default: ``None``.
+        """
+        assert num_negative_samples < cardinality
+
+        super().__init__()
+
+        self.cardinality = cardinality
+        self.out_feature_name = out_feature_name
+        self.num_negative_samples = num_negative_samples
+        self.generator = generator
+        self.mode = mode
+
+        self.frequencies = torch.nn.Buffer(
+            torch.zeros(cardinality, dtype = torch.int64)
+        )
+
+    def get_probas(self) -> torch.Tensor:
+        raw: torch.Tensor = 1.0 / (1.0 + self.frequencies)
+        thr: torch.Tensor = torch.max(self.frequencies)
+        mask: torch.Tensor = thr != self.frequencies
+        match self.mode:
+            case "softsum":
+                eps = torch.finfo(raw.dtype).eps
+                raw = torch.where(mask, raw, eps)
+                result: torch.Tensor = raw / torch.sum(raw)
+            case "softmax":
+                inf = torch.finfo(raw.dtype).min
+                raw = torch.where(mask, raw, inf)
+                result: torch.Tensor = func.softmax(raw, dim = -1)
+            case _:
+                msg: str = f"Unsupported mode: {self.mode}."
+                raise TypeError(msg)
+        return result
+    
+    def update_probas(self, selected: torch.Tensor) -> None:
+        device = self.frequencies.device
+        one = torch.ones(1, dtype = torch.int64, device = device)
+        self.frequencies.index_add_(-1, selected, one.expand(selected.numel()))
+
+    def forward(self, batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+        output_batch = dict(batch.items())
+
+        negatives = torch.multinomial(
+            input = self.get_probas(),
+            num_samples=self.num_negative_samples,
+            replacement=False,
+            generator=self.generator,
+        )
+
+        self.update_probas(negatives)
+
+        device = next(iter(output_batch.values())).device
+        output_batch[self.out_feature_name] = negatives.to(device)
         return output_batch
 
 
@@ -124,8 +318,8 @@ def __init__(
         num_negative_samples: int,
         sample_mask: torch.Tensor,
         *,
-        negative_selector_name: Optional[str] = "negative_selector",
-        out_feature_name: Optional[str] = "negative_labels",
+        negative_selector_name: str = "negative_selector",
+        out_feature_name: str = "negative_labels",
         generator: Optional[torch.Generator] = None,
     ) -> None:
         """
@@ -153,7 +347,9 @@ def __init__(
 
         super().__init__()
 
-        self.register_buffer("sample_mask", sample_mask.float())
+        self.sample_mask = torch.nn.Buffer(
+            sample_mask.float()
+        )
 
         self.num_negative_samples = num_negative_samples
         self.negative_selector_name = negative_selector_name

From 584b54a46fcafd0943b8178f21b5c178c17856f6 Mon Sep 17 00:00:00 2001
From: Nikita Kulikov <neilkulikov@gmail.com>
Date: Mon, 2 Feb 2026 07:16:13 +0300
Subject: [PATCH 2/5] Unnecessary files removed

---
 examples/09_sasrec_example_freq_sampled.ipynb | 4768 -----------------
 examples/09_sasrec_example_sampled.ipynb      | 4744 ----------------
 examples/09_sasrec_example_thr_sampled.ipynb  | 4768 -----------------
 3 files changed, 14280 deletions(-)
 delete mode 100644 examples/09_sasrec_example_freq_sampled.ipynb
 delete mode 100644 examples/09_sasrec_example_sampled.ipynb
 delete mode 100644 examples/09_sasrec_example_thr_sampled.ipynb

diff --git a/examples/09_sasrec_example_freq_sampled.ipynb b/examples/09_sasrec_example_freq_sampled.ipynb
deleted file mode 100644
index 4988e2512..000000000
--- a/examples/09_sasrec_example_freq_sampled.ipynb
+++ /dev/null
@@ -1,4768 +0,0 @@
-{
- "cells": [
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Example of SasRec training/inference with Parquet Module"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 1,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Seed set to 42\n"
-     ]
-    }
-   ],
-   "source": [
-    "from typing import Optional\n",
-    "\n",
-    "import lightning as L\n",
-    "import pandas as pd\n",
-    "\n",
-    "L.seed_everything(42)\n",
-    "\n",
-    "import warnings\n",
-    "warnings.filterwarnings(\"ignore\")"
-   ]
-  },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Preparing data\n",
-    "In this example, we will be using the MovieLens dataset, namely the 1m subset. It's demonstrated a simple case, so only item ids will be used as model input.\n",
-    "\n",
-    "---\n",
-    "**NOTE**\n",
-    "\n",
-    "Current implementation of SasRec is able to handle item and interactions features. It does not take into account user features. \n",
-    "\n",
-    "---"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 2,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "interactions = pd.read_csv(\"./data/ml1m_ratings.dat\", sep=\"\\t\", names=[\"user_id\", \"item_id\",\"rating\",\"timestamp\"])\n",
-    "interactions = interactions.drop(columns=[\"rating\"])"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 3,
-   "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>timestamp</th>\n",
-       "    </tr>\n",
-       "  </thead>\n",
-       "  <tbody>\n",
-       "    <tr>\n",
-       "      <th>1000138</th>\n",
-       "      <td>6040</td>\n",
-       "      <td>858</td>\n",
-       "      <td>0</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1000153</th>\n",
-       "      <td>6040</td>\n",
-       "      <td>2384</td>\n",
-       "      <td>1</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>999873</th>\n",
-       "      <td>6040</td>\n",
-       "      <td>593</td>\n",
-       "      <td>2</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1000007</th>\n",
-       "      <td>6040</td>\n",
-       "      <td>1961</td>\n",
-       "      <td>3</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1000192</th>\n",
-       "      <td>6040</td>\n",
-       "      <td>2019</td>\n",
-       "      <td>4</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>...</th>\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>446</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>825438</th>\n",
-       "      <td>4958</td>\n",
-       "      <td>1407</td>\n",
-       "      <td>447</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>825724</th>\n",
-       "      <td>4958</td>\n",
-       "      <td>3264</td>\n",
-       "      <td>448</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>825731</th>\n",
-       "      <td>4958</td>\n",
-       "      <td>2634</td>\n",
-       "      <td>449</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>825603</th>\n",
-       "      <td>4958</td>\n",
-       "      <td>1924</td>\n",
-       "      <td>450</td>\n",
-       "    </tr>\n",
-       "  </tbody>\n",
-       "</table>\n",
-       "<p>1000209 rows × 3 columns</p>\n",
-       "</div>"
-      ],
-      "text/plain": [
-       "         user_id  item_id  timestamp\n",
-       "1000138     6040      858          0\n",
-       "1000153     6040     2384          1\n",
-       "999873      6040      593          2\n",
-       "1000007     6040     1961          3\n",
-       "1000192     6040     2019          4\n",
-       "...          ...      ...        ...\n",
-       "825793      4958     2399        446\n",
-       "825438      4958     1407        447\n",
-       "825724      4958     3264        448\n",
-       "825731      4958     2634        449\n",
-       "825603      4958     1924        450\n",
-       "\n",
-       "[1000209 rows x 3 columns]"
-      ]
-     },
-     "execution_count": 3,
-     "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"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### Encode catagorical data.\n",
-    "To ensure all categorical data is fit for training, it needs to be encoded using the `LabelEncoder` class. Create an instance of the encoder, providing a `LabelEncodingRule` for each categorcial column in the dataset that will be used in model. Note that ids of users and ids of items are always used."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 4,
-   "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>timestamp</th>\n",
-       "      <th>user_id</th>\n",
-       "      <th>item_id</th>\n",
-       "    </tr>\n",
-       "  </thead>\n",
-       "  <tbody>\n",
-       "    <tr>\n",
-       "      <th>0</th>\n",
-       "      <td>12</td>\n",
-       "      <td>0</td>\n",
-       "      <td>0</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1</th>\n",
-       "      <td>68</td>\n",
-       "      <td>1</td>\n",
-       "      <td>0</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>2</th>\n",
-       "      <td>67</td>\n",
-       "      <td>2</td>\n",
-       "      <td>0</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>3</th>\n",
-       "      <td>12</td>\n",
-       "      <td>3</td>\n",
-       "      <td>0</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>4</th>\n",
-       "      <td>140</td>\n",
-       "      <td>4</td>\n",
-       "      <td>0</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>...</th>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1000204</th>\n",
-       "      <td>14</td>\n",
-       "      <td>4555</td>\n",
-       "      <td>3705</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1000205</th>\n",
-       "      <td>90</td>\n",
-       "      <td>2813</td>\n",
-       "      <td>3705</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1000206</th>\n",
-       "      <td>70</td>\n",
-       "      <td>2404</td>\n",
-       "      <td>3705</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1000207</th>\n",
-       "      <td>25</td>\n",
-       "      <td>5835</td>\n",
-       "      <td>3705</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1000208</th>\n",
-       "      <td>380</td>\n",
-       "      <td>979</td>\n",
-       "      <td>3705</td>\n",
-       "    </tr>\n",
-       "  </tbody>\n",
-       "</table>\n",
-       "<p>1000209 rows × 3 columns</p>\n",
-       "</div>"
-      ],
-      "text/plain": [
-       "         timestamp  user_id  item_id\n",
-       "0               12        0        0\n",
-       "1               68        1        0\n",
-       "2               67        2        0\n",
-       "3               12        3        0\n",
-       "4              140        4        0\n",
-       "...            ...      ...      ...\n",
-       "1000204         14     4555     3705\n",
-       "1000205         90     2813     3705\n",
-       "1000206         70     2404     3705\n",
-       "1000207         25     5835     3705\n",
-       "1000208        380      979     3705\n",
-       "\n",
-       "[1000209 rows x 3 columns]"
-      ]
-     },
-     "execution_count": 4,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "from replay.preprocessing.label_encoder import LabelEncoder, LabelEncodingRule\n",
-    "\n",
-    "encoder = LabelEncoder(\n",
-    "    [\n",
-    "        LabelEncodingRule(\"user_id\", default_value=\"last\"),\n",
-    "        LabelEncodingRule(\"item_id\", default_value=\"last\"),\n",
-    "    ]\n",
-    ")\n",
-    "interactions = interactions.sort_values(by=\"item_id\", ascending=True)\n",
-    "encoded_interactions = encoder.fit_transform(interactions)\n",
-    "encoded_interactions"
-   ]
-  },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### Split interactions into the train, validation and test datasets using LastNSplitter\n",
-    "We use widespread splitting strategy Last-One-Out. We filter out cold items and users for simplicity."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 5,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from replay.splitters import LastNSplitter\n",
-    "\n",
-    "splitter = LastNSplitter(\n",
-    "    N=1,\n",
-    "    divide_column=\"user_id\",\n",
-    "    query_column=\"user_id\",\n",
-    "    strategy=\"interactions\",\n",
-    "    drop_cold_users=True,\n",
-    "    drop_cold_items=True\n",
-    ")\n",
-    "\n",
-    "test_events, test_gt = splitter.split(encoded_interactions)\n",
-    "validation_events, validation_gt = splitter.split(test_events)\n",
-    "train_events = validation_events"
-   ]
-  },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### Dataset preprocessing (\"baking\")\n",
-    "SasRec expects each user in the batch to provide their events in form of a sequence. For this reason, the event splits must be properly processed using the `groupby_sequences` function provided by RePlay."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 6,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from replay.data.nn.utils import groupby_sequences\n",
-    "\n",
-    "\n",
-    "def bake_data(full_data):\n",
-    "    grouped_interactions = groupby_sequences(events=full_data, groupby_col=\"user_id\", sort_col=\"timestamp\")\n",
-    "    return grouped_interactions"
-   ]
-  },
-  {
-   "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>timestamp</th>\n",
-       "      <th>item_id</th>\n",
-       "    </tr>\n",
-       "  </thead>\n",
-       "  <tbody>\n",
-       "    <tr>\n",
-       "      <th>0</th>\n",
-       "      <td>0</td>\n",
-       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
-       "      <td>[2426, 822, 2733, 2587, 2937, 3618, 2943, 708,...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1</th>\n",
-       "      <td>1</td>\n",
-       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
-       "      <td>[3272, 3026, 2760, 851, 346, 3393, 1107, 515, ...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>2</th>\n",
-       "      <td>2</td>\n",
-       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
-       "      <td>[579, 1140, 1154, 2426, 1524, 1260, 2160, 2621...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>3</th>\n",
-       "      <td>3</td>\n",
-       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
-       "      <td>[1781, 2940, 2468, 890, 948, 106, 593, 309, 49...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>4</th>\n",
-       "      <td>4</td>\n",
-       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
-       "      <td>[1108, 2229, 21, 2435, 2142, 106, 1167, 593, 1...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>...</th>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6035</th>\n",
-       "      <td>6035</td>\n",
-       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
-       "      <td>[2426, 1279, 3151, 3321, 1178, 3301, 2501, 278...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6036</th>\n",
-       "      <td>6036</td>\n",
-       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
-       "      <td>[1592, 2302, 1633, 1813, 2879, 1482, 2651, 250...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6037</th>\n",
-       "      <td>6037</td>\n",
-       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
-       "      <td>[1971, 3500, 2077, 1666, 1399, 2651, 2748, 283...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6038</th>\n",
-       "      <td>6038</td>\n",
-       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
-       "      <td>[1486, 1485, 3384, 3512, 3302, 3126, 3650, 330...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6039</th>\n",
-       "      <td>6039</td>\n",
-       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
-       "      <td>[2432, 2960, 1848, 2114, 2142, 3091, 3248, 317...</td>\n",
-       "    </tr>\n",
-       "  </tbody>\n",
-       "</table>\n",
-       "<p>6040 rows × 3 columns</p>\n",
-       "</div>"
-      ],
-      "text/plain": [
-       "      user_id                                          timestamp  \\\n",
-       "0           0  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
-       "1           1  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
-       "2           2  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
-       "3           3  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
-       "4           4  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
-       "...       ...                                                ...   \n",
-       "6035     6035  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
-       "6036     6036  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
-       "6037     6037  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
-       "6038     6038  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
-       "6039     6039  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
-       "\n",
-       "                                                item_id  \n",
-       "0     [2426, 822, 2733, 2587, 2937, 3618, 2943, 708,...  \n",
-       "1     [3272, 3026, 2760, 851, 346, 3393, 1107, 515, ...  \n",
-       "2     [579, 1140, 1154, 2426, 1524, 1260, 2160, 2621...  \n",
-       "3     [1781, 2940, 2468, 890, 948, 106, 593, 309, 49...  \n",
-       "4     [1108, 2229, 21, 2435, 2142, 106, 1167, 593, 1...  \n",
-       "...                                                 ...  \n",
-       "6035  [2426, 1279, 3151, 3321, 1178, 3301, 2501, 278...  \n",
-       "6036  [1592, 2302, 1633, 1813, 2879, 1482, 2651, 250...  \n",
-       "6037  [1971, 3500, 2077, 1666, 1399, 2651, 2748, 283...  \n",
-       "6038  [1486, 1485, 3384, 3512, 3302, 3126, 3650, 330...  \n",
-       "6039  [2432, 2960, 1848, 2114, 2142, 3091, 3248, 317...  \n",
-       "\n",
-       "[6040 rows x 3 columns]"
-      ]
-     },
-     "execution_count": 7,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "train_events = bake_data(train_events)\n",
-    "\n",
-    "validation_events = bake_data(validation_events)\n",
-    "validation_gt = bake_data(validation_gt)\n",
-    "\n",
-    "test_events = bake_data(test_events)\n",
-    "test_gt = bake_data(test_gt)\n",
-    "\n",
-    "train_events"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "To ensure we don't have unknown users in ground truth, we join validation events and validation ground truth (also join test events and test ground truth correspondingly) by user ids to leave only the common ones.  "
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 8,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def add_gt_to_events(events_df, gt_df):\n",
-    "    gt_to_join = gt_df[[\"user_id\", \"item_id\"]].rename(columns={\"item_id\": \"ground_truth\"})\n",
-    "\n",
-    "    events_df = events_df.merge(gt_to_join, on=\"user_id\", how=\"inner\")\n",
-    "    return events_df\n",
-    "\n",
-    "validation_events = add_gt_to_events(validation_events, validation_gt)\n",
-    "test_events = add_gt_to_events(test_events, test_gt)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 9,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from pathlib import Path\n",
-    "\n",
-    "data_dir = Path(\"temp/data/\")\n",
-    "data_dir.mkdir(parents=True, exist_ok=True)\n",
-    "\n",
-    "TRAIN_PATH = data_dir / \"train.parquet\"\n",
-    "VAL_PATH = data_dir / \"val.parquet\"\n",
-    "PREDICT_PATH = data_dir / \"test.parquet\"\n",
-    "\n",
-    "ENCODER_PATH = data_dir / \"encoder\""
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 10,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "train_events.to_parquet(TRAIN_PATH)\n",
-    "validation_events.to_parquet(VAL_PATH)\n",
-    "test_events.to_parquet(PREDICT_PATH)\n",
-    "\n",
-    "encoder.save(ENCODER_PATH)"
-   ]
-  },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Prepare to model training\n",
-    "### 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. It is required by the SasRec model to correctly create embeddings for every source column. Note that user_id does not required in `TensorSchema`.\n",
-    "\n",
-    "Note that the **padding value** is the next value (item_id) after the last one. **Cardinality** is the number of unique values ​​given the padding value."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 11,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from replay.data import FeatureHint, FeatureType\n",
-    "from replay.data.nn import TensorFeatureInfo, TensorSchema\n",
-    "\n",
-    "\n",
-    "EMBEDDING_DIM = 64\n",
-    "\n",
-    "encoder = encoder.load(ENCODER_PATH)\n",
-    "NUM_UNIQUE_ITEMS = len(encoder.mapping[\"item_id\"])\n",
-    "\n",
-    "tensor_schema = TensorSchema(\n",
-    "    [\n",
-    "        TensorFeatureInfo(\n",
-    "            name=\"item_id\",\n",
-    "            is_seq=True,\n",
-    "            padding_value=NUM_UNIQUE_ITEMS,\n",
-    "            cardinality=NUM_UNIQUE_ITEMS + 1,  # taking into account padding\n",
-    "            embedding_dim=EMBEDDING_DIM,\n",
-    "            feature_type=FeatureType.CATEGORICAL,\n",
-    "            feature_hint=FeatureHint.ITEM_ID,\n",
-    "        )\n",
-    "    ]\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### Configure ParquetModule and transformation pipelines\n",
-    "\n",
-    "The `ParquetModule` class enables training of models on large datasets by reading data in batch-wise way. This class initialized with **paths to every data split, a metadata dict containing information about shape and padding value of every column and a dict of transforms**. `ParquetModule`'s  \"transform pipelines\" are stage-specific modules implementing additional preprocessing to be performed on batch level right before the forward pass.  \n",
-    "\n",
-    "For SasRec model, RePlay provides a function that generates a sequence of appropriate transforms for each data split named **make_default_sasrec_transforms**.\n",
-    "\n",
-    "Internally this function creates the following transforms:\n",
-    "1) Training:\n",
-    "    1. Create a target, which contains the shifted item sequence that represents the next item in the sequence (for the next item prediction task).\n",
-    "    2. Rename features to match it with expected format by the model during training.\n",
-    "    3. Unsqueeze target (*positive_labels*) and it's padding mask (*target_padding_mask*) for getting required shape of this tensors for loss computation.\n",
-    "    4. Group input features to be embed in expected format.\n",
-    "\n",
-    "2) Validation/Inference:\n",
-    "    1. Rename/group features to match it with expected format by the model during valdiation/inference.\n",
-    "\n",
-    "If a different set of transforms is required, you can create them yourself and submit them to the ParquetModule in the form of a dictionary where the key is the name of the split, and the value is the list of transforms. Available transforms are in the replay/nn/transforms/.\n",
-    "\n",
-    "**Note:** One of the transforms for the training data prepares the initial sequence for the task of Next Item Prediction so it shifts the sequence of items. For the final sequence length to be correct, you need to set shape of item_id in metadata as **model sequence length + shift**. Default shift value is 1."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 12,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import copy\n",
-    "\n",
-    "import torch\n",
-    "\n",
-    "from replay.data.nn import TensorSchema\n",
-    "from replay.nn.transform import GroupTransform, NextTokenTransform, RenameTransform, UnsqueezeTransform, FrequencyNegativeSamplingTransform"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 13,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def make_sasrec_transforms(\n",
-    "    tensor_schema: TensorSchema, query_column: str = \"query_id\", num_negative_samples: int = 128,\n",
-    ") -> dict[str, list[torch.nn.Module]]:\n",
-    "    item_column = tensor_schema.item_id_feature_name\n",
-    "    vocab_size = tensor_schema[item_column].cardinality\n",
-    "    train_transforms = [\n",
-    "        FrequencyNegativeSamplingTransform(vocab_size, num_negative_samples),\n",
-    "        NextTokenTransform(label_field=item_column, query_features=query_column, shift=1),\n",
-    "        RenameTransform(\n",
-    "            {\n",
-    "                query_column: \"query_id\",\n",
-    "                f\"{item_column}_mask\": \"padding_mask\",\n",
-    "                \"positive_labels_mask\": \"target_padding_mask\",\n",
-    "            }\n",
-    "        ),\n",
-    "        UnsqueezeTransform(\"target_padding_mask\", -1),\n",
-    "        UnsqueezeTransform(\"positive_labels\", -1),\n",
-    "        GroupTransform({\"feature_tensors\": [item_column]}),\n",
-    "    ]\n",
-    "\n",
-    "    val_transforms = [\n",
-    "        RenameTransform({query_column: \"query_id\", f\"{item_column}_mask\": \"padding_mask\"}),\n",
-    "        GroupTransform({\"feature_tensors\": [item_column]}),\n",
-    "    ]\n",
-    "    test_transforms = copy.deepcopy(val_transforms)\n",
-    "\n",
-    "    predict_transforms = copy.deepcopy(val_transforms)\n",
-    "\n",
-    "    transforms = {\n",
-    "        \"train\": train_transforms,\n",
-    "        \"validate\": val_transforms,\n",
-    "        \"test\": test_transforms,\n",
-    "        \"predict\": predict_transforms,\n",
-    "    }\n",
-    "\n",
-    "    return transforms\n",
-    "\n",
-    "transforms = make_sasrec_transforms(tensor_schema, query_column=\"user_id\")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 14,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "MAX_SEQ_LEN = 50\n",
-    "\n",
-    "def create_meta(shape: int, gt_shape: Optional[int] = None):\n",
-    "    meta = {\n",
-    "        \"user_id\": {},\n",
-    "        \"item_id\": {\"shape\": shape, \"padding\": tensor_schema[\"item_id\"].padding_value},\n",
-    "    }\n",
-    "    if gt_shape is not None:\n",
-    "        meta.update({\"ground_truth\": {\"shape\": gt_shape, \"padding\": -1}})\n",
-    "\n",
-    "    return meta\n",
-    "\n",
-    "train_metadata = {\n",
-    "    \"train\": create_meta(shape=MAX_SEQ_LEN+1),\n",
-    "    \"validate\": create_meta(shape=MAX_SEQ_LEN, gt_shape=1),\n",
-    "}"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 15,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from replay.data.nn import ParquetModule\n",
-    "\n",
-    "BATCH_SIZE = 32\n",
-    "\n",
-    "parquet_module = ParquetModule(\n",
-    "    train_path=TRAIN_PATH,\n",
-    "    validate_path=VAL_PATH,\n",
-    "    batch_size=BATCH_SIZE,\n",
-    "    metadata=train_metadata,\n",
-    "    transforms=transforms,\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Train model\n",
-    "### Create SasRec model instance and run the training stage using lightning\n",
-    "We may now train the model using the Lightning trainer class. \n",
-    "\n",
-    "RePlay's implementation of SasRec is designed in a modular, **block-based approach**. Instead of passing configuration parameters to the constructor, SasRec is now built by providing fully initialized components that makes the model more flexible and easier to extend.\n",
-    "\n",
-    "#### Default Configuration\n",
-    "\n",
-    "Default SasRec model may be created quickly via method **from_params**. Default model instance has CE loss, original SasRec transformer layes, and embeddings are aggregated via sum."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 16,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from replay.nn.sequential import SasRec\n",
-    "from typing import Literal\n",
-    "def make_sasrec(\n",
-    "    schema: TensorSchema,\n",
-    "    embedding_dim: int = 192,\n",
-    "    num_heads: int = 4,\n",
-    "    num_blocks: int = 2,\n",
-    "    max_sequence_length: int = 50,\n",
-    "    dropout: float = 0.3,\n",
-    "    excluded_features: Optional[list[str]] = None,\n",
-    "    categorical_list_feature_aggregation_method: Literal[\"sum\", \"mean\", \"max\"] = \"sum\",\n",
-    ") -> SasRec:\n",
-    "    from replay.nn.sequential.sasrec import SasRecBody, SasRecTransformerLayer\n",
-    "    from replay.nn.agg import SumAggregator\n",
-    "    from replay.nn.embedding import SequenceEmbedding\n",
-    "    from replay.nn.loss import CE, CESampled\n",
-    "    from replay.nn.mask import DefaultAttentionMask\n",
-    "    from replay.nn.sequential.sasrec.agg import PositionAwareAggregator\n",
-    "    from replay.nn.sequential.sasrec.transformer import SasRecTransformerLayer\n",
-    "    excluded_features = [\n",
-    "        schema.query_id_feature_name,\n",
-    "        schema.timestamp_feature_name,\n",
-    "        *(excluded_features or []),\n",
-    "    ]\n",
-    "    excluded_features = list(set(excluded_features))\n",
-    "    body = SasRecBody(\n",
-    "        embedder=SequenceEmbedding(\n",
-    "            schema=schema,\n",
-    "            categorical_list_feature_aggregation_method=categorical_list_feature_aggregation_method,\n",
-    "            excluded_features=excluded_features,\n",
-    "        ),\n",
-    "        embedding_aggregator=PositionAwareAggregator(\n",
-    "            embedding_aggregator=SumAggregator(embedding_dim=embedding_dim),\n",
-    "            max_sequence_length=max_sequence_length,\n",
-    "            dropout=dropout,\n",
-    "        ),\n",
-    "        attn_mask_builder=DefaultAttentionMask(\n",
-    "            reference_feature_name=schema.item_id_feature_name,\n",
-    "            num_heads=num_heads,\n",
-    "        ),\n",
-    "        encoder=SasRecTransformerLayer(\n",
-    "            embedding_dim=embedding_dim,\n",
-    "            num_heads=num_heads,\n",
-    "            num_blocks=num_blocks,\n",
-    "            dropout=dropout,\n",
-    "            activation=\"relu\",\n",
-    "        ),\n",
-    "        output_normalization=torch.nn.LayerNorm(embedding_dim),\n",
-    "    )\n",
-    "    padding_idx = schema.item_id_features.item().padding_value\n",
-    "    return SasRec(\n",
-    "        body=body,\n",
-    "        loss=CESampled(padding_idx=padding_idx),\n",
-    "    )"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 17,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "NUM_BLOCKS = 2\n",
-    "NUM_HEADS = 2\n",
-    "DROPOUT = 0.3\n",
-    "\n",
-    "sasrec = make_sasrec(\n",
-    "    schema=tensor_schema,\n",
-    "    embedding_dim=EMBEDDING_DIM,\n",
-    "    max_sequence_length=MAX_SEQ_LEN,\n",
-    "    num_heads=NUM_HEADS,\n",
-    "    num_blocks=NUM_BLOCKS,\n",
-    "    dropout=DROPOUT,\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "A universal PyTorch Lightning module is provided. It can work with any NN model."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 18,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from replay.nn.lightning.optimizer import OptimizerFactory\n",
-    "from replay.nn.lightning.scheduler import LRSchedulerFactory\n",
-    "from replay.nn.lightning import LightningModule\n",
-    "\n",
-    "model = LightningModule(\n",
-    "    sasrec,\n",
-    "    optimizer_factory=OptimizerFactory(),\n",
-    "    lr_scheduler_factory=LRSchedulerFactory(),\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "To facilitate training, we add the following callbacks:\n",
-    "1) `ModelCheckpoint` - to save the best trained model based on its Recall metric. It's a default Lightning Callback.\n",
-    "1) `ComputeMetricsCallback` - to display a detailed validation metric matrix after each epoch. It's a custom RePlay callback for computing recsys metrics on validation and test stages.\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 19,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "GPU available: True (cuda), used: True\n",
-      "TPU available: False, using: 0 TPU cores\n",
-      "You are using a CUDA device ('NVIDIA GeForce RTX 3060 Laptop GPU') 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 | Mode  | FLOPs\n",
-      "-------------------------------------------------\n",
-      "0 | model | SasRec | 291 K  | train | 0    \n",
-      "-------------------------------------------------\n",
-      "291 K     Trainable params\n",
-      "0         Non-trainable params\n",
-      "291 K     Total params\n",
-      "1.164     Total estimated model params size (MB)\n",
-      "39        Modules in train mode\n",
-      "0         Modules in eval mode\n",
-      "0         Total Flops\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "45904c0e26294a109f736787e2831727",
-       "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": "16229579b0fa4c3eadc8eeea029b139c",
-       "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": "b95d0d9187cd4511ba6122995385977c",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 0, global step 189: 'recall@10' reached 0.01921 (best 0.01921), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=0-step=189-v3.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.000662  0.004889  0.005938  0.003535\n",
-      "ndcg    0.000662  0.008170  0.012093  0.004836\n",
-      "recall  0.000662  0.019208  0.034940  0.008776\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f396e2c9f6d1415181d9cf9ad119e483",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 1, global step 378: 'recall@10' reached 0.02335 (best 0.02335), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=1-step=378-v2.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k             1        10        20         5\n",
-      "map     0.00149  0.006367  0.007697  0.005037\n",
-      "ndcg    0.00149  0.010302  0.015430  0.007117\n",
-      "recall  0.00149  0.023348  0.044213  0.013578\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "5efd6eba4c714bc7a47bab506d96d083",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 2, global step 567: 'recall@10' reached 0.03809 (best 0.03809), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=2-step=567.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.004968  0.013050  0.015700  0.011130\n",
-      "ndcg    0.004968  0.018854  0.028575  0.014137\n",
-      "recall  0.004968  0.038086  0.076668  0.023348\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "14b9ffb84489441fbb65ecd2bffab125",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 3, global step 756: 'recall@10' reached 0.05448 (best 0.05448), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=3-step=756-v3.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.004968  0.015252  0.018691  0.011594\n",
-      "ndcg    0.004968  0.024178  0.036971  0.015113\n",
-      "recall  0.004968  0.054479  0.105647  0.025998\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "4e6aea05fd6e4fd796fdcc9c8432ffe9",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 4, global step 945: 'recall@10' reached 0.08130 (best 0.08130), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=4-step=945-v4.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.008942  0.025228  0.029602  0.020003\n",
-      "ndcg    0.008942  0.038097  0.054455  0.025390\n",
-      "recall  0.008942  0.081305  0.146879  0.041894\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "d79a151cb404477283713ea09c68f907",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 5, global step 1134: 'recall@10' reached 0.09571 (best 0.09571), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=5-step=1134-v1.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.010101  0.029038  0.033716  0.022943\n",
-      "ndcg    0.010101  0.044299  0.061645  0.029304\n",
-      "recall  0.010101  0.095711  0.164928  0.048849\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "ab01d88752a54a229fa429899a4e2552",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 6, global step 1323: 'recall@10' reached 0.10714 (best 0.10714), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=6-step=1323-v2.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.012254  0.033783  0.039462  0.027742\n",
-      "ndcg    0.012254  0.050622  0.071662  0.035493\n",
-      "recall  0.012254  0.107137  0.191091  0.059281\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "14a8ee885fe8477fa24673a3227e7e9c",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 7, global step 1512: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.010929  0.031616  0.037486  0.025432\n",
-      "ndcg    0.010929  0.048030  0.069781  0.032787\n",
-      "recall  0.010929  0.103163  0.189932  0.055473\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "6b9e020670534db1879e97dd7a566422",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 8, global step 1701: 'recall@10' reached 0.12353 (best 0.12353), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=8-step=1701-v1.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.012916  0.038435  0.044025  0.031021\n",
-      "ndcg    0.012916  0.058060  0.078630  0.040108\n",
-      "recall  0.012916  0.123530  0.205332  0.068058\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "bed94580dcbd4656951d1efda7ba8a48",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 9, global step 1890: 'recall@10' reached 0.12899 (best 0.12899), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=9-step=1890-v2.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.013082  0.040335  0.046561  0.032693\n",
-      "ndcg    0.013082  0.060782  0.083663  0.042103\n",
-      "recall  0.013082  0.128995  0.219904  0.070873\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "e05eafc1c8854f8787b5e39e2c47dd76",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 10, global step 2079: 'recall@10' reached 0.12966 (best 0.12966), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=10-step=2079.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.014738  0.041342  0.047593  0.033750\n",
-      "ndcg    0.014738  0.061675  0.084694  0.043076\n",
-      "recall  0.014738  0.129657  0.221229  0.071701\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f9eee10fd1af441db81b11a8926c9d14",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 11, global step 2268: 'recall@10' reached 0.13396 (best 0.13396), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=11-step=2268.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.014075  0.040714  0.047416  0.032690\n",
-      "ndcg    0.014075  0.062142  0.086667  0.042514\n",
-      "recall  0.014075  0.133963  0.231164  0.072860\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "189ec86d37e0416583428616bbce96fc",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 12, global step 2457: 'recall@10' reached 0.14340 (best 0.14340), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=12-step=2457-v1.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016559  0.043821  0.050222  0.035111\n",
-      "ndcg    0.016559  0.066609  0.090205  0.045130\n",
-      "recall  0.016559  0.143401  0.237291  0.076172\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "c7b5622a9ed74116bfa6d8f35cba8025",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 13, global step 2646: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.014406  0.043623  0.050502  0.035351\n",
-      "ndcg    0.014406  0.066028  0.091376  0.045681\n",
-      "recall  0.014406  0.140917  0.241762  0.077331\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "72432093eb964a6183fc1dbd640a5450",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 14, global step 2835: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.015565  0.044625  0.051633  0.036168\n",
-      "ndcg    0.015565  0.067321  0.093215  0.046660\n",
-      "recall  0.015565  0.143236  0.246398  0.078987\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "4b1a70c1dcd4427da45471a9610505f7",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 15, global step 3024: 'recall@10' reached 0.15367 (best 0.15367), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=15-step=3024.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.015069  0.046444  0.053260  0.037230\n",
-      "ndcg    0.015069  0.071060  0.096111  0.048327\n",
-      "recall  0.015069  0.153668  0.253188  0.082464\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "dc15f2eb71fe47a28ca54c0bd9cf42f3",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 16, global step 3213: 'recall@10' reached 0.15946 (best 0.15946), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=16-step=3213.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.018546  0.050829  0.057741  0.041376\n",
-      "ndcg    0.018546  0.075784  0.101199  0.052429\n",
-      "recall  0.018546  0.159463  0.260474  0.086273\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "134594f99f4a4711a91047c3b4c1c1c0",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 17, global step 3402: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.018546  0.050609  0.057747  0.041723\n",
-      "ndcg    0.018546  0.074945  0.101494  0.053166\n",
-      "recall  0.018546  0.156152  0.262295  0.088260\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "952d8b7c4ecc4b8f845681f333e60e95",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 18, global step 3591: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.015565  0.048216  0.055341  0.039681\n",
-      "ndcg    0.015565  0.072890  0.099398  0.052029\n",
-      "recall  0.015565  0.154993  0.260970  0.090081\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "2e47543bfb714dcfb9ab93773989c6e3",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 19, global step 3780: 'recall@10' reached 0.16162 (best 0.16162), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=19-step=3780.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017553  0.050086  0.057070  0.040307\n",
-      "ndcg    0.017553  0.075715  0.101411  0.051752\n",
-      "recall  0.017553  0.161616  0.263785  0.086935\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "0186cb8fe6314be29e5c27ff61a4efb4",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 20, global step 3969: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016559  0.048569  0.055774  0.038668\n",
-      "ndcg    0.016559  0.074165  0.100638  0.049831\n",
-      "recall  0.016559  0.160126  0.265276  0.084120\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "ff8cbda5a5ad4131bc24bf0e21701493",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 21, global step 4158: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k             1        10        20         5\n",
-      "map     0.01689  0.050299  0.057754  0.041229\n",
-      "ndcg    0.01689  0.075501  0.103081  0.053162\n",
-      "recall  0.01689  0.159629  0.269581  0.089750\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "c7efc2e3954042dd8c399e271650103b",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 22, global step 4347: 'recall@10' reached 0.16360 (best 0.16360), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=22-step=4347-v1.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017718  0.051337  0.058542  0.041903\n",
-      "ndcg    0.017718  0.077220  0.103802  0.054198\n",
-      "recall  0.017718  0.163603  0.269415  0.092068\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "8fdf5663208e4b7189cb87eb315dd714",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 23, global step 4536: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017553  0.049966  0.057488  0.040051\n",
-      "ndcg    0.017553  0.075609  0.103244  0.051355\n",
-      "recall  0.017553  0.161616  0.271403  0.086107\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "7b221483a15448a186b77b3329be6ff6",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 24, global step 4725: 'recall@10' reached 0.16956 (best 0.16956), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=24-step=4725.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016559  0.050711  0.057813  0.040766\n",
-      "ndcg    0.016559  0.078002  0.104261  0.053405\n",
-      "recall  0.016559  0.169565  0.274218  0.092399\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "54f8fc2699aa419fa209a5ba25852555",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 25, global step 4914: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.018381  0.050549  0.058444  0.041033\n",
-      "ndcg    0.018381  0.075701  0.104654  0.052507\n",
-      "recall  0.018381  0.159795  0.274714  0.087763\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "399879ba18f7499d846c1f937a714ae6",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 26, global step 5103: 'recall@10' reached 0.17122 (best 0.17122), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=26-step=5103.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.022686  0.057648  0.064919  0.048236\n",
-      "ndcg    0.022686  0.083871  0.110727  0.060845\n",
-      "recall  0.022686  0.171220  0.278192  0.099520\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "7fd4a45308704782b46d80a418bbe30e",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 27, global step 5292: 'recall@10' reached 0.17188 (best 0.17188), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=27-step=5292.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.018049  0.053604  0.060891  0.043368\n",
-      "ndcg    0.018049  0.080857  0.107818  0.055870\n",
-      "recall  0.018049  0.171883  0.279351  0.094221\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "fa5f1113fac746b0ac41368de2eca760",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 28, global step 5481: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016228  0.052327  0.059730  0.042565\n",
-      "ndcg    0.016228  0.079886  0.107125  0.055913\n",
-      "recall  0.016228  0.171883  0.280179  0.097036\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "a34e71864baf44dab20ed841ae6e720e",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 29, global step 5670: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.015731  0.050238  0.057931  0.040238\n",
-      "ndcg    0.015731  0.077079  0.105369  0.052521\n",
-      "recall  0.015731  0.166915  0.279351  0.090247\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f9cb3a2036894111b2ea320f1b666937",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 30, global step 5859: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016559  0.051478  0.059405  0.041271\n",
-      "ndcg    0.016559  0.078418  0.107606  0.053619\n",
-      "recall  0.016559  0.168405  0.284484  0.091571\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "dd15b49180b14ee9bdfc03e1d9f1440c",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 31, global step 6048: 'recall@10' reached 0.17519 (best 0.17519), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=31-step=6048.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016062  0.051997  0.059447  0.041563\n",
-      "ndcg    0.016062  0.080327  0.107738  0.054703\n",
-      "recall  0.016062  0.175195  0.284153  0.095214\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "2087d0ec66e245d68755fe0d4b38be3a",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 32, global step 6237: 'recall@10' reached 0.17834 (best 0.17834), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=32-step=6237.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016725  0.053134  0.060705  0.042214\n",
-      "ndcg    0.016725  0.081918  0.109782  0.055210\n",
-      "recall  0.016725  0.178341  0.289121  0.095214\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "fd91e9acc7094ed1b7baa2f1c257a447",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 33, global step 6426: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.019208  0.053199  0.061121  0.042943\n",
-      "ndcg    0.019208  0.080245  0.109425  0.055177\n",
-      "recall  0.019208  0.170889  0.286968  0.092896\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "b41a72bd49c541c88a0068c1eb3c8940",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 34, global step 6615: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016228  0.052733  0.060214  0.042899\n",
-      "ndcg    0.016228  0.080353  0.108022  0.056273\n",
-      "recall  0.016228  0.172380  0.282663  0.097367\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "7574b671f1144c35ad72d2bb149072a4",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 35, global step 6804: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k             1        10        20         5\n",
-      "map     0.01954  0.055007  0.062359  0.044064\n",
-      "ndcg    0.01954  0.083274  0.110562  0.056605\n",
-      "recall  0.01954  0.178010  0.286968  0.095214\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "b47e4be5150e4f8bacaf33c8b9f16b16",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 36, global step 6993: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.020036  0.056306  0.064246  0.046409\n",
-      "ndcg    0.020036  0.084078  0.113253  0.059811\n",
-      "recall  0.020036  0.176685  0.292598  0.101010\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "3dcb7bbf088b490ebd44a7421f69b641",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 37, global step 7182: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.019871  0.057126  0.065215  0.047538\n",
-      "ndcg    0.019871  0.084864  0.114584  0.061411\n",
-      "recall  0.019871  0.177016  0.295082  0.103991\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "96902e23ab3d49b1beca203004cbab7d",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 38, global step 7371: 'recall@10' reached 0.18132 (best 0.18132), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=38-step=7371.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k            1        10        20         5\n",
-      "map     0.0154  0.054742  0.062340  0.044704\n",
-      "ndcg    0.0154  0.084020  0.111946  0.059455\n",
-      "recall  0.0154  0.181321  0.292267  0.104819\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "fd9bd09986c240529c23d02d6dca3201",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 39, global step 7560: 'recall@10' reached 0.18231 (best 0.18231), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=39-step=7560.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017387  0.054622  0.062271  0.043890\n",
-      "ndcg    0.017387  0.083950  0.112185  0.057401\n",
-      "recall  0.017387  0.182315  0.294751  0.099023\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "0aba8305e71f4f1c957ec35d79fc0acb",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 40, global step 7749: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k            1        10        20         5\n",
-      "map     0.0154  0.052420  0.060409  0.041273\n",
-      "ndcg    0.0154  0.082044  0.111393  0.054579\n",
-      "recall  0.0154  0.181487  0.298063  0.095546\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "11610e6768a442a7aac578162ca26123",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 41, global step 7938: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017387  0.055064  0.063020  0.044130\n",
-      "ndcg    0.017387  0.084157  0.113600  0.057531\n",
-      "recall  0.017387  0.181321  0.298725  0.098692\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "afac63fbd8b54249bc6e0ebf214ae721",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 42, global step 8127: 'recall@10' reached 0.18364 (best 0.18364), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=42-step=8127.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.019374  0.056773  0.064316  0.046059\n",
-      "ndcg    0.019374  0.085958  0.113787  0.059584\n",
-      "recall  0.019374  0.183640  0.294420  0.101176\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "ef019fe931f742b3b811e73fd5404c52",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 43, global step 8316: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.018381  0.055013  0.062569  0.044560\n",
-      "ndcg    0.018381  0.083226  0.110987  0.057615\n",
-      "recall  0.018381  0.177513  0.287796  0.097698\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "6ec9a487bb04449db455a3b620156a77",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 44, global step 8505: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.019871  0.056811  0.064680  0.046062\n",
-      "ndcg    0.019871  0.085587  0.114373  0.059273\n",
-      "recall  0.019871  0.181818  0.295910  0.099851\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "5b06405c1da54f31823138198940eda6",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 45, global step 8694: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016062  0.053028  0.060992  0.042634\n",
-      "ndcg    0.016062  0.081447  0.110764  0.055989\n",
-      "recall  0.016062  0.176354  0.292929  0.097036\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "9f9ccc6e42734a7eab2c3b6097f56427",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 46, global step 8883: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.015565  0.054815  0.062818  0.044467\n",
-      "ndcg    0.015565  0.084146  0.113348  0.058769\n",
-      "recall  0.015565  0.181818  0.297400  0.102666\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "8f0b2c3bfb9b4dfab3b788b72ba05835",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 47, global step 9072: 'recall@10' reached 0.18662 (best 0.18662), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=47-step=9072.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.018877  0.058107  0.065509  0.047707\n",
-      "ndcg    0.018877  0.087807  0.115050  0.062437\n",
-      "recall  0.018877  0.186620  0.294916  0.107799\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "c5a662fbe23e433bb394c0ea3fe13f0c",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 48, global step 9261: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016725  0.053884  0.061396  0.043225\n",
-      "ndcg    0.016725  0.083373  0.110910  0.057023\n",
-      "recall  0.016725  0.182149  0.291439  0.099520\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "a05aba0cde6e40b584a4578dfe033a85",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 49, global step 9450: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.019705  0.057117  0.064926  0.046454\n",
-      "ndcg    0.019705  0.086163  0.114734  0.060090\n",
-      "recall  0.019705  0.183143  0.296407  0.102004\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "2efa59edb69541259ad8a15ce10f5b1e",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 50, global step 9639: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.022189  0.059775  0.067521  0.049652\n",
-      "ndcg    0.022189  0.088257  0.116693  0.063292\n",
-      "recall  0.022189  0.183308  0.296241  0.105150\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f85b90dc5a9b49da8739cf0df654d539",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 51, global step 9828: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.020864  0.058991  0.066927  0.049070\n",
-      "ndcg    0.020864  0.087316  0.116491  0.062868\n",
-      "recall  0.020864  0.181653  0.297566  0.105150\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "474196c5121f4c1eb1eed65c9f975492",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 52, global step 10017: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.019208  0.054836  0.063047  0.044936\n",
-      "ndcg    0.019208  0.082000  0.111981  0.057885\n",
-      "recall  0.019208  0.172545  0.291273  0.097698\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "3c85df11f2ea4dbbae2be2ff3327b1c0",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 53, global step 10206: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.015234  0.051760  0.059551  0.041094\n",
-      "ndcg    0.015234  0.080313  0.109110  0.054204\n",
-      "recall  0.015234  0.175857  0.290611  0.094552\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "6935921a319d4f9dbb3322c7bfade032",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 54, global step 10395: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016393  0.053267  0.061428  0.042606\n",
-      "ndcg    0.016393  0.082624  0.112685  0.056081\n",
-      "recall  0.016393  0.181156  0.300712  0.097533\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "a20b778f33914f6a859ac8c9af047756",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 55, global step 10584: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.018546  0.056375  0.064499  0.045653\n",
-      "ndcg    0.018546  0.085650  0.115357  0.059239\n",
-      "recall  0.018546  0.183640  0.301374  0.101010\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "74efdae045ca49fdbdac049832ffe873",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 56, global step 10773: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k             1        10        20         5\n",
-      "map     0.01689  0.054781  0.063420  0.044130\n",
-      "ndcg    0.01689  0.083637  0.115169  0.057295\n",
-      "recall  0.01689  0.180162  0.305017  0.097533\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "c9fed52a17f74d13aaab5606b8c22b21",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 57, global step 10962: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.014903  0.052826  0.061114  0.042604\n",
-      "ndcg    0.014903  0.081359  0.111826  0.056267\n",
-      "recall  0.014903  0.176519  0.297566  0.098195\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "8385fa92b9fa4fbbad147fe2276ecd05",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 58, global step 11151: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.018546  0.057551  0.065000  0.046760\n",
-      "ndcg    0.018546  0.087028  0.114631  0.060664\n",
-      "recall  0.018546  0.185296  0.295413  0.103328\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "1e41ff6e0789411d94ea038d498d9e1f",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 59, global step 11340: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016725  0.054240  0.062219  0.043672\n",
-      "ndcg    0.016725  0.083596  0.112902  0.057523\n",
-      "recall  0.016725  0.181818  0.298228  0.100182\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "19d7ff7655534dcb83a44980e6db8490",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 60, global step 11529: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.018877  0.056656  0.064225  0.045791\n",
-      "ndcg    0.018877  0.085705  0.113617  0.059086\n",
-      "recall  0.018877  0.182812  0.293923  0.099851\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "d138b9ee480745289b9fb32d37479313",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 61, global step 11718: 'recall@10' reached 0.18894 (best 0.18894), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=61-step=11718.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.020202  0.058858  0.066230  0.047933\n",
-      "ndcg    0.020202  0.088808  0.116065  0.061998\n",
-      "recall  0.020202  0.188939  0.297566  0.105315\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "ab3628938d5549658e2072ffac8c8cd3",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 62, global step 11907: 'recall@10' reached 0.19126 (best 0.19126), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=62-step=11907.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.018546  0.059156  0.066260  0.047778\n",
-      "ndcg    0.018546  0.089650  0.115582  0.061979\n",
-      "recall  0.018546  0.191257  0.293923  0.105481\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "5fb68ad004cd45dcaed411451fccdd7a",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 63, global step 12096: 'recall@10' reached 0.19424 (best 0.19424), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=63-step=12096.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017056  0.058524  0.066103  0.047138\n",
-      "ndcg    0.017056  0.089873  0.117581  0.062228\n",
-      "recall  0.017056  0.194237  0.304024  0.108627\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "71e338441644421ea8b6cc32c086b71b",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 64, global step 12285: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.018381  0.057568  0.065550  0.046550\n",
-      "ndcg    0.018381  0.087703  0.117039  0.060806\n",
-      "recall  0.018381  0.188276  0.304852  0.104653\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "0381d441937f40219cbedef04f8235bc",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 65, global step 12474: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k             1        10        20         5\n",
-      "map     0.01689  0.057109  0.065127  0.046473\n",
-      "ndcg    0.01689  0.087242  0.116738  0.061132\n",
-      "recall  0.01689  0.187614  0.304852  0.106143\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "5120fc8bdf784305890f9db8a9c6fa80",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 66, global step 12663: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.018049  0.057167  0.065304  0.046158\n",
-      "ndcg    0.018049  0.087499  0.117282  0.060454\n",
-      "recall  0.018049  0.188939  0.307004  0.104488\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "189c0b78fcc24681b130c8c53c0aa0cb",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 67, global step 12852: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017387  0.056557  0.064928  0.046536\n",
-      "ndcg    0.017387  0.085917  0.116763  0.061370\n",
-      "recall  0.017387  0.183474  0.306177  0.106971\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "e15930826aa840c4bc6e03a49be0848d",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 68, global step 13041: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.018712  0.058088  0.066686  0.047259\n",
-      "ndcg    0.018712  0.088043  0.119697  0.061554\n",
-      "recall  0.018712  0.187945  0.313794  0.105481\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "0c8d861c579a4e43ac269be8985f0af3",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 69, global step 13230: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.018049  0.058269  0.066586  0.047334\n",
-      "ndcg    0.018049  0.088878  0.119519  0.062153\n",
-      "recall  0.018049  0.190926  0.312800  0.107799\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "ea017f9572c046a5ab02de1d9196badc",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 70, global step 13419: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016725  0.057414  0.065719  0.046779\n",
-      "ndcg    0.016725  0.088837  0.119433  0.062701\n",
-      "recall  0.016725  0.193575  0.315284  0.111939\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "946138e70b474424ac7b15fabbcb1b15",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 71, global step 13608: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017884  0.056873  0.064915  0.045822\n",
-      "ndcg    0.017884  0.087012  0.116591  0.060047\n",
-      "recall  0.017884  0.187614  0.305183  0.103825\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "37b783951bcd4fe49afa8ad617c72fdb",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 72, global step 13797: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.019043  0.057263  0.065566  0.046616\n",
-      "ndcg    0.019043  0.086345  0.116833  0.060380\n",
-      "recall  0.019043  0.183308  0.304355  0.102666\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "89240546fba2486eace45af1403ffc4e",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 73, global step 13986: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.018215  0.057662  0.066092  0.047373\n",
-      "ndcg    0.018215  0.086925  0.117706  0.061920\n",
-      "recall  0.018215  0.184136  0.306011  0.106640\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "34f2d52527524ccb9e4caafdbf45a8a7",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 74, global step 14175: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.018215  0.058467  0.066697  0.047516\n",
-      "ndcg    0.018215  0.089375  0.119492  0.062387\n",
-      "recall  0.018215  0.192582  0.311972  0.108130\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "9eb272372ba74074979223fb36988b4b",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 75, global step 14364: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017387  0.057744  0.065851  0.046274\n",
-      "ndcg    0.017387  0.088852  0.118646  0.060779\n",
-      "recall  0.017387  0.192747  0.311144  0.105315\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "e7fb52f34f9849fe80f317eff347ac59",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 76, global step 14553: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.015897  0.055235  0.063557  0.044367\n",
-      "ndcg    0.015897  0.085566  0.116100  0.058943\n",
-      "recall  0.015897  0.186786  0.307998  0.103825\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "e3905ba364ce4b79bec284d9f3ba399a",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 77, global step 14742: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016725  0.057168  0.065456  0.046627\n",
-      "ndcg    0.016725  0.087777  0.118245  0.061825\n",
-      "recall  0.016725  0.189767  0.310813  0.108627\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "5d9d0a90829a47be8a39188f73fcb299",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 78, global step 14931: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016725  0.054820  0.063159  0.043650\n",
-      "ndcg    0.016725  0.084686  0.115602  0.057321\n",
-      "recall  0.016725  0.184633  0.307998  0.099354\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "d1f3a1f2e246405eaa8a9db0a95a511a",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 79, global step 15120: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k            1        10        20         5\n",
-      "map     0.0154  0.056311  0.064352  0.045432\n",
-      "ndcg    0.0154  0.087242  0.116821  0.060765\n",
-      "recall  0.0154  0.190098  0.307667  0.107965\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "9fa3134251a44d71b04280c72af0ddc5",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 80, global step 15309: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017056  0.056116  0.064667  0.045543\n",
-      "ndcg    0.017056  0.085447  0.116845  0.059621\n",
-      "recall  0.017056  0.183143  0.307832  0.102832\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "837c34003f17463cb66971c10ddce1e6",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 81, global step 15498: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017221  0.055801  0.064094  0.044936\n",
-      "ndcg    0.017221  0.085296  0.115712  0.058833\n",
-      "recall  0.017221  0.183640  0.304355  0.101507\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "bfab9a3923474fd4930e468a5144466f",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 82, global step 15687: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016062  0.054903  0.063218  0.044500\n",
-      "ndcg    0.016062  0.083626  0.114177  0.058224\n",
-      "recall  0.016062  0.179169  0.300546  0.100182\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "ce807241730e4067b08757e709957866",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 83, global step 15876: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.015897  0.057228  0.065394  0.046249\n",
-      "ndcg    0.015897  0.088010  0.117992  0.060994\n",
-      "recall  0.015897  0.190594  0.309654  0.106143\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "da634848c3d24e2e94fc916a025e7bbe",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 84, global step 16065: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016725  0.056139  0.064478  0.044845\n",
-      "ndcg    0.016725  0.086567  0.117230  0.059108\n",
-      "recall  0.016725  0.188111  0.309985  0.102997\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "1a6a134c1f7a47d7845ac0b7070b2c7b",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 85, global step 16254: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017056  0.056356  0.064636  0.046164\n",
-      "ndcg    0.017056  0.085823  0.116328  0.060727\n",
-      "recall  0.017056  0.183971  0.305349  0.105481\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "8c0c921e05674658b3bb48f8f101abc9",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 86, global step 16443: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.018877  0.057997  0.065875  0.046260\n",
-      "ndcg    0.018877  0.088555  0.117546  0.059837\n",
-      "recall  0.018877  0.190926  0.306177  0.101507\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "19135c0fba704ad0ae7838c12ed0fed5",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 87, global step 16632: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017221  0.056635  0.064463  0.046288\n",
-      "ndcg    0.017221  0.087059  0.116068  0.061620\n",
-      "recall  0.017221  0.188442  0.304189  0.108958\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "5bf0b56c43d24c5d930359679b722e04",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 88, global step 16821: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016725  0.057663  0.065748  0.046454\n",
-      "ndcg    0.016725  0.088640  0.118382  0.061238\n",
-      "recall  0.016725  0.191919  0.310151  0.106640\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "0223aa1f4bc645d9aab2100e1c4f7384",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 89, global step 17010: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016062  0.057177  0.065153  0.045805\n",
-      "ndcg    0.016062  0.088354  0.117906  0.060564\n",
-      "recall  0.016062  0.192250  0.310151  0.105812\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "28bf241cbefb422281ec0406264c576a",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 90, global step 17199: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k             1        10        20         5\n",
-      "map     0.01689  0.057756  0.066077  0.046741\n",
-      "ndcg    0.01689  0.088707  0.119294  0.061811\n",
-      "recall  0.01689  0.191754  0.313297  0.108130\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "d12545edc74045718845f0fd8bff93cd",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 91, global step 17388: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016228  0.057542  0.065642  0.046476\n",
-      "ndcg    0.016228  0.088482  0.118328  0.061535\n",
-      "recall  0.016228  0.191422  0.310151  0.107799\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "b6753a36ce234372b66a2d48f49f1035",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 92, global step 17577: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016559  0.056973  0.064812  0.045885\n",
-      "ndcg    0.016559  0.087680  0.116507  0.060459\n",
-      "recall  0.016559  0.190098  0.304686  0.105150\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "dbdee467d75c4e36ad40a23e875994b6",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 93, global step 17766: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016559  0.058180  0.066174  0.047397\n",
-      "ndcg    0.016559  0.089539  0.118949  0.063218\n",
-      "recall  0.016559  0.193741  0.310647  0.111939\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "2399141e31cf415f9ee3a4ad926ce1dd",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 94, global step 17955: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.015731  0.056793  0.065030  0.045725\n",
-      "ndcg    0.015731  0.087689  0.117979  0.060803\n",
-      "recall  0.015731  0.190429  0.310813  0.107137\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "39f19e1dc4174c0bb397e17bd243bdb6",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 95, global step 18144: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.018546  0.057964  0.066046  0.047199\n",
-      "ndcg    0.018546  0.088238  0.118000  0.061845\n",
-      "recall  0.018546  0.189270  0.307667  0.106971\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "1c38e55824ff41bbaf7520ddfa469142",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 96, global step 18333: 'recall@10' reached 0.19622 (best 0.19622), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=96-step=18333.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k             1        10        20         5\n",
-      "map     0.01954  0.060201  0.068178  0.048722\n",
-      "ndcg    0.01954  0.091547  0.121043  0.063432\n",
-      "recall  0.01954  0.196225  0.313794  0.108627\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "04808d64d1004f1aa9ce59a5e66148f4",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 97, global step 18522: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017221  0.057618  0.065656  0.046542\n",
-      "ndcg    0.017221  0.088212  0.118036  0.061226\n",
-      "recall  0.017221  0.190098  0.309157  0.106309\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "57c0e23d1c8945d4b7f02d50241b88fb",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 98, global step 18711: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.014903  0.055248  0.063857  0.043799\n",
-      "ndcg    0.014903  0.086091  0.117646  0.058227\n",
-      "recall  0.014903  0.188939  0.314125  0.102500\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "9da10aefb0d24ca7be02704f6c74a4c6",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 99, global step 18900: 'recall@10' was not in top 1\n",
-      "`Trainer.fit` stopped: `max_epochs=100` reached.\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k             1        10        20         5\n",
-      "map     0.01689  0.057811  0.065843  0.047138\n",
-      "ndcg    0.01689  0.088620  0.118022  0.062352\n",
-      "recall  0.01689  0.191257  0.307832  0.109124\n",
-      "\n"
-     ]
-    }
-   ],
-   "source": [
-    "from lightning.pytorch.callbacks import ModelCheckpoint\n",
-    "from lightning.pytorch.loggers import CSVLogger\n",
-    "\n",
-    "from replay.nn.lightning.callback import ComputeMetricsCallback\n",
-    "\n",
-    "\n",
-    "checkpoint_callback = ModelCheckpoint(\n",
-    "    dirpath=\"sasrec/checkpoints\",\n",
-    "    save_top_k=1,\n",
-    "    verbose=True,\n",
-    "    monitor=\"recall@10\",\n",
-    "    mode=\"max\",\n",
-    ")\n",
-    "\n",
-    "validation_metrics_callback = ComputeMetricsCallback(\n",
-    "    metrics=[\"map\", \"ndcg\", \"recall\"],\n",
-    "    ks=[1, 5, 10, 20],\n",
-    "    item_count=NUM_UNIQUE_ITEMS,\n",
-    ")\n",
-    "\n",
-    "csv_logger = CSVLogger(save_dir=\"sasrec/logs/train\", name=\"SasRec-example\")\n",
-    "\n",
-    "trainer = L.Trainer(\n",
-    "    max_epochs=100,\n",
-    "    callbacks=[checkpoint_callback, validation_metrics_callback],\n",
-    "    logger=csv_logger,\n",
-    ")\n",
-    "\n",
-    "trainer.fit(model, datamodule=parquet_module)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 20,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAhQAAAGnCAYAAADvzM6jAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjgsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvwVt1zgAAAAlwSFlzAAASdAAAEnQB3mYfeAAALWJJREFUeJzt3Xt8VPWd//H3hEAmZBoEIWS4La4BpYE0UDOUPBCWooiggpcUYZHCClRDpaB4iaxr0YpRQLeWUlC7hXLxAnVBagW1jwR5CCTFgtUVGyKLEEwMFCEkJAGS8/uDX2YZJpfJfJOZyczr+XjkEfI93/Od7yczZ3jnnDPn2CzLsgQAAGAgKtgTAAAAbR+BAgAAGCNQAAAAYwQKAABgjEABAACMESgAAIAxAgUAADBGoAAAAMYIFAAAwBiBAgAAGIsO9gRay6lTp7Rjxw717t1bMTExwZ4OAABtRnV1tY4ePaqRI0fqiiuu8GmdsA0UO3bs0MSJE4M9DQAA2qzNmzdrwoQJPvUN20DRu3dvSRd/GUlJSUGeTcsqLy9Xfn6+XC6XHA5HsKcTENRMzeEo0uqVqLmt1FxYWKiJEye6/y/1RdgGirrDHElJSUpOTg7ybFpWWVmZSkpKNGDAAMXHxwd7OgFBzdQcjiKtXoma21rNzTllgJMyAQCAMQIFAAAwRqAAAADGCBQAAMAYgQIAABgjUAAAAGMECgAAYMyvQPGXv/xFP/3pT5WcnKy4uDj16dNHP/rRj1RQUODRb/r06bLZbF5f1157rdeYtbW1ev7553XVVVfJbrcrJSVFr732mn9VAQCAgPLrwlbPPfecPvroI2VkZCglJUUlJSVavny5hgwZoj179mjgwIHuvjExMXr11Vc91u/UqZPXmAsXLlR2drZmzZqltLQ0bdmyRVOmTJHNZtPdd9/tzzQBAECA+BUoHnzwQW3YsEEdOnRwt02aNEmDBg1Sdna21q1b938PEB2tqVOnNjresWPHtGzZMs2ZM0fLly+XJM2cOVMjR47Uww8/rIyMDLVr186fqQIAgADw65BHenq6R5iQpH79+ik5OVkHDhzw6l9TU6OysrIGx9uyZYvOnz+vzMxMd5vNZtP999+voqIi7d69259pAgCAAGmxkzIty9I333yjrl27erSfPXtW8fHx6tSpk7p06aI5c+aovLzco8++ffsUFxenAQMGeLS7XC73cgAAELpa7OZg69ev17Fjx/TUU0+525xOpx555BENGTJEtbW12rZtm1asWKFPPvlEubm5io6++PDFxcXq3r27bDabx5hOp1OS9PXXXzf62KWlpTp+/LhHW2FhoaSLd3lrbO9IW1RRUeHxPRJQc2SItJojrV6JmtuKy//w90WLBIovvvhCc+bM0bBhw/TjH//Y3f7ss8969Lv77rvVv39/LVy4UJs2bXKfbFlZWVnvHc3sdrt7eWNWrFihRYsW1bssPz9fJSUlzaqnrcjPzw/2FAKOmiNDpNUcafVK1Bzqjhw50ux1jANFSUmJxo8fr06dOmnTpk1Nnjw5f/58PfHEE/rggw/cgSI2NlbV1dVefauqqtzLG5OZmamMjAyPtrp7ubtcLq9DKW1dRUWF8vPz5XK5FBcXF+zpBAQ1h2fNY1780OPnrnZLs66t1StfROlElU3vzR8RpJkFRiQ8x5ej5rZRc33nQzbFKFCcPn1aN998s06dOqWdO3eqR48eTa4TGxurK6+8UidPnnS3OZ1O5eTkyLIsj8MexcXFktTkuAkJCUpISKh3mcPhaHP3n/dVXFxc2NbWEGoOLyWVtnrbT1TZVFJpC9u6LxfOz3FDqDm0ORyOZq/j90mZVVVVuvXWW1VQUKA//vGP+u53v+vTemfOnNGJEyfUrVs3d1tqaqrOnj3rlYjy8vLcywEAQOjyK1DU1NRo0qRJ2r17tzZu3Khhw4Z59amqqtKZM2e82p9++mlZlqWxY8e62yZMmKD27dtrxYoV7jbLsrRy5Ur17NlT6enp/kwTAAAEiF+HPB566CG9/fbbuvXWW3Xy5EmPC1lJ0tSpU1VSUqLBgwdr8uTJ7kttb9++XX/60580duxYTZgwwd2/V69emjdvnpYsWaLz588rLS1Nmzdv1s6dO7V+/XouagUAQIjzK1Ds379fkrR161Zt3brVa/nUqVN1xRVX6JZbbtH777+vNWvWqKamRklJSVq8eLEWLFigqCjPnSPZ2dnq3LmzVq1apdWrV6tfv35at26dpkyZ4s8UAQBAAPkVKHJzc5vsc8UVV2jt2rU+jxkVFaWsrCxlZWX5MyUAABBE3L4cAAAYI1AAAABjBAoAAGCMQAEAAIwRKAAAgDECBQAAMEagAAAAxggUAADAGIECAAAYI1AAAABjfl16GwBCRd/H3ml0+eHs8QGaCRDZ2EMBAACMESgAAIAxAgUAADBGoAAAAMYIFAAAwBiBAgAAGCNQAAAAYwQKAABgjEABAACMESgAAIAxAgUAADBGoAAAAMYIFAAAwBiBAgAAGCNQAAAAYwQKAABgjEABAACMESgAAIAxAgUAADBGoAAAAMYIFAAAwBiBAgAAGCNQAAAAYwQKAABgjEABAACMESgAAIAxAgUAADBGoAAAAMYIFAAAwBiBAgAAGCNQAAAAYwQKAABgjEABAACMESgAAIAxAgUAADBGoAAAAMYIFAAAwBiBAgAAGCNQAAAAYwQKAABgjEABAACMESgAAIAxAgUAADBGoAAAAMYIFAAAwBiBAgAAGPMrUPzlL3/RT3/6UyUnJysuLk59+vTRj370IxUUFHj1PXDggMaOHSuHw6EuXbronnvu0fHjx7361dbW6vnnn9dVV10lu92ulJQUvfbaa/5MDwAABFi0Pys999xz+uijj5SRkaGUlBSVlJRo+fLlGjJkiPbs2aOBAwdKkoqKijRixAh16tRJixcvVnl5uZYuXapPP/1U+fn56tChg3vMhQsXKjs7W7NmzVJaWpq2bNmiKVOmyGaz6e67726ZagEAQKvwK1A8+OCD2rBhg0cgmDRpkgYNGqTs7GytW7dOkrR48WJVVFTo448/Vp8+fSRJLpdLN954o1avXq3Zs2dLko4dO6Zly5Zpzpw5Wr58uSRp5syZGjlypB5++GFlZGSoXbt2RoUCAIDW49chj/T0dI8wIUn9+vVTcnKyDhw44G77wx/+oFtuucUdJiTphhtuUP/+/fXmm2+627Zs2aLz588rMzPT3Waz2XT//ferqKhIu3fv9meaAAAgQFrspEzLsvTNN9+oa9euki7udSgtLdV1113n1dflcmnfvn3un/ft26e4uDgNGDDAq1/dcgAAELr8OuRRn/Xr1+vYsWN66qmnJEnFxcWSJKfT6dXX6XTq5MmTqq6uVkxMjIqLi9W9e3fZbDavfpL09ddfN/rYpaWlXid6FhYWSpLKy8tVVlbmX1EhqqKiwuN7JKDm8JQYa3n83NVueXz3Zdu9fIzLhfL2HwnP8eWouW0oLy9v9jotEii++OILzZkzR8OGDdOPf/xjSVJlZaUkKSYmxqu/3W5394mJiXF/b6xfY1asWKFFixbVuyw/P18lJSW+F9OG5OfnB3sKAUfN4SUrtf72WdfWSpJycnL8HqOOL2MEWzg/xw2h5tB25MiRZq9jHChKSko0fvx4derUSZs2bXKfPBkbGytJqq6u9lqnqqrKo09sbKxP/RqSmZmpjIwMj7bCwkJNnDhRLpfL61BKW1dRUaH8/Hy5XC7FxcUFezoBQc3hWfOYFz/0+Lmr3dKsa2v1yhdROlFl03vzRzR7jMv5MkawRMJzfDlqbhs1X3o+pK+MAsXp06d1880369SpU9q5c6d69OjhXlZ3uKLu0MeliouL1aVLF/deCafTqZycHFmW5XHYo27dS8etT0JCghISEupd5nA4FB8f37zC2oi4uLiwra0h1BwYfR97p8k+h7PHGz9OSaWt3vYTVTaVVNp8qruhMeq0hdcLr+vI0JZqdjgczV7H75Myq6qqdOutt6qgoEB//OMf9d3vftdjec+ePdWtWzft3bvXa938/Hylpqa6f05NTdXZs2e9ElFeXp57OQAACF1+7aGoqanRpEmTtHv3bm3ZskXDhg2rt9+dd96pNWvW6OjRo+rdu7ck6c9//rMKCgo0f/58d78JEyZo/vz5WrFihfs6FJZlaeXKlerZs6fS09P9mSYA+KSpPTItsTcGCHd+BYqHHnpIb7/9tm699VadPHnSfSGrOlOnTpUkPf7449q4caNGjRqln/3sZyovL9eSJUs0aNAgzZgxw92/V69emjdvnpYsWaLz588rLS1Nmzdv1s6dO7V+/XouagXAb74cvgFgzq9AsX//fknS1q1btXXrVq/ldYGid+/e2rFjhx588EE99thj6tChg8aPH69ly5Z5faojOztbnTt31qpVq7R69Wr169dP69at05QpU/yZIgAACCC/AkVubq7PfZOTk7V9+/Ym+0VFRSkrK0tZWVn+TAkAAAQRty8HAADGCBQAAMAYgQIAABgjUAAAAGMtdnMwAIhkXMsCkY49FAAAwBh7KAD4hb/IAVyKPRQAAMAYgQIAABgjUAAAAGMECgAAYIxAAQAAjBEoAACAMQIFAAAwRqAAAADGCBQAAMAYV8oEELKauhongNDBHgoAAGCMQAEAAIwRKAAAgDECBQAAMEagAAAAxggUAADAGIECAAAYI1AAAABjBAoAAGCMQAEAAIwRKAAAgDECBQAAMEagAAAAxggUAADAGIECAAAYI1AAAABjBAoAAGCMQAEAAIwRKAAAgDECBQAAMEagAAAAxqKDPQEALa/vY+80uvxw9vgAzQRApGAPBQAAMEagAAAAxggUAADAGIECAAAYI1AAAABjBAoAAGCMQAEAAIwRKAAAgDECBQAAMEagAAAAxggUAADAGIECAAAYI1AAAABjBAoAAGCMQAEAAIwRKAAAgDECBQAAMEagAAAAxqL9XbG8vFxLlixRXl6e8vPz9e233+p3v/udpk+f7tFv+vTpWrNmjdf611xzjb744guPttraWi1dulS/+c1vVFxcrP79+ysrK0uTJ0/2d5oAgqTvY+8EewoAAsjvQHHixAk99dRT6tOnj773ve8pNze3wb4xMTF69dVXPdo6derk1W/hwoXKzs7WrFmzlJaWpi1btmjKlCmy2Wy6++67/Z0qAABoZX4HCqfTqeLiYiUmJmrv3r1KS0tr+EGiozV16tRGxzt27JiWLVumOXPmaPny5ZKkmTNnauTIkXr44YeVkZGhdu3a+TtdAADQivw+hyImJkaJiYk+96+pqVFZWVmDy7ds2aLz588rMzPT3Waz2XT//ferqKhIu3fv9neqAACglfm9h6I5zp49q/j4eJ09e1adO3fW5MmT9dxzz8nhcLj77Nu3T3FxcRowYIDHui6Xy718+PDh9Y5fWlqq48ePe7QVFhZKuniuR2NBpi2qqKjw+B4JqLl5EmOtRpc3tU00tX5r6Wq3PL6HCl/eQ/z5nfO6jgxtseby8vJmr9PqgcLpdOqRRx7RkCFDVFtbq23btmnFihX65JNPlJubq+joi1MoLi5W9+7dZbPZvNaXpK+//rrBx1ixYoUWLVpU77L8/HyVlJS0UDWhJT8/P9hTCDhq9k1WauPLc3JyjNZvbbOurQ3uBC7T1O9LMvud87qODG2p5iNHjjR7nVYPFM8++6zHz3fffbf69++vhQsXatOmTe6TLSsrKxUTE+O1vt1udy9vSGZmpjIyMjzaCgsLNXHiRLlcLq+9Hm1dRUWF8vPz5XK5FBcXF+zpBAQ1N6/mMS9+2Ojy9+aPMFq/tXS1W5p1ba1e+SJKJ6psTa8QIE39viT/fue8rqk5VB04cKDZ6wTkkMfl5s+fryeeeEIffPCBO1DExsaqurraq29VVZV7eUMSEhKUkJBQ7zKHw6H4+PgWmHXoiYuLC9vaGkLNvimpbPw/46bGa2r91naiyhb0OVzKl9+/ye+c13VkaEs1X3pKgq+CcmGr2NhYXXnllTp58qS7zel0qqSkRJbleRyyuLhYktSjR4+AzhEAAPguKIHizJkzOnHihLp16+ZuS01N1dmzZ712s+Tl5bmXAwCA0NSqgaKqqkpnzpzxan/66adlWZbGjh3rbpswYYLat2+vFStWuNssy9LKlSvVs2dPpaent+ZUAQCAAaNzKJYvX65Tp065P4GxdetWFRUVSZIeeOABffvttxo8eLAmT56sa6+9VpK0fft2/elPf9LYsWM1YcIE91i9evXSvHnztGTJEp0/f15paWnavHmzdu7cqfXr13NRKwAAQphRoFi6dKm++uor989vvfWW3nrrLUnS1KlTdcUVV+iWW27R+++/rzVr1qimpkZJSUlavHixFixYoKgozx0k2dnZ6ty5s1atWqXVq1erX79+WrdunaZMmWIyTQAA0MqMAsXhw4eb7LN27Vqfx4uKilJWVpaysrIMZgUAAAItKB8bBRBc3AkUQEsLyqc8AABAeCFQAAAAYwQKAABgjEABAACMESgAAIAxAgUAADBGoAAAAMYIFAAAwBiBAgAAGCNQAAAAYwQKAABgjEABAACMESgAAIAxAgUAADBGoAAAAMYIFAAAwBiBAgAAGCNQAAAAYwQKAABgjEABAACMESgAAIAxAgUAADBGoAAAAMaigz0BAIgEfR97x6stMdZSVqo05sUPtefJWwI/KaAFsYcCAAAYI1AAAABjBAoAAGCMQAEAAIwRKAAAgDECBQAAMMbHRoEQU/fxwks/UlhSaXMvP5w9PkgzA4CGESgAoAn1XUMCgCcOeQAAAGMECgAAYIxDHgAQAnw5rML5Mwhl7KEAAADGCBQAAMAYgQIAABgjUAAAAGMECgAAYIxAAQAAjBEoAACAMQIFAAAwxoWtgDaG+0oACEXsoQAAAMYIFAAAwBiBAgAAGCNQAAAAYwQKAABgjEABAACMESgAAIAxrkMBAGHCl2uUHM4eH4CZIBKxhwIAABgjUAAAAGMECgAAYMzvQFFeXq4nn3xSY8eOVZcuXWSz2bR69ep6+x44cEBjx46Vw+FQly5ddM899+j48eNe/Wpra/X888/rqquukt1uV0pKil577TV/pwgAAALE70Bx4sQJPfXUUzpw4IC+973vNdivqKhII0aMUGFhoRYvXqwFCxbonXfe0Y033qhz58559F24cKEeffRR3XjjjfrVr36lPn36aMqUKXr99df9nSYAAAgAvz/l4XQ6VVxcrMTERO3du1dpaWn19lu8eLEqKir08ccfq0+fPpIkl8ulG2+8UatXr9bs2bMlSceOHdOyZcs0Z84cLV++XJI0c+ZMjRw5Ug8//LAyMjLUrl07f6cLAABakd97KGJiYpSYmNhkvz/84Q+65ZZb3GFCkm644Qb1799fb775prtty5YtOn/+vDIzM91tNptN999/v4qKirR7925/pwoAAFpZq56UeezYMZWWluq6667zWuZyubRv3z73z/v27VNcXJwGDBjg1a9uOQAACE2temGr4uJiSRcPj1zO6XTq5MmTqq6uVkxMjIqLi9W9e3fZbDavfpL09ddfN/g4paWlXid5FhYWSrp48mhZWZlRHaGmoqLC43skiKSaE2MtSVJXu+f3SBBpNTe33qbey+peOyZjtLZI2pbrtMWay8vLm71OqwaKyspKSRcPj1zObre7+8TExLi/N9avIStWrNCiRYvqXZafn6+SkpJmz70tyM/PD/YUAi4Sas5K9fx51rW1QZlHMEVazb7Wm5OT0+jyy187/owRKJGwLV+uLdV85MiRZq/TqoEiNjZWklRdXe21rKqqyqNPbGysT/3qk5mZqYyMDI+2wsJCTZw4US6Xy+swSltXUVGh/Px8uVwuxcXFBXs6ARFJNY958UNJF/9qnXVtrV75IkonqmxNrBUeIq3m5tb73vwRjS6ve+2YjNHaImlbrtMWaz5w4ECz12nVQFF3uKLu0MeliouL1aVLF/deCafTqZycHFmW5XHYo27dHj16NPg4CQkJSkhIqHeZw+FQfHy83zWEsri4uLCtrSGRUHNJped/LCeqbF5t4S7Sava13qZe+y0xRqBEwrZ8ubZUs8PhaPY6rXpSZs+ePdWtWzft3bvXa1l+fr5SU1PdP6empurs2bNeqSgvL8+9HAAAhKZWv9vonXfeqTVr1ujo0aPq3bu3JOnPf/6zCgoKNH/+fHe/CRMmaP78+VqxYoX7OhSWZWnlypXq2bOn0tPTW3uqQKvz5W6QANAWGQWK5cuX69SpU+5PYGzdulVFRUWSpAceeECdOnXS448/ro0bN2rUqFH62c9+pvLyci1ZskSDBg3SjBkz3GP16tVL8+bN05IlS3T+/HmlpaVp8+bN2rlzp9avX89FrQAACGFGgWLp0qX66quv3D+/9dZbeuuttyRJU6dOVadOndS7d2/t2LFDDz74oB577DF16NBB48eP17Jly7w+1ZGdna3OnTtr1apVWr16tfr166d169ZpypQpJtMEgLDAHi6EMqNAcfjwYZ/6JScna/v27U32i4qKUlZWlrKyskymBQAAAozblwMAAGMECgAAYIxAAQAAjBEoAACAMQIFAAAwRqAAAADGCBQAAMAYgQIAABgjUAAAAGMECgAAYKzV7zYKRBLutQAgUrGHAgAAGCNQAAAAYwQKAABgjEABAACMESgAAIAxAgUAADBGoAAAAMYIFAAAwBiBAgAAGCNQAAAAYwQKAABgjEABAACMESgAAIAxAgUAADBGoAAAAMYIFAAAwBiBAgAAGCNQAAAAYwQKAABgjEABAACMESgAAIAxAgUAADBGoAAAAMYIFAAAwBiBAgAAGCNQAAAAYwQKAABgjEABAACMRQd7AgCA0NH3sXea7HM4e3wAZoK2hkAB/H9NvZHyJgoADeOQBwAAMEagAAAAxggUAADAGIECAAAYI1AAAABjBAoAAGCMQAEAAIwRKAAAgDEubAUAEcSXK2EC/mAPBQAAMEagAAAAxggUAADAGIECAAAY46RMAECLauzEz8RYS1mpgZsLAoc9FAAAwBiBAgAAGGv1QJGbmyubzVbv1549ezz67tq1S8OHD1fHjh2VmJiouXPnqry8vLWnCAAADAXsHIq5c+cqLS3Noy0pKcn97/3792v06NEaMGCAXnjhBRUVFWnp0qU6ePCg3n333UBNEwAA+CFggeL666/XXXfd1eDyxx9/XJ07d1Zubq7i4+MlSX379tWsWbP03nvvacyYMYGaKgAAaKaAnkNx5swZXbhwwau9rKxM77//vqZOneoOE5I0bdo0ORwOvfnmm4GcJgAAaKaABYoZM2YoPj5edrtdo0aN0t69e93LPv30U124cEHXXXedxzodOnRQamqq9u3bF6hpAgAAP7T6IY8OHTrozjvv1Lhx49S1a1d9/vnnWrp0qa6//nrt2rVLgwcPVnFxsSTJ6XR6re90OrVz585GH6O0tFTHjx/3aCssLJQklZeXq6ysrIWqCQ0VFRUe3yNBIGpOjLUaXe7L66ipMZqjq93y+B4JIq3mtlpvU9tCY9tBXa28f4U2fz4QYbMsK+Cv5MLCQqWkpGjEiBHatm2b1q5dq2nTpikvL08ul8uj77Rp0/T222/r1KlTDY7385//XIsWLap32UsvvaQ+ffq05PQBAAhrR44c0dy5c/XZZ58pOTnZp3WCcqXMpKQkTZgwQW+99ZZqamoUGxsrSaqurvbqW1VV5V7ekMzMTGVkZHi0FRYWauLEiXK5XBowYEDLTT4EVFRUKD8/Xy6XS3FxccGeTkAEouYxL37Y6PL35o8wHqM5utotzbq2Vq98EaUTVbYWGzeURVrNbbXepraFxraDupp5/wptBw4caPY6Qbv0du/evXXu3DlVVFS4D3XUHfq4VHFxsXr06NHoWAkJCUpISKh3mcPh8DjRM5zExcWFbW0Nac2aSyobf0P35XGbGsMfJ6psrTJuKIu0mttavU1tC77UwvtXaHM4HM1eJ2hXyjx06JDsdrscDocGDhyo6OhojxM1JencuXPav3+/UlNTgzNJAADgk1YPFJefLClJn3zyid5++22NGTNGUVFR6tSpk2644QatW7dOZ86ccfdbu3atysvLvQ5nAACA0NLqhzwmTZqk2NhYpaenKyEhQZ9//rlefvlldezYUdnZ2e5+zzzzjNLT0zVy5EjNnj1bRUVFWrZsmcaMGaOxY8e29jTRxjV2d0NJOpw9PkAzAYDI1Op7KCZOnKgTJ07ohRdeUGZmpt544w3dcccd2rt3r8fJkkOGDNEHH3yg2NhYzZ8/Xy+//LLuvfdebdq0qbWnCAAADLX6Hoq5c+dq7ty5PvUdPny4Pvroo1aeEeCfpvaCAEAk4/blAADAGIECAAAYI1AAAABjBAoAAGCMQAEAAIwRKAAAgDECBQAAMBa0m4MBANomrsmC+rCHAgAAGCNQAAAAYwQKAABgjEABAACMESgAAIAxAgUAADBGoAAAAMYIFAAAwBgXtkJE4EI8QNviyzZ7OHt8AGYCXxEoAAABN+bFD1VSaQv2NNCCOOQBAACMESgAAIAxAgUAADBGoAAAAMYIFAAAwBiBAgAAGCNQAAAAYwQKAABgjAtbIeT1fewdJcZaykrlYjgAEKrYQwEAAIwRKAAAgDECBQAAMEagAAAAxggUAADAGIECAAAYI1AAAABjBAoAAGCMQAEAAIwRKAAAgDEuvY2g6/vYO8GeAgDAEIECANAmNfXHyOHs8QGaCSQOeQAAgBZAoAAAAMYIFAAAwBiBAgAAGCNQAAAAYwQKAABgjI+NAgDCEh8rDSz2UAAAAGMECgAAYIxAAQAAjBEoAACAMU7KjGC+3JSrqZOWuLEXgHDGiZ2+Yw8FAAAwRqAAAADGCBQAAMAY51C0YWNe/FAllbYGl3NsDwAQKOyhAAAAxggUAADAWEgd8qiurtZ//Md/aO3atfr222+VkpKiX/ziF7rxxhuDPTUAALz48tH5vz1+fQBmEnwhFSimT5+uTZs2ad68eerXr59Wr16tcePGKScnR8OHDw/29AAAYSRQ19EZ8+KHykpt+Ly3cDnfLWQCRX5+vl5//XUtWbJECxYskCRNmzZNAwcO1COPPKJdu3YFeYYAAKAhIRMoNm3apHbt2mn27NnuNrvdrnvvvVePP/64jh49qt69ewdxhi2nJa5Q2VKPAwAIrkD9n9DaQiZQ7Nu3T/3791d8fLxHu8vlkiTt37+/wUBRWlqq48ePe7R9/vnnkqS//e1vKi8vb7F5zvr9XuMx4pvuory8vAaXVVZW6siRI4o5HaX4auPpNCplzq8bXe5LLS0h5px05EhtQGoOFdQc7Nm0vkirV6Jmf2tu6r34lWnX+TdwAw4dOiTp4rmNPrNCRHJysvXDH/7Qq/1//ud/LEnWypUrG1z3ySeftCTxxRdffPHFF18t+LV582af/x8PmT0UlZWViomJ8Wq32+3u5Q3JzMxURkaGR1tZWZkKCgo0aNCgesdtywoLCzVx4kRt3rxZSUlJwZ5OQFAzNYejSKtXoua2UnN1dbWOHj2qkSNH+rxOyASK2NjYenetVFVVuZc3JCEhQQkJCV7tw4YNa7kJhqCkpCQlJycHexoBRc2RIdJqjrR6JWpuC4YMGdKs/iFzYSun06ni4mKv9rq2Hj16BHpKAADARyETKFJTU1VQUKCysjKP9rqTE1NTU4MwKwAA4IuQCRR33XWXampq9PLLL7vbqqur9bvf/U5Dhw4Nm4+MAgAQjkLmHIqhQ4cqIyNDWVlZKi0tVVJSktasWaPDhw/rt7/9bbCnF1K6deumJ598Ut26dQv2VAKGmiNDpNUcafVK1BzObJZlWcGeRJ2qqio98cQTWrdunfteHk8//bRuuummYE8NAAA0IqQCBQAAaJtC5hwKAADQdhEoAACAMQIFAAAwRqAAAADGCBRBlpubK5vNVu/Xnj17JEmHDx9usI/NZtOsWbOaNV6o+Otf/6rbbrtNXbp0UceOHTVw4EC99NJLHn127dql4cOHq2PHjkpMTNTcuXPrvXtsdXW1Hn30UfXo0UOxsbEaOnSo3n///UCV4pPG6j179qx+/etfa8yYMXI6nfrOd76jwYMH6ze/+Y1qamo8xmns9fD6668Ho7QGNfUc/8u//Eu9dYwdO9ZrrLbwHEuN1xxu2/L06dMbrefYsWPuvuGwLftSb7huy74ImetQRLq5c+cqLS3No63uJjLdunXT2rVrvdbZtm2b1q9frzFjxjRrvFDw3nvv6dZbb9XgwYP1xBNPyOFw6Msvv1RRUZG7z/79+zV69GgNGDBAL7zwgoqKirR06VIdPHhQ7777rsd406dP16ZNmzRv3jz169dPq1ev1rhx45STk6Phw4cHujwvTdV76NAhPfDAAxo9erQefPBBxcfHa/v27crMzNSePXu0Zs0arzEnT56scePGebSF0v1rfHmOJalXr1569tlnPdrqu9R+qD/HUtM1h9u2/JOf/EQ33HCDR5tlWbrvvvvUt29f9ezZU1L4bMu+1PvZZ5+F3bbss2beZRwtLCcnx5Jkbdy4sdnrjh492oqPj7cqKytbZLxAOX36tNW9e3fr9ttvt2pqahrsd/PNN1tOp9M6ffq0u+2VV16xJFnbt293t+Xl5VmSrCVLlrjbKisrrauvvtoaNmxY6xTRDL7Ue/z4ceuzzz7zap8xY4YlyTp48KC77X//93+96g01vj7HI0eOtJKTk5scL9SfY8vyveb6tNVtuT47d+60JFnPPPOMuy1ctuX6XF5vuG3LzcEhjxBy5swZXbhwwae+xcXFysnJ0R133OG+xbvJeIG0YcMGffPNN3rmmWcUFRWliooK1dbWevQpKyvT+++/r6lTpyo+Pt7dPm3aNDkcDr355pvutk2bNqldu3aaPXu2u81ut+vee+/V7t27dfTo0dYvqhG+1Nu1a9d670J4++23S5IOHDhQ79gVFRU6d+5cy0/akC81X+rChQv17v6uE+rPsdT8muu05W25Phs2bJDNZtOUKVMkhde2XJ/L6w23bbk5CBQhYsaMGYqPj5fdbteoUaO0d+/eRvu//vrrqq2t1b/+67+2yHiB9MEHHyg+Pl7Hjh3TNddcI4fDofj4eN1///3u29V/+umnunDhgq677jqPdTt06KDU1FTt27fP3bZv3z7179/f481Kklwul6SLu1uDyZd6G1JSUiLp4pvU5RYtWiSHwyG73a60tDS99957rTJ/fzSn5oKCAsXFxek73/mOEhMT9cQTT+j8+fMefUL9OZb8f57b8rZ8ufPnz+vNN99Uenq6+vbtKym8tuXL1VdvQ9rqttwcnEMRZB06dNCdd96pcePGqWvXrvr888+1dOlSXX/99dq1a5cGDx5c73rr16+X0+nUD3/4wxYZL5AOHjyoCxcuaMKECbr33nv17LPPKjc3V7/61a906tQpvfbaa+7b1judTq/1nU6ndu7c6f65uLi4wX6S9PXXX7dSJb7xpd76nDt3Tv/5n/+pq666yuMYelRUlMaMGaPbb79dPXv21KFDh/TCCy/o5ptv1ttvv63x48cHqrQG+Vrz1VdfrVGjRmnQoEGqqKjQpk2b9Itf/EIFBQV644033OOF+nMs+f88t+Vt+XLbt2/XP/7xD49wFE7b8uXqq7c+bXlbbpZgH3OBt4MHD1qxsbHWTTfdVO/yv//975Yka/78+S0yXqD98z//syXJuu+++zzaf/KTn1iSrIKCAuv3v/+9JcnKy8vzWv+ee+6xOnXq5DHezTff7NXvyy+/tCRZL774YkuX0Cy+1FufWbNmWZKsd955p8nH+Mc//mF1797duuaaa1pkzqb8rdmy/q/u3bt3e4wXys+xZflXc1vfli83efJkq3379taJEyfcbeG0LV+uvnrr05a35ebgkEcISkpK0oQJE5STk+P1MSPp4l80kppMxb6OF2ixsbGSLp7ZfKm6Y5C7d+9296murvZav6qqyr28bryG+l36eMHiS72XW7JkiV555RU9/fTTXmd/16dLly6aMWOG/v73v3t9iiIY/Km5zkMPPSTp4iGES8cL5ef40jk0p+a2vi1fqry8XFu2bNFNN92kK6+80t0eTtvypRqq93JtfVtuDgJFiOrdu7fOnTuniooKr2UbNmzQNddco+9///stMl6g1X0ksHv37h7tCQkJkqRvv/3WvYuzbnfppYqLiz0+Vuh0Ohvsd+njBYsv9V5q9erVevTRR3Xffffp3//9331+nN69e0uSTp48aTLdFtHcmi9VXx2h/hxfOofm1NzWt+VLbd68WWfPnvUKR+G0LV+qoXovFQ7bcnMQKELUoUOHZLfb5XA4PNrz8vJUWFjo8180TY0XDHVvnpde9Eb6v+Oj3bp108CBAxUdHe11Atq5c+e0f/9+paamuttSU1NVUFCgsrIyj755eXnu5cHkS711tmzZopkzZ+qOO+7Qr3/962Y9zqFDh7zGC5bm1Hy5+uoI9edYan7N4bAtX2r9+vVyOBy67bbbPNrDaVu+VEP11gmXbblZgn3MJdKVlpZ6te3fv99q3769ddttt3ktmzt3riXJKiwsbJHxguGvf/2rJcmaMmWKR/vkyZOt6Oho69ixY5ZlWdbYsWMtp9NplZWVufu8+uqrliTr3Xffdbft2bPH67PcVVVVVlJSkjV06NBWrqZpvta7Y8cOy263W6NGjbKqqqoaHK++57ioqMjq3LmzlZKS0rKT95MvNZ8+fdqrztraWmvSpEmWJOvjjz92t4f6c2xZvj/PdcJhW65TWlpqRUdHW/fcc0+9y8NlW67TVL3htC03B5/yCLJJkyYpNjZW6enpSkhI0Oeff66XX35ZHTt2VHZ2tkffmpoavfHGG/rBD36gq6++2ni8YBk8eLD+7d/+Tf/1X/+lCxcuaOTIkcrNzdXGjRuVlZXl3q35zDPPKD09XSNHjtTs2bNVVFSkZcuWacyYMR6XZh46dKgyMjKUlZWl0tJSJSUlac2aNTp8+LB++9vfBqtMN1/q/eqrr3TbbbfJZrPprrvu0saNGz3GSElJUUpKiiTpkUce0ZdffqnRo0erR48eOnz4sFatWqWKigr98pe/DEaJXnypOTc3V5MnT9bkyZOVlJSkyspK/fd//7c++ugjzZ49W0OGDHGPF+rPseT761oKn225zhtvvKELFy40uLclXLblOo3VG27bcrMEO9FEul/+8peWy+WyunTpYkVHR1tOp9OaOnWqx9XU6mzbts2SZL300kstMl4wnTt3zvr5z39u/dM//ZPVvn17Kykpqd4zuHfu3Gmlp6dbdrvd6tatmzVnzhyPv3LqVFZWWgsWLLASExOtmJgYKy0tzdq2bVsAKvFNU/XWXRWxoa8nn3zS3XfDhg3WiBEjrG7dulnR0dFW165drdtvv93jL/pQ0FTNhw4dsjIyMqy+fftadrvd6tixo/X973/fWrlypVVbW+s1Xqg/x5bl++s6nLZly7KsH/zgB1ZCQoJ14cKFBvuEy7ZsWY3XG47bsq9slmVZgY8xAAAgnHBSJgAAMEagAAAAxggUAADAGIECAAAYI1AAAABjBAoAAGCMQAEAAIwRKAAAgDECBQAAMEagAAAAxggUAADAGIECAAAYI1AAAABjBAoAAGCMQAEAAIz9P5Cu58xcNqARAAAAAElFTkSuQmCC",
-      "text/plain": [
-       "<Figure size 600x480 with 1 Axes>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    }
-   ],
-   "source": [
-    "import matplotlib.pyplot as plt\n",
-    "\n",
-    "plt.figure(figsize = (5, 4), dpi = 120)\n",
-    "plt.hist(transforms[\"train\"][0].frequencies.cpu().numpy(), bins = 50)\n",
-    "plt.grid()\n"
-   ]
-  },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Now we can get the best model path stored in the checkpoint callback."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 21,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "'/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=96-step=18333.ckpt'"
-      ]
-     },
-     "execution_count": 21,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "best_model_path = checkpoint_callback.best_model_path\n",
-    "best_model_path"
-   ]
-  },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Inference\n",
-    "\n",
-    "To obtain model scores, we will load the weights from the best checkpoint. To do this, we use the `LightningModule`, provide there the path to the checkpoint and the model instance."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 22,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import replay\n",
-    "torch.serialization.add_safe_globals([\n",
-    "    replay.nn.lightning.optimizer.OptimizerFactory,\n",
-    "    replay.nn.lightning.scheduler.LRSchedulerFactory\n",
-    "])"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 23,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "sasrec = make_sasrec(\n",
-    "    schema=tensor_schema,\n",
-    "    embedding_dim=EMBEDDING_DIM,\n",
-    "    max_sequence_length=MAX_SEQ_LEN,\n",
-    "    num_heads=NUM_HEADS,\n",
-    "    num_blocks=NUM_BLOCKS,\n",
-    "    dropout=DROPOUT,\n",
-    ")\n",
-    "\n",
-    "best_model = LightningModule.load_from_checkpoint(best_model_path, model=sasrec)\n",
-    "best_model.eval();"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Configure `ParquetModule` for inference"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 24,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "inference_metadata = {\"predict\": create_meta(shape=MAX_SEQ_LEN)}\n",
-    "\n",
-    "parquet_module = ParquetModule(\n",
-    "    predict_path=PREDICT_PATH,\n",
-    "    batch_size=BATCH_SIZE,\n",
-    "    metadata=inference_metadata,\n",
-    "    transforms=transforms,\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "During inference, we can use `TopItemsCallback`. Such callback allows you to get scores for each user throughout the entire catalog and get recommendations in the form of ids of items with the highest score values.\n",
-    "\n",
-    "\n",
-    "Recommendations can be fetched in four formats: PySpark DataFrame, Pandas DataFrame, Polars DataFrame or raw PyTorch tensors. Each of the types corresponds a callback. In this example, we'll be using the `PandasTopItemsCallback`."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 25,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "💡 Tip: For seamless cloud uploads and versioning, try installing [litmodels](https://pypi.org/project/litmodels/) to enable LitModelCheckpoint, which syncs automatically with the Lightning model registry.\n",
-      "GPU available: True (cuda), used: True\n",
-      "TPU available: False, using: 0 TPU cores\n",
-      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "9b94b2f1745140e5ae3bc1d63bdbbe4e",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Predicting: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    }
-   ],
-   "source": [
-    "from replay.nn.lightning.callback import PandasTopItemsCallback\n",
-    "\n",
-    "csv_logger = CSVLogger(save_dir=\"sasrec/logs/test\", name=\"SasRec-example\")\n",
-    "\n",
-    "TOPK = [1, 5, 10, 20]\n",
-    "\n",
-    "pandas_prediction_callback = PandasTopItemsCallback(\n",
-    "    top_k=max(TOPK),\n",
-    "    query_column=\"user_id\",\n",
-    "    item_column=\"item_id\",\n",
-    "    rating_column=\"score\",\n",
-    ")\n",
-    "\n",
-    "trainer = L.Trainer(callbacks=[pandas_prediction_callback], logger=csv_logger, inference_mode=True)\n",
-    "trainer.predict(best_model, datamodule=parquet_module, return_predictions=False)\n",
-    "\n",
-    "pandas_res = pandas_prediction_callback.get_result()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 26,
-   "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>486</td>\n",
-       "      <td>6.457348</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>0</th>\n",
-       "      <td>0</td>\n",
-       "      <td>224</td>\n",
-       "      <td>6.311622</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>0</th>\n",
-       "      <td>0</td>\n",
-       "      <td>210</td>\n",
-       "      <td>6.299765</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>0</th>\n",
-       "      <td>0</td>\n",
-       "      <td>101</td>\n",
-       "      <td>6.193144</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>0</th>\n",
-       "      <td>0</td>\n",
-       "      <td>2142</td>\n",
-       "      <td>6.179489</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>...</th>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6037</th>\n",
-       "      <td>6039</td>\n",
-       "      <td>84</td>\n",
-       "      <td>10.029899</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6037</th>\n",
-       "      <td>6039</td>\n",
-       "      <td>2634</td>\n",
-       "      <td>9.860862</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6037</th>\n",
-       "      <td>6039</td>\n",
-       "      <td>1490</td>\n",
-       "      <td>9.738086</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6037</th>\n",
-       "      <td>6039</td>\n",
-       "      <td>2633</td>\n",
-       "      <td>9.715384</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6037</th>\n",
-       "      <td>6039</td>\n",
-       "      <td>2497</td>\n",
-       "      <td>9.684793</td>\n",
-       "    </tr>\n",
-       "  </tbody>\n",
-       "</table>\n",
-       "<p>120760 rows × 3 columns</p>\n",
-       "</div>"
-      ],
-      "text/plain": [
-       "      user_id item_id      score\n",
-       "0           0     486   6.457348\n",
-       "0           0     224   6.311622\n",
-       "0           0     210   6.299765\n",
-       "0           0     101   6.193144\n",
-       "0           0    2142   6.179489\n",
-       "...       ...     ...        ...\n",
-       "6037     6039      84  10.029899\n",
-       "6037     6039    2634   9.860862\n",
-       "6037     6039    1490   9.738086\n",
-       "6037     6039    2633   9.715384\n",
-       "6037     6039    2497   9.684793\n",
-       "\n",
-       "[120760 rows x 3 columns]"
-      ]
-     },
-     "execution_count": 26,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "pandas_res"
-   ]
-  },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### Calculating metrics\n",
-    "\n",
-    "*test_gt* is already encoded, so we can use it for computing metrics."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 27,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from replay.metrics import MAP, OfflineMetrics, Precision, Recall\n",
-    "from replay.metrics.torch_metrics_builder import metrics_to_df"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 28,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "result_metrics = OfflineMetrics(\n",
-    "    [Recall(TOPK), Precision(TOPK), MAP(TOPK)],\n",
-    "    query_column=\"user_id\",\n",
-    "    rating_column=\"score\",\n",
-    ")(pandas_res, test_gt.explode(\"item_id\"))"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 29,
-   "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>5</th>\n",
-       "    </tr>\n",
-       "  </thead>\n",
-       "  <tbody>\n",
-       "    <tr>\n",
-       "      <th>MAP</th>\n",
-       "      <td>0.020371</td>\n",
-       "      <td>0.061189</td>\n",
-       "      <td>0.068575</td>\n",
-       "      <td>0.051101</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>Precision</th>\n",
-       "      <td>0.020371</td>\n",
-       "      <td>0.018731</td>\n",
-       "      <td>0.014790</td>\n",
-       "      <td>0.022093</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>Recall</th>\n",
-       "      <td>0.020371</td>\n",
-       "      <td>0.187314</td>\n",
-       "      <td>0.295793</td>\n",
-       "      <td>0.110467</td>\n",
-       "    </tr>\n",
-       "  </tbody>\n",
-       "</table>\n",
-       "</div>"
-      ],
-      "text/plain": [
-       "k                 1        10        20         5\n",
-       "MAP        0.020371  0.061189  0.068575  0.051101\n",
-       "Precision  0.020371  0.018731  0.014790  0.022093\n",
-       "Recall     0.020371  0.187314  0.295793  0.110467"
-      ]
-     },
-     "execution_count": 29,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "metrics_to_df(result_metrics)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Let's call the `inverse_transform` encoder's function to get the final dataframe with recommendations"
-   ]
-  },
-  {
-   "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></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>2012</td>\n",
-       "      <td>500</td>\n",
-       "      <td>6.457348</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>0</th>\n",
-       "      <td>2012</td>\n",
-       "      <td>231</td>\n",
-       "      <td>6.311622</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>0</th>\n",
-       "      <td>2012</td>\n",
-       "      <td>216</td>\n",
-       "      <td>6.299765</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>0</th>\n",
-       "      <td>2012</td>\n",
-       "      <td>104</td>\n",
-       "      <td>6.193144</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>0</th>\n",
-       "      <td>2012</td>\n",
-       "      <td>2335</td>\n",
-       "      <td>6.179489</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>...</th>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6037</th>\n",
-       "      <td>5727</td>\n",
-       "      <td>86</td>\n",
-       "      <td>10.029899</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6037</th>\n",
-       "      <td>5727</td>\n",
-       "      <td>2841</td>\n",
-       "      <td>9.860862</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6037</th>\n",
-       "      <td>5727</td>\n",
-       "      <td>1623</td>\n",
-       "      <td>9.738086</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6037</th>\n",
-       "      <td>5727</td>\n",
-       "      <td>2840</td>\n",
-       "      <td>9.715384</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6037</th>\n",
-       "      <td>5727</td>\n",
-       "      <td>2702</td>\n",
-       "      <td>9.684793</td>\n",
-       "    </tr>\n",
-       "  </tbody>\n",
-       "</table>\n",
-       "<p>120760 rows × 3 columns</p>\n",
-       "</div>"
-      ],
-      "text/plain": [
-       "      user_id  item_id      score\n",
-       "0        2012      500   6.457348\n",
-       "0        2012      231   6.311622\n",
-       "0        2012      216   6.299765\n",
-       "0        2012      104   6.193144\n",
-       "0        2012     2335   6.179489\n",
-       "...       ...      ...        ...\n",
-       "6037     5727       86  10.029899\n",
-       "6037     5727     2841   9.860862\n",
-       "6037     5727     1623   9.738086\n",
-       "6037     5727     2840   9.715384\n",
-       "6037     5727     2702   9.684793\n",
-       "\n",
-       "[120760 rows x 3 columns]"
-      ]
-     },
-     "execution_count": 30,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "encoder.inverse_transform(pandas_res)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "new_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.12.3"
-  },
-  "orig_nbformat": 4
- },
- "nbformat": 4,
- "nbformat_minor": 2
-}
diff --git a/examples/09_sasrec_example_sampled.ipynb b/examples/09_sasrec_example_sampled.ipynb
deleted file mode 100644
index 6672edd85..000000000
--- a/examples/09_sasrec_example_sampled.ipynb
+++ /dev/null
@@ -1,4744 +0,0 @@
-{
- "cells": [
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Example of SasRec training/inference with Parquet Module"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 1,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Seed set to 42\n"
-     ]
-    }
-   ],
-   "source": [
-    "from typing import Optional\n",
-    "\n",
-    "import lightning as L\n",
-    "import pandas as pd\n",
-    "\n",
-    "L.seed_everything(42)\n",
-    "\n",
-    "import warnings\n",
-    "warnings.filterwarnings(\"ignore\")"
-   ]
-  },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Preparing data\n",
-    "In this example, we will be using the MovieLens dataset, namely the 1m subset. It's demonstrated a simple case, so only item ids will be used as model input.\n",
-    "\n",
-    "---\n",
-    "**NOTE**\n",
-    "\n",
-    "Current implementation of SasRec is able to handle item and interactions features. It does not take into account user features. \n",
-    "\n",
-    "---"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 2,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "interactions = pd.read_csv(\"./data/ml1m_ratings.dat\", sep=\"\\t\", names=[\"user_id\", \"item_id\",\"rating\",\"timestamp\"])\n",
-    "interactions = interactions.drop(columns=[\"rating\"])"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 3,
-   "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>timestamp</th>\n",
-       "    </tr>\n",
-       "  </thead>\n",
-       "  <tbody>\n",
-       "    <tr>\n",
-       "      <th>1000138</th>\n",
-       "      <td>6040</td>\n",
-       "      <td>858</td>\n",
-       "      <td>0</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1000153</th>\n",
-       "      <td>6040</td>\n",
-       "      <td>2384</td>\n",
-       "      <td>1</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>999873</th>\n",
-       "      <td>6040</td>\n",
-       "      <td>593</td>\n",
-       "      <td>2</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1000007</th>\n",
-       "      <td>6040</td>\n",
-       "      <td>1961</td>\n",
-       "      <td>3</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1000192</th>\n",
-       "      <td>6040</td>\n",
-       "      <td>2019</td>\n",
-       "      <td>4</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>...</th>\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>446</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>825438</th>\n",
-       "      <td>4958</td>\n",
-       "      <td>1407</td>\n",
-       "      <td>447</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>825724</th>\n",
-       "      <td>4958</td>\n",
-       "      <td>3264</td>\n",
-       "      <td>448</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>825731</th>\n",
-       "      <td>4958</td>\n",
-       "      <td>2634</td>\n",
-       "      <td>449</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>825603</th>\n",
-       "      <td>4958</td>\n",
-       "      <td>1924</td>\n",
-       "      <td>450</td>\n",
-       "    </tr>\n",
-       "  </tbody>\n",
-       "</table>\n",
-       "<p>1000209 rows × 3 columns</p>\n",
-       "</div>"
-      ],
-      "text/plain": [
-       "         user_id  item_id  timestamp\n",
-       "1000138     6040      858          0\n",
-       "1000153     6040     2384          1\n",
-       "999873      6040      593          2\n",
-       "1000007     6040     1961          3\n",
-       "1000192     6040     2019          4\n",
-       "...          ...      ...        ...\n",
-       "825793      4958     2399        446\n",
-       "825438      4958     1407        447\n",
-       "825724      4958     3264        448\n",
-       "825731      4958     2634        449\n",
-       "825603      4958     1924        450\n",
-       "\n",
-       "[1000209 rows x 3 columns]"
-      ]
-     },
-     "execution_count": 3,
-     "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"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### Encode catagorical data.\n",
-    "To ensure all categorical data is fit for training, it needs to be encoded using the `LabelEncoder` class. Create an instance of the encoder, providing a `LabelEncodingRule` for each categorcial column in the dataset that will be used in model. Note that ids of users and ids of items are always used."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 4,
-   "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>timestamp</th>\n",
-       "      <th>user_id</th>\n",
-       "      <th>item_id</th>\n",
-       "    </tr>\n",
-       "  </thead>\n",
-       "  <tbody>\n",
-       "    <tr>\n",
-       "      <th>0</th>\n",
-       "      <td>12</td>\n",
-       "      <td>0</td>\n",
-       "      <td>0</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1</th>\n",
-       "      <td>68</td>\n",
-       "      <td>1</td>\n",
-       "      <td>0</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>2</th>\n",
-       "      <td>67</td>\n",
-       "      <td>2</td>\n",
-       "      <td>0</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>3</th>\n",
-       "      <td>12</td>\n",
-       "      <td>3</td>\n",
-       "      <td>0</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>4</th>\n",
-       "      <td>140</td>\n",
-       "      <td>4</td>\n",
-       "      <td>0</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>...</th>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1000204</th>\n",
-       "      <td>14</td>\n",
-       "      <td>4555</td>\n",
-       "      <td>3705</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1000205</th>\n",
-       "      <td>90</td>\n",
-       "      <td>2813</td>\n",
-       "      <td>3705</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1000206</th>\n",
-       "      <td>70</td>\n",
-       "      <td>2404</td>\n",
-       "      <td>3705</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1000207</th>\n",
-       "      <td>25</td>\n",
-       "      <td>5835</td>\n",
-       "      <td>3705</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1000208</th>\n",
-       "      <td>380</td>\n",
-       "      <td>979</td>\n",
-       "      <td>3705</td>\n",
-       "    </tr>\n",
-       "  </tbody>\n",
-       "</table>\n",
-       "<p>1000209 rows × 3 columns</p>\n",
-       "</div>"
-      ],
-      "text/plain": [
-       "         timestamp  user_id  item_id\n",
-       "0               12        0        0\n",
-       "1               68        1        0\n",
-       "2               67        2        0\n",
-       "3               12        3        0\n",
-       "4              140        4        0\n",
-       "...            ...      ...      ...\n",
-       "1000204         14     4555     3705\n",
-       "1000205         90     2813     3705\n",
-       "1000206         70     2404     3705\n",
-       "1000207         25     5835     3705\n",
-       "1000208        380      979     3705\n",
-       "\n",
-       "[1000209 rows x 3 columns]"
-      ]
-     },
-     "execution_count": 4,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "from replay.preprocessing.label_encoder import LabelEncoder, LabelEncodingRule\n",
-    "\n",
-    "encoder = LabelEncoder(\n",
-    "    [\n",
-    "        LabelEncodingRule(\"user_id\", default_value=\"last\"),\n",
-    "        LabelEncodingRule(\"item_id\", default_value=\"last\"),\n",
-    "    ]\n",
-    ")\n",
-    "interactions = interactions.sort_values(by=\"item_id\", ascending=True)\n",
-    "encoded_interactions = encoder.fit_transform(interactions)\n",
-    "encoded_interactions"
-   ]
-  },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### Split interactions into the train, validation and test datasets using LastNSplitter\n",
-    "We use widespread splitting strategy Last-One-Out. We filter out cold items and users for simplicity."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 5,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from replay.splitters import LastNSplitter\n",
-    "\n",
-    "splitter = LastNSplitter(\n",
-    "    N=1,\n",
-    "    divide_column=\"user_id\",\n",
-    "    query_column=\"user_id\",\n",
-    "    strategy=\"interactions\",\n",
-    "    drop_cold_users=True,\n",
-    "    drop_cold_items=True\n",
-    ")\n",
-    "\n",
-    "test_events, test_gt = splitter.split(encoded_interactions)\n",
-    "validation_events, validation_gt = splitter.split(test_events)\n",
-    "train_events = validation_events"
-   ]
-  },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### Dataset preprocessing (\"baking\")\n",
-    "SasRec expects each user in the batch to provide their events in form of a sequence. For this reason, the event splits must be properly processed using the `groupby_sequences` function provided by RePlay."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 6,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from replay.data.nn.utils import groupby_sequences\n",
-    "\n",
-    "\n",
-    "def bake_data(full_data):\n",
-    "    grouped_interactions = groupby_sequences(events=full_data, groupby_col=\"user_id\", sort_col=\"timestamp\")\n",
-    "    return grouped_interactions"
-   ]
-  },
-  {
-   "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>timestamp</th>\n",
-       "      <th>item_id</th>\n",
-       "    </tr>\n",
-       "  </thead>\n",
-       "  <tbody>\n",
-       "    <tr>\n",
-       "      <th>0</th>\n",
-       "      <td>0</td>\n",
-       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
-       "      <td>[2426, 822, 2733, 2587, 2937, 3618, 2943, 708,...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1</th>\n",
-       "      <td>1</td>\n",
-       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
-       "      <td>[3272, 3026, 2760, 851, 346, 3393, 1107, 515, ...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>2</th>\n",
-       "      <td>2</td>\n",
-       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
-       "      <td>[579, 1140, 1154, 2426, 1524, 1260, 2160, 2621...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>3</th>\n",
-       "      <td>3</td>\n",
-       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
-       "      <td>[1781, 2940, 2468, 890, 948, 106, 593, 309, 49...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>4</th>\n",
-       "      <td>4</td>\n",
-       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
-       "      <td>[1108, 2229, 21, 2435, 2142, 106, 1167, 593, 1...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>...</th>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6035</th>\n",
-       "      <td>6035</td>\n",
-       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
-       "      <td>[2426, 1279, 3151, 3321, 1178, 3301, 2501, 278...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6036</th>\n",
-       "      <td>6036</td>\n",
-       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
-       "      <td>[1592, 2302, 1633, 1813, 2879, 1482, 2651, 250...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6037</th>\n",
-       "      <td>6037</td>\n",
-       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
-       "      <td>[1971, 3500, 2077, 1666, 1399, 2651, 2748, 283...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6038</th>\n",
-       "      <td>6038</td>\n",
-       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
-       "      <td>[1486, 1485, 3384, 3512, 3302, 3126, 3650, 330...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6039</th>\n",
-       "      <td>6039</td>\n",
-       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
-       "      <td>[2432, 2960, 1848, 2114, 2142, 3091, 3248, 317...</td>\n",
-       "    </tr>\n",
-       "  </tbody>\n",
-       "</table>\n",
-       "<p>6040 rows × 3 columns</p>\n",
-       "</div>"
-      ],
-      "text/plain": [
-       "      user_id                                          timestamp  \\\n",
-       "0           0  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
-       "1           1  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
-       "2           2  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
-       "3           3  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
-       "4           4  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
-       "...       ...                                                ...   \n",
-       "6035     6035  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
-       "6036     6036  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
-       "6037     6037  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
-       "6038     6038  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
-       "6039     6039  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
-       "\n",
-       "                                                item_id  \n",
-       "0     [2426, 822, 2733, 2587, 2937, 3618, 2943, 708,...  \n",
-       "1     [3272, 3026, 2760, 851, 346, 3393, 1107, 515, ...  \n",
-       "2     [579, 1140, 1154, 2426, 1524, 1260, 2160, 2621...  \n",
-       "3     [1781, 2940, 2468, 890, 948, 106, 593, 309, 49...  \n",
-       "4     [1108, 2229, 21, 2435, 2142, 106, 1167, 593, 1...  \n",
-       "...                                                 ...  \n",
-       "6035  [2426, 1279, 3151, 3321, 1178, 3301, 2501, 278...  \n",
-       "6036  [1592, 2302, 1633, 1813, 2879, 1482, 2651, 250...  \n",
-       "6037  [1971, 3500, 2077, 1666, 1399, 2651, 2748, 283...  \n",
-       "6038  [1486, 1485, 3384, 3512, 3302, 3126, 3650, 330...  \n",
-       "6039  [2432, 2960, 1848, 2114, 2142, 3091, 3248, 317...  \n",
-       "\n",
-       "[6040 rows x 3 columns]"
-      ]
-     },
-     "execution_count": 7,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "train_events = bake_data(train_events)\n",
-    "\n",
-    "validation_events = bake_data(validation_events)\n",
-    "validation_gt = bake_data(validation_gt)\n",
-    "\n",
-    "test_events = bake_data(test_events)\n",
-    "test_gt = bake_data(test_gt)\n",
-    "\n",
-    "train_events"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "To ensure we don't have unknown users in ground truth, we join validation events and validation ground truth (also join test events and test ground truth correspondingly) by user ids to leave only the common ones.  "
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 8,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def add_gt_to_events(events_df, gt_df):\n",
-    "    gt_to_join = gt_df[[\"user_id\", \"item_id\"]].rename(columns={\"item_id\": \"ground_truth\"})\n",
-    "\n",
-    "    events_df = events_df.merge(gt_to_join, on=\"user_id\", how=\"inner\")\n",
-    "    return events_df\n",
-    "\n",
-    "validation_events = add_gt_to_events(validation_events, validation_gt)\n",
-    "test_events = add_gt_to_events(test_events, test_gt)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 9,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from pathlib import Path\n",
-    "\n",
-    "data_dir = Path(\"temp/data/\")\n",
-    "data_dir.mkdir(parents=True, exist_ok=True)\n",
-    "\n",
-    "TRAIN_PATH = data_dir / \"train.parquet\"\n",
-    "VAL_PATH = data_dir / \"val.parquet\"\n",
-    "PREDICT_PATH = data_dir / \"test.parquet\"\n",
-    "\n",
-    "ENCODER_PATH = data_dir / \"encoder\""
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 10,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "train_events.to_parquet(TRAIN_PATH)\n",
-    "validation_events.to_parquet(VAL_PATH)\n",
-    "test_events.to_parquet(PREDICT_PATH)\n",
-    "\n",
-    "encoder.save(ENCODER_PATH)"
-   ]
-  },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Prepare to model training\n",
-    "### 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. It is required by the SasRec model to correctly create embeddings for every source column. Note that user_id does not required in `TensorSchema`.\n",
-    "\n",
-    "Note that the **padding value** is the next value (item_id) after the last one. **Cardinality** is the number of unique values ​​given the padding value."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 11,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from replay.data import FeatureHint, FeatureType\n",
-    "from replay.data.nn import TensorFeatureInfo, TensorSchema\n",
-    "\n",
-    "\n",
-    "EMBEDDING_DIM = 64\n",
-    "\n",
-    "encoder = encoder.load(ENCODER_PATH)\n",
-    "NUM_UNIQUE_ITEMS = len(encoder.mapping[\"item_id\"])\n",
-    "\n",
-    "tensor_schema = TensorSchema(\n",
-    "    [\n",
-    "        TensorFeatureInfo(\n",
-    "            name=\"item_id\",\n",
-    "            is_seq=True,\n",
-    "            padding_value=NUM_UNIQUE_ITEMS,\n",
-    "            cardinality=NUM_UNIQUE_ITEMS + 1,  # taking into account padding\n",
-    "            embedding_dim=EMBEDDING_DIM,\n",
-    "            feature_type=FeatureType.CATEGORICAL,\n",
-    "            feature_hint=FeatureHint.ITEM_ID,\n",
-    "        )\n",
-    "    ]\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### Configure ParquetModule and transformation pipelines\n",
-    "\n",
-    "The `ParquetModule` class enables training of models on large datasets by reading data in batch-wise way. This class initialized with **paths to every data split, a metadata dict containing information about shape and padding value of every column and a dict of transforms**. `ParquetModule`'s  \"transform pipelines\" are stage-specific modules implementing additional preprocessing to be performed on batch level right before the forward pass.  \n",
-    "\n",
-    "For SasRec model, RePlay provides a function that generates a sequence of appropriate transforms for each data split named **make_default_sasrec_transforms**.\n",
-    "\n",
-    "Internally this function creates the following transforms:\n",
-    "1) Training:\n",
-    "    1. Create a target, which contains the shifted item sequence that represents the next item in the sequence (for the next item prediction task).\n",
-    "    2. Rename features to match it with expected format by the model during training.\n",
-    "    3. Unsqueeze target (*positive_labels*) and it's padding mask (*target_padding_mask*) for getting required shape of this tensors for loss computation.\n",
-    "    4. Group input features to be embed in expected format.\n",
-    "\n",
-    "2) Validation/Inference:\n",
-    "    1. Rename/group features to match it with expected format by the model during valdiation/inference.\n",
-    "\n",
-    "If a different set of transforms is required, you can create them yourself and submit them to the ParquetModule in the form of a dictionary where the key is the name of the split, and the value is the list of transforms. Available transforms are in the replay/nn/transforms/.\n",
-    "\n",
-    "**Note:** One of the transforms for the training data prepares the initial sequence for the task of Next Item Prediction so it shifts the sequence of items. For the final sequence length to be correct, you need to set shape of item_id in metadata as **model sequence length + shift**. Default shift value is 1."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 12,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import copy\n",
-    "\n",
-    "import torch\n",
-    "\n",
-    "from replay.data.nn import TensorSchema\n",
-    "from replay.nn.transform import GroupTransform, NextTokenTransform, RenameTransform, UnsqueezeTransform, UniformNegativeSamplingTransform"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 20,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def make_sasrec_transforms(\n",
-    "    tensor_schema: TensorSchema, query_column: str = \"query_id\", num_negative_samples: int = 128,\n",
-    ") -> dict[str, list[torch.nn.Module]]:\n",
-    "    item_column = tensor_schema.item_id_feature_name\n",
-    "    vocab_size = tensor_schema[item_column].cardinality\n",
-    "    train_transforms = [\n",
-    "        UniformNegativeSamplingTransform(vocab_size, num_negative_samples),\n",
-    "        NextTokenTransform(label_field=item_column, query_features=query_column, shift=1),\n",
-    "        RenameTransform(\n",
-    "            {\n",
-    "                query_column: \"query_id\",\n",
-    "                f\"{item_column}_mask\": \"padding_mask\",\n",
-    "                \"positive_labels_mask\": \"target_padding_mask\",\n",
-    "            }\n",
-    "        ),\n",
-    "        UnsqueezeTransform(\"target_padding_mask\", -1),\n",
-    "        UnsqueezeTransform(\"positive_labels\", -1),\n",
-    "        GroupTransform({\"feature_tensors\": [item_column]}),\n",
-    "    ]\n",
-    "\n",
-    "    val_transforms = [\n",
-    "        RenameTransform({query_column: \"query_id\", f\"{item_column}_mask\": \"padding_mask\"}),\n",
-    "        GroupTransform({\"feature_tensors\": [item_column]}),\n",
-    "    ]\n",
-    "    test_transforms = copy.deepcopy(val_transforms)\n",
-    "\n",
-    "    predict_transforms = copy.deepcopy(val_transforms)\n",
-    "\n",
-    "    transforms = {\n",
-    "        \"train\": train_transforms,\n",
-    "        \"validate\": val_transforms,\n",
-    "        \"test\": test_transforms,\n",
-    "        \"predict\": predict_transforms,\n",
-    "    }\n",
-    "\n",
-    "    return transforms\n",
-    "\n",
-    "transforms = make_sasrec_transforms(tensor_schema, query_column=\"user_id\")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 21,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "MAX_SEQ_LEN = 50\n",
-    "\n",
-    "def create_meta(shape: int, gt_shape: Optional[int] = None):\n",
-    "    meta = {\n",
-    "        \"user_id\": {},\n",
-    "        \"item_id\": {\"shape\": shape, \"padding\": tensor_schema[\"item_id\"].padding_value},\n",
-    "    }\n",
-    "    if gt_shape is not None:\n",
-    "        meta.update({\"ground_truth\": {\"shape\": gt_shape, \"padding\": -1}})\n",
-    "\n",
-    "    return meta\n",
-    "\n",
-    "train_metadata = {\n",
-    "    \"train\": create_meta(shape=MAX_SEQ_LEN+1),\n",
-    "    \"validate\": create_meta(shape=MAX_SEQ_LEN, gt_shape=1),\n",
-    "}"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 22,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from replay.data.nn import ParquetModule\n",
-    "\n",
-    "BATCH_SIZE = 32\n",
-    "\n",
-    "parquet_module = ParquetModule(\n",
-    "    train_path=TRAIN_PATH,\n",
-    "    validate_path=VAL_PATH,\n",
-    "    batch_size=BATCH_SIZE,\n",
-    "    metadata=train_metadata,\n",
-    "    transforms=transforms,\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Train model\n",
-    "### Create SasRec model instance and run the training stage using lightning\n",
-    "We may now train the model using the Lightning trainer class. \n",
-    "\n",
-    "RePlay's implementation of SasRec is designed in a modular, **block-based approach**. Instead of passing configuration parameters to the constructor, SasRec is now built by providing fully initialized components that makes the model more flexible and easier to extend.\n",
-    "\n",
-    "#### Default Configuration\n",
-    "\n",
-    "Default SasRec model may be created quickly via method **from_params**. Default model instance has CE loss, original SasRec transformer layes, and embeddings are aggregated via sum."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 23,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from replay.nn.sequential import SasRec\n",
-    "from typing import Literal\n",
-    "def make_sasrec(\n",
-    "    schema: TensorSchema,\n",
-    "    embedding_dim: int = 192,\n",
-    "    num_heads: int = 4,\n",
-    "    num_blocks: int = 2,\n",
-    "    max_sequence_length: int = 50,\n",
-    "    dropout: float = 0.3,\n",
-    "    excluded_features: Optional[list[str]] = None,\n",
-    "    categorical_list_feature_aggregation_method: Literal[\"sum\", \"mean\", \"max\"] = \"sum\",\n",
-    ") -> SasRec:\n",
-    "    from replay.nn.sequential.sasrec import SasRecBody, SasRecTransformerLayer\n",
-    "    from replay.nn.agg import SumAggregator\n",
-    "    from replay.nn.embedding import SequenceEmbedding\n",
-    "    from replay.nn.loss import CE, CESampled\n",
-    "    from replay.nn.mask import DefaultAttentionMask\n",
-    "    from replay.nn.sequential.sasrec.agg import PositionAwareAggregator\n",
-    "    from replay.nn.sequential.sasrec.transformer import SasRecTransformerLayer\n",
-    "    excluded_features = [\n",
-    "        schema.query_id_feature_name,\n",
-    "        schema.timestamp_feature_name,\n",
-    "        *(excluded_features or []),\n",
-    "    ]\n",
-    "    excluded_features = list(set(excluded_features))\n",
-    "    body = SasRecBody(\n",
-    "        embedder=SequenceEmbedding(\n",
-    "            schema=schema,\n",
-    "            categorical_list_feature_aggregation_method=categorical_list_feature_aggregation_method,\n",
-    "            excluded_features=excluded_features,\n",
-    "        ),\n",
-    "        embedding_aggregator=PositionAwareAggregator(\n",
-    "            embedding_aggregator=SumAggregator(embedding_dim=embedding_dim),\n",
-    "            max_sequence_length=max_sequence_length,\n",
-    "            dropout=dropout,\n",
-    "        ),\n",
-    "        attn_mask_builder=DefaultAttentionMask(\n",
-    "            reference_feature_name=schema.item_id_feature_name,\n",
-    "            num_heads=num_heads,\n",
-    "        ),\n",
-    "        encoder=SasRecTransformerLayer(\n",
-    "            embedding_dim=embedding_dim,\n",
-    "            num_heads=num_heads,\n",
-    "            num_blocks=num_blocks,\n",
-    "            dropout=dropout,\n",
-    "            activation=\"relu\",\n",
-    "        ),\n",
-    "        output_normalization=torch.nn.LayerNorm(embedding_dim),\n",
-    "    )\n",
-    "    padding_idx = schema.item_id_features.item().padding_value\n",
-    "    return SasRec(\n",
-    "        body=body,\n",
-    "        loss=CESampled(padding_idx=padding_idx),\n",
-    "    )"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 24,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "NUM_BLOCKS = 2\n",
-    "NUM_HEADS = 2\n",
-    "DROPOUT = 0.3\n",
-    "\n",
-    "sasrec = make_sasrec(\n",
-    "    schema=tensor_schema,\n",
-    "    embedding_dim=EMBEDDING_DIM,\n",
-    "    max_sequence_length=MAX_SEQ_LEN,\n",
-    "    num_heads=NUM_HEADS,\n",
-    "    num_blocks=NUM_BLOCKS,\n",
-    "    dropout=DROPOUT,\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "A universal PyTorch Lightning module is provided. It can work with any NN model."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 25,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from replay.nn.lightning.optimizer import OptimizerFactory\n",
-    "from replay.nn.lightning.scheduler import LRSchedulerFactory\n",
-    "from replay.nn.lightning import LightningModule\n",
-    "\n",
-    "model = LightningModule(\n",
-    "    sasrec,\n",
-    "    optimizer_factory=OptimizerFactory(),\n",
-    "    lr_scheduler_factory=LRSchedulerFactory(),\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "To facilitate training, we add the following callbacks:\n",
-    "1) `ModelCheckpoint` - to save the best trained model based on its Recall metric. It's a default Lightning Callback.\n",
-    "1) `ComputeMetricsCallback` - to display a detailed validation metric matrix after each epoch. It's a custom RePlay callback for computing recsys metrics on validation and test stages.\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 26,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "GPU available: True (cuda), used: True\n",
-      "TPU available: False, using: 0 TPU cores\n",
-      "You are using a CUDA device ('NVIDIA GeForce RTX 3060 Laptop GPU') 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 | Mode  | FLOPs\n",
-      "-------------------------------------------------\n",
-      "0 | model | SasRec | 291 K  | train | 0    \n",
-      "-------------------------------------------------\n",
-      "291 K     Trainable params\n",
-      "0         Non-trainable params\n",
-      "291 K     Total params\n",
-      "1.164     Total estimated model params size (MB)\n",
-      "39        Modules in train mode\n",
-      "0         Modules in eval mode\n",
-      "0         Total Flops\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "bdbe6004cd2b40f09355b6ed047313af",
-       "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": "5f67652c4a59456087990b5eb2a048c6",
-       "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": "cde930c138ce478fb113c82a14740877",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 0, global step 189: 'recall@10' reached 0.01507 (best 0.01507), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=0-step=189-v3.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.001821  0.005005  0.005654  0.004220\n",
-      "ndcg    0.001821  0.007341  0.009693  0.005424\n",
-      "recall  0.001821  0.015069  0.024342  0.009107\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "e2c4de9f0d544a3c9fe936be40e767ea",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 1, global step 378: 'recall@10' reached 0.02302 (best 0.02302), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=1-step=378-v2.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.002981  0.007374  0.008710  0.006044\n",
-      "ndcg    0.002981  0.010976  0.016031  0.007759\n",
-      "recall  0.002981  0.023017  0.043385  0.013082\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "ebeeac2ff096409c9d92c8f500f5a4fa",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 2, global step 567: 'recall@10' reached 0.02981 (best 0.02981), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=2-step=567.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.003809  0.009334  0.011016  0.007628\n",
-      "ndcg    0.003809  0.014037  0.020154  0.009877\n",
-      "recall  0.003809  0.029806  0.053982  0.016890\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "af33cdb8efe947fb96e680d605c57ad8",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 3, global step 756: 'recall@10' reached 0.04554 (best 0.04554), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=3-step=756-v3.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.004637  0.013212  0.015986  0.010153\n",
-      "ndcg    0.004637  0.020583  0.030981  0.013059\n",
-      "recall  0.004637  0.045537  0.087266  0.022024\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "a80a1830e850491ba4bb54053894886b",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 4, global step 945: 'recall@10' reached 0.06143 (best 0.06143), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=4-step=945-v4.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.008611  0.020454  0.024102  0.016998\n",
-      "ndcg    0.008611  0.029885  0.043547  0.021429\n",
-      "recall  0.008611  0.061434  0.116244  0.035105\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "2e60d216a807444693f4880a5d3026d9",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 5, global step 1134: 'recall@10' reached 0.08147 (best 0.08147), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=5-step=1134-v1.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.006789  0.023326  0.027358  0.018168\n",
-      "ndcg    0.006789  0.036654  0.051677  0.023948\n",
-      "recall  0.006789  0.081470  0.141580  0.041729\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "9a55138f358d4264b3f326e2b087c7f7",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 6, global step 1323: 'recall@10' reached 0.09141 (best 0.09141), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=6-step=1323-v2.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.008776  0.027117  0.031929  0.021949\n",
-      "ndcg    0.008776  0.041864  0.059872  0.028990\n",
-      "recall  0.008776  0.091406  0.163603  0.050836\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "3be95e5408e844238aea4dad092d3f6c",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 7, global step 1512: 'recall@10' reached 0.10383 (best 0.10383), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=7-step=1512.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k             1        10        20         5\n",
-      "map     0.01126  0.032126  0.037192  0.026006\n",
-      "ndcg    0.01126  0.048599  0.067473  0.033582\n",
-      "recall  0.01126  0.103825  0.179334  0.056963\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "9fac6ccdf1cf496fa5a226221d9e35ee",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 8, global step 1701: 'recall@10' reached 0.11244 (best 0.11244), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=8-step=1701.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.012088  0.033863  0.039145  0.026379\n",
-      "ndcg    0.012088  0.051835  0.071279  0.033566\n",
-      "recall  0.012088  0.112436  0.189767  0.055638\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "200aaa74132545ebb133abb0f7132f38",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 9, global step 1890: 'recall@10' reached 0.12038 (best 0.12038), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=9-step=1890-v2.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.014075  0.037768  0.043720  0.030035\n",
-      "ndcg    0.014075  0.056689  0.078432  0.037731\n",
-      "recall  0.014075  0.120384  0.206491  0.061268\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "d95e0c88f6664255a963b4dfd25fb420",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 10, global step 2079: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.014572  0.038940  0.045264  0.032039\n",
-      "ndcg    0.014572  0.057550  0.080690  0.040461\n",
-      "recall  0.014572  0.119887  0.211624  0.066236\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "d0e225b1fda2442894cbded20f774fbe",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 11, global step 2268: 'recall@10' reached 0.13313 (best 0.13313), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=11-step=2268.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.014572  0.041118  0.047172  0.032944\n",
-      "ndcg    0.014572  0.062227  0.084644  0.042165\n",
-      "recall  0.014572  0.133135  0.222553  0.070541\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "289a63e4e8d248e9830917c74921b213",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 12, global step 2457: 'recall@10' reached 0.13578 (best 0.13578), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=12-step=2457-v1.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.012916  0.041002  0.048065  0.032950\n",
-      "ndcg    0.012916  0.062827  0.088817  0.043172\n",
-      "recall  0.012916  0.135784  0.239112  0.074681\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "da5831ed697c40a5be21af8ed077c0af",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 13, global step 2646: 'recall@10' reached 0.13761 (best 0.13761), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=13-step=2646.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016393  0.043745  0.049922  0.035814\n",
-      "ndcg    0.016393  0.065313  0.088050  0.045814\n",
-      "recall  0.016393  0.137606  0.228018  0.076668\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "df8c7feb16c040749320c57ac1fb6929",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 14, global step 2835: 'recall@10' reached 0.14572 (best 0.14572), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=14-step=2835.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.013744  0.044572  0.051397  0.036082\n",
-      "ndcg    0.013744  0.067861  0.093062  0.046887\n",
-      "recall  0.013744  0.145719  0.246067  0.079980\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "24ca4ca06969457cbadeb87d9b786a1b",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 15, global step 3024: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.014572  0.043037  0.049347  0.033872\n",
-      "ndcg    0.014572  0.066414  0.089676  0.043945\n",
-      "recall  0.014572  0.145057  0.237622  0.075012\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "9a5ee338856b4205ba96c417dc1cd897",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 16, global step 3213: 'recall@10' reached 0.15301 (best 0.15301), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=16-step=3213.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.020533  0.050175  0.056921  0.041307\n",
-      "ndcg    0.020533  0.073800  0.098705  0.052068\n",
-      "recall  0.020533  0.153005  0.252194  0.085279\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "88de3cabb37446c4bd994c385fb71254",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 17, global step 3402: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016228  0.044922  0.052359  0.035591\n",
-      "ndcg    0.016228  0.068935  0.096304  0.045990\n",
-      "recall  0.016228  0.149859  0.258652  0.078159\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "1424cea483354fcd960f0b2fa686d6ff",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 18, global step 3591: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017056  0.048151  0.055104  0.039250\n",
-      "ndcg    0.017056  0.072077  0.097874  0.050346\n",
-      "recall  0.017056  0.152012  0.255009  0.084451\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "3ec2ba11a2da4485965830c2ef84fef7",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 19, global step 3780: 'recall@10' reached 0.16443 (best 0.16443), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=19-step=3780.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017553  0.050482  0.057555  0.040688\n",
-      "ndcg    0.017553  0.076663  0.102664  0.052659\n",
-      "recall  0.017553  0.164431  0.267760  0.089584\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "019723cee5394e998e598c7a50892aa4",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 20, global step 3969: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.015731  0.048458  0.055842  0.039204\n",
-      "ndcg    0.015731  0.074221  0.101518  0.051256\n",
-      "recall  0.015731  0.160623  0.269415  0.088425\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "6041f647e1434cb3892b2a2f0f3dcecd",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 21, global step 4158: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.018049  0.049611  0.056758  0.040399\n",
-      "ndcg    0.018049  0.074154  0.100495  0.051515\n",
-      "recall  0.018049  0.156317  0.261136  0.085610\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "6ee1a62f7622477bb8648edcd1b92588",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 22, global step 4347: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016393  0.050516  0.057873  0.040887\n",
-      "ndcg    0.016393  0.076624  0.103727  0.053190\n",
-      "recall  0.016393  0.163603  0.271403  0.090909\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "3f8eab96a86443af85c2f62d01a9cf5a",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 23, global step 4536: 'recall@10' reached 0.16559 (best 0.16559), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=23-step=4536.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.014903  0.050339  0.057592  0.040376\n",
-      "ndcg    0.014903  0.076891  0.103476  0.052462\n",
-      "recall  0.014903  0.165590  0.271071  0.089419\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "9c8ded0ade1b4879986fd73a0ebebc2c",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 24, global step 4725: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k             1        10        20         5\n",
-      "map     0.01689  0.049176  0.056484  0.040109\n",
-      "ndcg    0.01689  0.074093  0.100858  0.052004\n",
-      "recall  0.01689  0.157145  0.263289  0.088591\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "c979af1f9bbc4295a23eb8b2472fb96d",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 25, global step 4914: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.018049  0.051505  0.058838  0.042118\n",
-      "ndcg    0.018049  0.077604  0.104511  0.054477\n",
-      "recall  0.018049  0.164928  0.271734  0.092565\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "3d0e8715d4a249f1966939aecd397b8a",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 26, global step 5103: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017718  0.051380  0.059348  0.042295\n",
-      "ndcg    0.017718  0.076988  0.106090  0.054788\n",
-      "recall  0.017718  0.162279  0.277529  0.093227\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "73181d4a5c4e4013b7a8eeebf742ce9e",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 27, global step 5292: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017221  0.048906  0.057236  0.039648\n",
-      "ndcg    0.017221  0.074121  0.104692  0.051467\n",
-      "recall  0.017221  0.158470  0.279848  0.087928\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "1634e7c2003d44d697030ab44288f0b1",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 28, global step 5481: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k             1        10        20         5\n",
-      "map     0.01954  0.051349  0.059358  0.041872\n",
-      "ndcg    0.01954  0.076578  0.106111  0.053326\n",
-      "recall  0.01954  0.161119  0.278689  0.088591\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "908c27a408b84ef7b6c91700b34f1f0d",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 29, global step 5670: 'recall@10' reached 0.17221 (best 0.17221), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=29-step=5670.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.021527  0.055773  0.063681  0.046164\n",
-      "ndcg    0.021527  0.082563  0.111508  0.058804\n",
-      "recall  0.021527  0.172214  0.286968  0.097698\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "b51365e160ab4865a27cf655e17707a5",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 30, global step 5859: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.018877  0.053691  0.061561  0.044171\n",
-      "ndcg    0.018877  0.080600  0.109718  0.057056\n",
-      "recall  0.018877  0.170558  0.286637  0.096705\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "8547295f7ccb4fd790366aab9fe9de78",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 31, global step 6048: 'recall@10' reached 0.17404 (best 0.17404), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=31-step=6048.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.014903  0.050917  0.058470  0.040520\n",
-      "ndcg    0.014903  0.079214  0.107110  0.053566\n",
-      "recall  0.014903  0.174035  0.285147  0.093724\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "71f02a600023460aa36238e118f89453",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 32, global step 6237: 'recall@10' reached 0.17553 (best 0.17553), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=32-step=6237.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.018049  0.053705  0.061464  0.043180\n",
-      "ndcg    0.018049  0.081750  0.110451  0.056098\n",
-      "recall  0.018049  0.175526  0.289949  0.095877\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "a94c40bad33b4fef92fbe3b7ca2dd7a9",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 33, global step 6426: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.019208  0.053874  0.061629  0.043597\n",
-      "ndcg    0.019208  0.081296  0.109911  0.056049\n",
-      "recall  0.019208  0.173207  0.287134  0.094386\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "93e62f688cae4954826ce5f65bc33ea6",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 34, global step 6615: 'recall@10' reached 0.18016 (best 0.18016), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=34-step=6615.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017553  0.054073  0.061606  0.043559\n",
-      "ndcg    0.017553  0.083073  0.111055  0.057211\n",
-      "recall  0.017553  0.180162  0.291936  0.099354\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "4bf8a5210e174568bd46212f9ea3fa7c",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 35, global step 6804: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.018215  0.055078  0.062641  0.044364\n",
-      "ndcg    0.018215  0.082983  0.110672  0.056984\n",
-      "recall  0.018215  0.176023  0.285809  0.095546\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "c43ebbf0977a431d8a961330f45570c5",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 36, global step 6993: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017221  0.052581  0.060386  0.042071\n",
-      "ndcg    0.017221  0.080544  0.109221  0.054559\n",
-      "recall  0.017221  0.174367  0.288293  0.092896\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "2d4241ed36f44607be6c6153931ac151",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 37, global step 7182: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017884  0.053871  0.061073  0.043766\n",
-      "ndcg    0.017884  0.082064  0.108563  0.057241\n",
-      "recall  0.017884  0.176354  0.281669  0.098857\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "aba0b7085a074368884b0651c8a2a0b3",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 38, global step 7371: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.018546  0.054212  0.061996  0.044508\n",
-      "ndcg    0.018546  0.081737  0.110419  0.057725\n",
-      "recall  0.018546  0.173704  0.287796  0.098361\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "feba795a523446eb8e5e6cd96bad39ba",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 39, global step 7560: 'recall@10' reached 0.18082 (best 0.18082), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=39-step=7560.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k             1        10        20         5\n",
-      "map     0.02103  0.057426  0.064922  0.046829\n",
-      "ndcg    0.02103  0.085795  0.113317  0.059729\n",
-      "recall  0.02103  0.180825  0.290114  0.099354\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "83bf44f90b024f2aa4f5b69bbc18b023",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 40, global step 7749: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.020699  0.055658  0.063230  0.044966\n",
-      "ndcg    0.020699  0.083080  0.111028  0.057106\n",
-      "recall  0.020699  0.174863  0.286140  0.094386\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "4abdd8633d3e453d93a2af9a1d6a9ccc",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 41, global step 7938: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016725  0.052379  0.060527  0.042126\n",
-      "ndcg    0.016725  0.079931  0.110088  0.054617\n",
-      "recall  0.016725  0.172214  0.292433  0.092896\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "b20c77f4d192498d81c2ed6baa4c811d",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 42, global step 8127: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.021692  0.058423  0.066561  0.048173\n",
-      "ndcg    0.021692  0.086242  0.116025  0.061053\n",
-      "recall  0.021692  0.179169  0.297235  0.100513\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "c4929ab9384a4a6ba9bf29a9c939b1ab",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 43, global step 8316: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.019705  0.055943  0.063840  0.045391\n",
-      "ndcg    0.019705  0.084424  0.113300  0.058431\n",
-      "recall  0.019705  0.179831  0.294254  0.098526\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "9ffbd31684ad4a6aa515233f9e8ef5aa",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 44, global step 8505: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.015069  0.053308  0.061374  0.043291\n",
-      "ndcg    0.015069  0.081830  0.111383  0.057192\n",
-      "recall  0.015069  0.176850  0.294088  0.099851\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "e6f025667c594ba0aa1029fc126858f2",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 45, global step 8694: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.018049  0.055185  0.062849  0.044538\n",
-      "ndcg    0.018049  0.083015  0.111301  0.057084\n",
-      "recall  0.018049  0.175857  0.288458  0.095380\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "ab51906b826a401396d92dd47c51028d",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 46, global step 8883: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017387  0.054915  0.062717  0.044329\n",
-      "ndcg    0.017387  0.083518  0.112327  0.057648\n",
-      "recall  0.017387  0.179003  0.293757  0.098526\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "e68d1d8754fa43958a9f7eb8b86ddfab",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 47, global step 9072: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017553  0.055809  0.064044  0.045537\n",
-      "ndcg    0.017553  0.084333  0.114317  0.059376\n",
-      "recall  0.017553  0.179169  0.297731  0.101838\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "dd0ec52c482f4fc789b0458eb0c2da5d",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 48, global step 9261: 'recall@10' reached 0.18596 (best 0.18596), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=48-step=9261.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.022355  0.060448  0.068135  0.049738\n",
-      "ndcg    0.022355  0.089385  0.117671  0.063264\n",
-      "recall  0.022355  0.185958  0.298394  0.104819\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "1f2fd19e798845c8959bd278a1e47c99",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 49, global step 9450: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.020533  0.058365  0.065847  0.047737\n",
-      "ndcg    0.020533  0.087662  0.115249  0.061533\n",
-      "recall  0.020533  0.185627  0.295413  0.103991\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "51c229a5a2ac4551a01b109eee0f72c9",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 50, global step 9639: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.019705  0.057252  0.065311  0.046550\n",
-      "ndcg    0.019705  0.086161  0.115841  0.059957\n",
-      "recall  0.019705  0.182812  0.300878  0.101176\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "c1bf323886a147aab5eb81a49910d99c",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 51, global step 9828: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.019871  0.056373  0.064782  0.046465\n",
-      "ndcg    0.019871  0.084487  0.115268  0.060018\n",
-      "recall  0.019871  0.178341  0.300381  0.101672\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "5a0a40c3297446219f782ccaddb48f36",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 52, global step 10017: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017056  0.055635  0.063238  0.044657\n",
-      "ndcg    0.017056  0.085148  0.113217  0.058444\n",
-      "recall  0.017056  0.183640  0.295413  0.100845\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "66e75dd4c62d419e8b41a635baeea9b6",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 53, global step 10206: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016725  0.056399  0.064312  0.046448\n",
-      "ndcg    0.016725  0.085609  0.114614  0.061160\n",
-      "recall  0.016725  0.182646  0.297731  0.106309\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "7b552a2b6ae343ab914f022626199df5",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 54, global step 10395: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.020864  0.056358  0.064559  0.045976\n",
-      "ndcg    0.020864  0.084551  0.114658  0.059012\n",
-      "recall  0.020864  0.179003  0.298559  0.099189\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "7d7d5e67c4264a8d936a886955a1e468",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 55, global step 10584: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.015897  0.052936  0.060713  0.042005\n",
-      "ndcg    0.015897  0.082446  0.111144  0.055636\n",
-      "recall  0.015897  0.181321  0.295579  0.097698\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "567aa7273f024f0496d7b105c0e5745f",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 56, global step 10773: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.018215  0.055970  0.064171  0.045659\n",
-      "ndcg    0.018215  0.084831  0.115129  0.059707\n",
-      "recall  0.018215  0.180990  0.301706  0.102997\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "effe5af26f79489f9fe7ced8ee84dbd6",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 57, global step 10962: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.019705  0.057038  0.065075  0.046796\n",
-      "ndcg    0.019705  0.085973  0.115791  0.060577\n",
-      "recall  0.019705  0.182812  0.301871  0.102997\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "5c3359ce59cd45e883e49fd6b07b642d",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 58, global step 11151: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.015731  0.054501  0.062389  0.043956\n",
-      "ndcg    0.015731  0.084609  0.113613  0.058502\n",
-      "recall  0.015731  0.185296  0.300546  0.103328\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "77112b36b0154d3cb233771cac236769",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 59, global step 11340: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.019705  0.057285  0.065306  0.046415\n",
-      "ndcg    0.019705  0.086774  0.116052  0.060147\n",
-      "recall  0.019705  0.185461  0.301374  0.102500\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "32f38bfe60b24b90b0a6d006a0babd31",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 60, global step 11529: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.018049  0.056947  0.064405  0.046498\n",
-      "ndcg    0.018049  0.086052  0.113448  0.060260\n",
-      "recall  0.018049  0.183143  0.291936  0.102335\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f094d09d3e7b453480b5b485f25692f3",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 61, global step 11718: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017553  0.057223  0.064880  0.046870\n",
-      "ndcg    0.017553  0.086840  0.115219  0.061361\n",
-      "recall  0.017553  0.185461  0.298725  0.105812\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f4948db841b44c168157b0b11e79efa4",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 62, global step 11907: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017718  0.056284  0.064414  0.045866\n",
-      "ndcg    0.017718  0.085691  0.115490  0.060021\n",
-      "recall  0.017718  0.183805  0.302037  0.103494\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "5fc0767f2e6a4b4ea45162db8c71ad1c",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 63, global step 12096: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.014075  0.053323  0.061450  0.042706\n",
-      "ndcg    0.014075  0.083044  0.113033  0.057016\n",
-      "recall  0.014075  0.182149  0.301540  0.101010\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "ffa03ec64fe842c2a11b1d4b2336c1b4",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 64, global step 12285: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016559  0.055950  0.063635  0.045606\n",
-      "ndcg    0.016559  0.085525  0.113860  0.060119\n",
-      "recall  0.016559  0.183971  0.296738  0.104653\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "1da052c1711a48eb9179718b1e8075cf",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 65, global step 12474: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017718  0.056836  0.065167  0.046782\n",
-      "ndcg    0.017718  0.085591  0.116301  0.060821\n",
-      "recall  0.017718  0.181321  0.303527  0.103825\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "b3993c1d275a4924b6731e4499bde5c6",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 66, global step 12663: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.015897  0.054146  0.063006  0.043884\n",
-      "ndcg    0.015897  0.083296  0.116044  0.058099\n",
-      "recall  0.015897  0.180493  0.310979  0.101838\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "832191caf90d44c7b8e8223afb823b76",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 67, global step 12852: 'recall@10' reached 0.19159 (best 0.19159), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=67-step=12852.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017056  0.057459  0.065661  0.045971\n",
-      "ndcg    0.017056  0.088389  0.118391  0.060327\n",
-      "recall  0.017056  0.191588  0.310482  0.104322\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "12bcca54c067420494528199338a1e56",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 68, global step 13041: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017718  0.058559  0.066859  0.048270\n",
-      "ndcg    0.017718  0.088819  0.119630  0.063497\n",
-      "recall  0.017718  0.189435  0.312469  0.110283\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "361a005333cf4228bd88446f5d3c7efd",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 69, global step 13230: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017387  0.057117  0.065430  0.046705\n",
-      "ndcg    0.017387  0.087054  0.117606  0.061578\n",
-      "recall  0.017387  0.186620  0.307998  0.107303\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "816b4845818f4e1ab83abd1cef17d94d",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 70, global step 13419: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.019043  0.057870  0.066302  0.047017\n",
-      "ndcg    0.019043  0.087867  0.118650  0.061382\n",
-      "recall  0.019043  0.187945  0.309820  0.105647\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "a7c54580cab949ed9162e5e2163f5144",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 71, global step 13608: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.015731  0.054928  0.063489  0.044635\n",
-      "ndcg    0.015731  0.084751  0.116133  0.059363\n",
-      "recall  0.015731  0.184136  0.308660  0.104653\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "cd4ae17c40d34d28a0d326139782a376",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 72, global step 13797: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.018381  0.056898  0.065084  0.046131\n",
-      "ndcg    0.018381  0.086764  0.116853  0.060242\n",
-      "recall  0.018381  0.186620  0.306177  0.103660\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "cb7d5821d4b24d928d834ac8f92c37b6",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 73, global step 13986: 'recall@10' reached 0.19308 (best 0.19308), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=73-step=13986.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k             1        10        20         5\n",
-      "map     0.01954  0.060165  0.068179  0.048888\n",
-      "ndcg    0.01954  0.090827  0.120453  0.063258\n",
-      "recall  0.01954  0.193078  0.311144  0.107303\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "8c82b678f2374af7859f13db93cc66a1",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 74, global step 14175: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.020202  0.059898  0.068066  0.049208\n",
-      "ndcg    0.020202  0.090335  0.120156  0.063797\n",
-      "recall  0.020202  0.192085  0.310151  0.108627\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "bff6682c9771473eba533531e1d63df0",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 75, global step 14364: 'recall@10' reached 0.19540 (best 0.19540), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=75-step=14364.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017553  0.059239  0.067168  0.047867\n",
-      "ndcg    0.017553  0.090614  0.119642  0.062577\n",
-      "recall  0.017553  0.195397  0.310482  0.107634\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "56513ed0de4f488e8a67ddab39bff555",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 76, global step 14553: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017553  0.057571  0.066031  0.046285\n",
-      "ndcg    0.017553  0.088607  0.119752  0.060844\n",
-      "recall  0.017553  0.192416  0.316278  0.105647\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "6c21be0f8bf64ff5a9a0433140db4afb",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 77, global step 14742: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k             1        10        20         5\n",
-      "map     0.02252  0.059148  0.067676  0.048570\n",
-      "ndcg    0.02252  0.087980  0.119156  0.062026\n",
-      "recall  0.02252  0.184468  0.307998  0.103494\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "67f774849e614a5e951e72ddc8847b0d",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 78, global step 14931: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.020864  0.059902  0.068102  0.048949\n",
-      "ndcg    0.020864  0.089627  0.119711  0.062921\n",
-      "recall  0.020864  0.188773  0.308164  0.105812\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "489fccb9e3774e5e9924cb0ae484229a",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 79, global step 15120: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.018049  0.057458  0.065510  0.046583\n",
-      "ndcg    0.018049  0.087580  0.117393  0.060762\n",
-      "recall  0.018049  0.188276  0.307170  0.104322\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "752a04f3c8ce4e36a5675c8a76c95503",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 80, global step 15309: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.018546  0.057481  0.065625  0.046603\n",
-      "ndcg    0.018546  0.087740  0.117695  0.061046\n",
-      "recall  0.018546  0.188773  0.307832  0.105481\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "9e682711164c445ba54cfe4a84262d8c",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 81, global step 15498: 'recall@10' reached 0.19672 (best 0.19672), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=81-step=15498.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.019705  0.060155  0.067919  0.048319\n",
-      "ndcg    0.019705  0.091545  0.120256  0.062412\n",
-      "recall  0.019705  0.196721  0.311144  0.105647\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f0da4eba218f4fffb4deb95efa8df791",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 82, global step 15687: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.018381  0.057389  0.065948  0.046868\n",
-      "ndcg    0.018381  0.087341  0.118833  0.061256\n",
-      "recall  0.018381  0.187448  0.312635  0.105481\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "1c626b061743416287ba6fc1aa65e8c6",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 83, global step 15876: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017387  0.057521  0.065933  0.046556\n",
-      "ndcg    0.017387  0.088367  0.119428  0.061458\n",
-      "recall  0.017387  0.191257  0.314953  0.107303\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "fd43c2909b574172a1726662394f32f7",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 84, global step 16065: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.019374  0.058125  0.066462  0.047210\n",
-      "ndcg    0.019374  0.087957  0.118482  0.061159\n",
-      "recall  0.019374  0.187614  0.308660  0.103991\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "33d3e2b831274e0b9783039482899584",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 85, global step 16254: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016725  0.055607  0.064210  0.044947\n",
-      "ndcg    0.016725  0.085456  0.117166  0.059262\n",
-      "recall  0.016725  0.185130  0.311310  0.103328\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "45f8fd0d5bd34318afd47a5b5b700d20",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 86, global step 16443: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016559  0.057457  0.065991  0.046661\n",
-      "ndcg    0.016559  0.087773  0.119155  0.061410\n",
-      "recall  0.016559  0.188607  0.313297  0.106640\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "bc750dde2e4d42eb8cf7655a39a465df",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 87, global step 16632: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.015731  0.057037  0.065072  0.045678\n",
-      "ndcg    0.015731  0.088516  0.118281  0.060867\n",
-      "recall  0.015731  0.193410  0.312138  0.107634\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "9b25c05436d24e3d80a56010edcff25b",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 88, global step 16821: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.018712  0.059290  0.067419  0.048380\n",
-      "ndcg    0.018712  0.090193  0.120127  0.063372\n",
-      "recall  0.018712  0.193244  0.312303  0.109455\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "89bed0fc3db14129badb8168c668c450",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 89, global step 17010: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.019208  0.060027  0.067680  0.048468\n",
-      "ndcg    0.019208  0.091291  0.119452  0.062846\n",
-      "recall  0.019208  0.195893  0.307832  0.106971\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "a66ad86b8b664ca6a4defd3c7cf97be1",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 90, global step 17199: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017553  0.059004  0.066841  0.047811\n",
-      "ndcg    0.017553  0.089908  0.118707  0.062456\n",
-      "recall  0.017553  0.192913  0.307336  0.107303\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "e6f035c7845b41f28da5d971ad4fc92f",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 91, global step 17388: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.019871  0.059618  0.068323  0.048620\n",
-      "ndcg    0.019871  0.089395  0.121484  0.062390\n",
-      "recall  0.019871  0.188773  0.316443  0.104488\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "782c0bc4a498460ca213a3930447fdc1",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 92, global step 17577: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016393  0.056838  0.065549  0.046489\n",
-      "ndcg    0.016393  0.087088  0.118991  0.061682\n",
-      "recall  0.016393  0.187779  0.314290  0.108462\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "448e00539d8c4ebcba2e51d0b5dbbb9e",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 93, global step 17766: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.020864  0.060046  0.068195  0.049431\n",
-      "ndcg    0.020864  0.090098  0.120281  0.063916\n",
-      "recall  0.020864  0.190429  0.310813  0.108462\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "2fea31da94254d8d8f7b87f88743a754",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 94, global step 17955: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.018215  0.059519  0.067797  0.048775\n",
-      "ndcg    0.018215  0.090521  0.120919  0.064276\n",
-      "recall  0.018215  0.193575  0.314290  0.111939\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "310abc5fa09f4f4597dba48212181900",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 95, global step 18144: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016393  0.057420  0.065561  0.045932\n",
-      "ndcg    0.016393  0.088958  0.118821  0.060803\n",
-      "recall  0.016393  0.194237  0.312800  0.106475\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "b3c83c7aa7c74abb8db94651d34e8a98",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 96, global step 18333: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016725  0.056762  0.065223  0.045708\n",
-      "ndcg    0.016725  0.087184  0.118141  0.060252\n",
-      "recall  0.016725  0.188607  0.311310  0.104984\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "54598b8e0f5b4bd3a2ed8a7c2ccbc8ae",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 97, global step 18522: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016228  0.056993  0.065326  0.045982\n",
-      "ndcg    0.016228  0.087804  0.118434  0.060850\n",
-      "recall  0.016228  0.190429  0.312138  0.106475\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "836e2a6ef8564b2bbaa9311a90e8817d",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 98, global step 18711: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016393  0.056134  0.064101  0.044497\n",
-      "ndcg    0.016393  0.087227  0.116437  0.058923\n",
-      "recall  0.016393  0.191091  0.307004  0.103328\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "e6338f1e7e1f409b9bb55d350f538a20",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 99, global step 18900: 'recall@10' was not in top 1\n",
-      "`Trainer.fit` stopped: `max_epochs=100` reached.\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017884  0.058572  0.066894  0.048223\n",
-      "ndcg    0.017884  0.088838  0.119484  0.063391\n",
-      "recall  0.017884  0.189435  0.311310  0.109952\n",
-      "\n"
-     ]
-    }
-   ],
-   "source": [
-    "from lightning.pytorch.callbacks import ModelCheckpoint\n",
-    "from lightning.pytorch.loggers import CSVLogger\n",
-    "\n",
-    "from replay.nn.lightning.callback import ComputeMetricsCallback\n",
-    "\n",
-    "\n",
-    "checkpoint_callback = ModelCheckpoint(\n",
-    "    dirpath=\"sasrec/checkpoints\",\n",
-    "    save_top_k=1,\n",
-    "    verbose=True,\n",
-    "    monitor=\"recall@10\",\n",
-    "    mode=\"max\",\n",
-    ")\n",
-    "\n",
-    "validation_metrics_callback = ComputeMetricsCallback(\n",
-    "    metrics=[\"map\", \"ndcg\", \"recall\"],\n",
-    "    ks=[1, 5, 10, 20],\n",
-    "    item_count=NUM_UNIQUE_ITEMS,\n",
-    ")\n",
-    "\n",
-    "csv_logger = CSVLogger(save_dir=\"sasrec/logs/train\", name=\"SasRec-example\")\n",
-    "\n",
-    "trainer = L.Trainer(\n",
-    "    max_epochs=100,\n",
-    "    callbacks=[checkpoint_callback, validation_metrics_callback],\n",
-    "    logger=csv_logger,\n",
-    ")\n",
-    "\n",
-    "trainer.fit(model, datamodule=parquet_module)"
-   ]
-  },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Now we can get the best model path stored in the checkpoint callback."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 27,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "'/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=81-step=15498.ckpt'"
-      ]
-     },
-     "execution_count": 27,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "best_model_path = checkpoint_callback.best_model_path\n",
-    "best_model_path"
-   ]
-  },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Inference\n",
-    "\n",
-    "To obtain model scores, we will load the weights from the best checkpoint. To do this, we use the `LightningModule`, provide there the path to the checkpoint and the model instance."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 33,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import replay\n",
-    "torch.serialization.add_safe_globals([\n",
-    "    replay.nn.lightning.optimizer.OptimizerFactory,\n",
-    "    replay.nn.lightning.scheduler.LRSchedulerFactory\n",
-    "])"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 34,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "sasrec = make_sasrec(\n",
-    "    schema=tensor_schema,\n",
-    "    embedding_dim=EMBEDDING_DIM,\n",
-    "    max_sequence_length=MAX_SEQ_LEN,\n",
-    "    num_heads=NUM_HEADS,\n",
-    "    num_blocks=NUM_BLOCKS,\n",
-    "    dropout=DROPOUT,\n",
-    ")\n",
-    "\n",
-    "best_model = LightningModule.load_from_checkpoint(best_model_path, model=sasrec)\n",
-    "best_model.eval();"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Configure `ParquetModule` for inference"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 44,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "inference_metadata = {\"predict\": create_meta(shape=MAX_SEQ_LEN)}\n",
-    "\n",
-    "parquet_module = ParquetModule(\n",
-    "    predict_path=PREDICT_PATH,\n",
-    "    batch_size=BATCH_SIZE,\n",
-    "    metadata=inference_metadata,\n",
-    "    transforms=transforms,\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "During inference, we can use `TopItemsCallback`. Such callback allows you to get scores for each user throughout the entire catalog and get recommendations in the form of ids of items with the highest score values.\n",
-    "\n",
-    "\n",
-    "Recommendations can be fetched in four formats: PySpark DataFrame, Pandas DataFrame, Polars DataFrame or raw PyTorch tensors. Each of the types corresponds a callback. In this example, we'll be using the `PandasTopItemsCallback`."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 46,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "💡 Tip: For seamless cloud uploads and versioning, try installing [litmodels](https://pypi.org/project/litmodels/) to enable LitModelCheckpoint, which syncs automatically with the Lightning model registry.\n",
-      "GPU available: True (cuda), used: True\n",
-      "TPU available: False, using: 0 TPU cores\n",
-      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "cd1373dad90640fe8cbc523662ccb563",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Predicting: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    }
-   ],
-   "source": [
-    "from replay.nn.lightning.callback import PandasTopItemsCallback\n",
-    "\n",
-    "csv_logger = CSVLogger(save_dir=\"sasrec/logs/test\", name=\"SasRec-example\")\n",
-    "\n",
-    "TOPK = [1, 5, 10, 20]\n",
-    "\n",
-    "pandas_prediction_callback = PandasTopItemsCallback(\n",
-    "    top_k=max(TOPK),\n",
-    "    query_column=\"user_id\",\n",
-    "    item_column=\"item_id\",\n",
-    "    rating_column=\"score\",\n",
-    ")\n",
-    "\n",
-    "trainer = L.Trainer(callbacks=[pandas_prediction_callback], logger=csv_logger, inference_mode=True)\n",
-    "trainer.predict(best_model, datamodule=parquet_module, return_predictions=False)\n",
-    "\n",
-    "pandas_res = pandas_prediction_callback.get_result()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 47,
-   "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>224</td>\n",
-       "      <td>7.242342</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>0</th>\n",
-       "      <td>0</td>\n",
-       "      <td>572</td>\n",
-       "      <td>6.818249</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>0</th>\n",
-       "      <td>0</td>\n",
-       "      <td>486</td>\n",
-       "      <td>6.81148</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>0</th>\n",
-       "      <td>0</td>\n",
-       "      <td>1371</td>\n",
-       "      <td>6.534966</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>0</th>\n",
-       "      <td>0</td>\n",
-       "      <td>210</td>\n",
-       "      <td>6.52649</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>...</th>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6037</th>\n",
-       "      <td>6039</td>\n",
-       "      <td>2497</td>\n",
-       "      <td>10.457304</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6037</th>\n",
-       "      <td>6039</td>\n",
-       "      <td>3503</td>\n",
-       "      <td>10.305973</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6037</th>\n",
-       "      <td>6039</td>\n",
-       "      <td>2601</td>\n",
-       "      <td>10.280416</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6037</th>\n",
-       "      <td>6039</td>\n",
-       "      <td>2750</td>\n",
-       "      <td>10.01198</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6037</th>\n",
-       "      <td>6039</td>\n",
-       "      <td>1490</td>\n",
-       "      <td>9.916577</td>\n",
-       "    </tr>\n",
-       "  </tbody>\n",
-       "</table>\n",
-       "<p>120760 rows × 3 columns</p>\n",
-       "</div>"
-      ],
-      "text/plain": [
-       "      user_id item_id      score\n",
-       "0           0     224   7.242342\n",
-       "0           0     572   6.818249\n",
-       "0           0     486    6.81148\n",
-       "0           0    1371   6.534966\n",
-       "0           0     210    6.52649\n",
-       "...       ...     ...        ...\n",
-       "6037     6039    2497  10.457304\n",
-       "6037     6039    3503  10.305973\n",
-       "6037     6039    2601  10.280416\n",
-       "6037     6039    2750   10.01198\n",
-       "6037     6039    1490   9.916577\n",
-       "\n",
-       "[120760 rows x 3 columns]"
-      ]
-     },
-     "execution_count": 47,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "pandas_res"
-   ]
-  },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### Calculating metrics\n",
-    "\n",
-    "*test_gt* is already encoded, so we can use it for computing metrics."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 48,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from replay.metrics import MAP, OfflineMetrics, Precision, Recall\n",
-    "from replay.metrics.torch_metrics_builder import metrics_to_df"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 49,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "result_metrics = OfflineMetrics(\n",
-    "    [Recall(TOPK), Precision(TOPK), MAP(TOPK)],\n",
-    "    query_column=\"user_id\",\n",
-    "    rating_column=\"score\",\n",
-    ")(pandas_res, test_gt.explode(\"item_id\"))"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 50,
-   "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>5</th>\n",
-       "    </tr>\n",
-       "  </thead>\n",
-       "  <tbody>\n",
-       "    <tr>\n",
-       "      <th>MAP</th>\n",
-       "      <td>0.016065</td>\n",
-       "      <td>0.054039</td>\n",
-       "      <td>0.061749</td>\n",
-       "      <td>0.043969</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>Precision</th>\n",
-       "      <td>0.016065</td>\n",
-       "      <td>0.017655</td>\n",
-       "      <td>0.014458</td>\n",
-       "      <td>0.020073</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>Recall</th>\n",
-       "      <td>0.016065</td>\n",
-       "      <td>0.176549</td>\n",
-       "      <td>0.289169</td>\n",
-       "      <td>0.100364</td>\n",
-       "    </tr>\n",
-       "  </tbody>\n",
-       "</table>\n",
-       "</div>"
-      ],
-      "text/plain": [
-       "k                 1        10        20         5\n",
-       "MAP        0.016065  0.054039  0.061749  0.043969\n",
-       "Precision  0.016065  0.017655  0.014458  0.020073\n",
-       "Recall     0.016065  0.176549  0.289169  0.100364"
-      ]
-     },
-     "execution_count": 50,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "metrics_to_df(result_metrics)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Let's call the `inverse_transform` encoder's function to get the final dataframe with recommendations"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 51,
-   "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>2012</td>\n",
-       "      <td>231</td>\n",
-       "      <td>7.242342</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>0</th>\n",
-       "      <td>2012</td>\n",
-       "      <td>586</td>\n",
-       "      <td>6.818249</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>0</th>\n",
-       "      <td>2012</td>\n",
-       "      <td>500</td>\n",
-       "      <td>6.81148</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>0</th>\n",
-       "      <td>2012</td>\n",
-       "      <td>1485</td>\n",
-       "      <td>6.534966</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>0</th>\n",
-       "      <td>2012</td>\n",
-       "      <td>216</td>\n",
-       "      <td>6.52649</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>...</th>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6037</th>\n",
-       "      <td>5727</td>\n",
-       "      <td>2702</td>\n",
-       "      <td>10.457304</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6037</th>\n",
-       "      <td>5727</td>\n",
-       "      <td>3745</td>\n",
-       "      <td>10.305973</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6037</th>\n",
-       "      <td>5727</td>\n",
-       "      <td>2806</td>\n",
-       "      <td>10.280416</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6037</th>\n",
-       "      <td>5727</td>\n",
-       "      <td>2961</td>\n",
-       "      <td>10.01198</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6037</th>\n",
-       "      <td>5727</td>\n",
-       "      <td>1623</td>\n",
-       "      <td>9.916577</td>\n",
-       "    </tr>\n",
-       "  </tbody>\n",
-       "</table>\n",
-       "<p>120760 rows × 3 columns</p>\n",
-       "</div>"
-      ],
-      "text/plain": [
-       "      user_id  item_id      score\n",
-       "0        2012      231   7.242342\n",
-       "0        2012      586   6.818249\n",
-       "0        2012      500    6.81148\n",
-       "0        2012     1485   6.534966\n",
-       "0        2012      216    6.52649\n",
-       "...       ...      ...        ...\n",
-       "6037     5727     2702  10.457304\n",
-       "6037     5727     3745  10.305973\n",
-       "6037     5727     2806  10.280416\n",
-       "6037     5727     2961   10.01198\n",
-       "6037     5727     1623   9.916577\n",
-       "\n",
-       "[120760 rows x 3 columns]"
-      ]
-     },
-     "execution_count": 51,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "encoder.inverse_transform(pandas_res)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  }
- ],
- "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.11.14"
-  },
-  "orig_nbformat": 4
- },
- "nbformat": 4,
- "nbformat_minor": 2
-}
diff --git a/examples/09_sasrec_example_thr_sampled.ipynb b/examples/09_sasrec_example_thr_sampled.ipynb
deleted file mode 100644
index b27ec3f8d..000000000
--- a/examples/09_sasrec_example_thr_sampled.ipynb
+++ /dev/null
@@ -1,4768 +0,0 @@
-{
- "cells": [
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Example of SasRec training/inference with Parquet Module"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 1,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Seed set to 42\n"
-     ]
-    }
-   ],
-   "source": [
-    "from typing import Optional\n",
-    "\n",
-    "import lightning as L\n",
-    "import pandas as pd\n",
-    "\n",
-    "L.seed_everything(42)\n",
-    "\n",
-    "import warnings\n",
-    "warnings.filterwarnings(\"ignore\")"
-   ]
-  },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Preparing data\n",
-    "In this example, we will be using the MovieLens dataset, namely the 1m subset. It's demonstrated a simple case, so only item ids will be used as model input.\n",
-    "\n",
-    "---\n",
-    "**NOTE**\n",
-    "\n",
-    "Current implementation of SasRec is able to handle item and interactions features. It does not take into account user features. \n",
-    "\n",
-    "---"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 2,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "interactions = pd.read_csv(\"./data/ml1m_ratings.dat\", sep=\"\\t\", names=[\"user_id\", \"item_id\",\"rating\",\"timestamp\"])\n",
-    "interactions = interactions.drop(columns=[\"rating\"])"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 3,
-   "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>timestamp</th>\n",
-       "    </tr>\n",
-       "  </thead>\n",
-       "  <tbody>\n",
-       "    <tr>\n",
-       "      <th>1000138</th>\n",
-       "      <td>6040</td>\n",
-       "      <td>858</td>\n",
-       "      <td>0</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1000153</th>\n",
-       "      <td>6040</td>\n",
-       "      <td>2384</td>\n",
-       "      <td>1</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>999873</th>\n",
-       "      <td>6040</td>\n",
-       "      <td>593</td>\n",
-       "      <td>2</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1000007</th>\n",
-       "      <td>6040</td>\n",
-       "      <td>1961</td>\n",
-       "      <td>3</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1000192</th>\n",
-       "      <td>6040</td>\n",
-       "      <td>2019</td>\n",
-       "      <td>4</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>...</th>\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>446</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>825438</th>\n",
-       "      <td>4958</td>\n",
-       "      <td>1407</td>\n",
-       "      <td>447</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>825724</th>\n",
-       "      <td>4958</td>\n",
-       "      <td>3264</td>\n",
-       "      <td>448</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>825731</th>\n",
-       "      <td>4958</td>\n",
-       "      <td>2634</td>\n",
-       "      <td>449</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>825603</th>\n",
-       "      <td>4958</td>\n",
-       "      <td>1924</td>\n",
-       "      <td>450</td>\n",
-       "    </tr>\n",
-       "  </tbody>\n",
-       "</table>\n",
-       "<p>1000209 rows × 3 columns</p>\n",
-       "</div>"
-      ],
-      "text/plain": [
-       "         user_id  item_id  timestamp\n",
-       "1000138     6040      858          0\n",
-       "1000153     6040     2384          1\n",
-       "999873      6040      593          2\n",
-       "1000007     6040     1961          3\n",
-       "1000192     6040     2019          4\n",
-       "...          ...      ...        ...\n",
-       "825793      4958     2399        446\n",
-       "825438      4958     1407        447\n",
-       "825724      4958     3264        448\n",
-       "825731      4958     2634        449\n",
-       "825603      4958     1924        450\n",
-       "\n",
-       "[1000209 rows x 3 columns]"
-      ]
-     },
-     "execution_count": 3,
-     "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"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### Encode catagorical data.\n",
-    "To ensure all categorical data is fit for training, it needs to be encoded using the `LabelEncoder` class. Create an instance of the encoder, providing a `LabelEncodingRule` for each categorcial column in the dataset that will be used in model. Note that ids of users and ids of items are always used."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 4,
-   "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>timestamp</th>\n",
-       "      <th>user_id</th>\n",
-       "      <th>item_id</th>\n",
-       "    </tr>\n",
-       "  </thead>\n",
-       "  <tbody>\n",
-       "    <tr>\n",
-       "      <th>0</th>\n",
-       "      <td>12</td>\n",
-       "      <td>0</td>\n",
-       "      <td>0</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1</th>\n",
-       "      <td>68</td>\n",
-       "      <td>1</td>\n",
-       "      <td>0</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>2</th>\n",
-       "      <td>67</td>\n",
-       "      <td>2</td>\n",
-       "      <td>0</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>3</th>\n",
-       "      <td>12</td>\n",
-       "      <td>3</td>\n",
-       "      <td>0</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>4</th>\n",
-       "      <td>140</td>\n",
-       "      <td>4</td>\n",
-       "      <td>0</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>...</th>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1000204</th>\n",
-       "      <td>14</td>\n",
-       "      <td>4555</td>\n",
-       "      <td>3705</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1000205</th>\n",
-       "      <td>90</td>\n",
-       "      <td>2813</td>\n",
-       "      <td>3705</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1000206</th>\n",
-       "      <td>70</td>\n",
-       "      <td>2404</td>\n",
-       "      <td>3705</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1000207</th>\n",
-       "      <td>25</td>\n",
-       "      <td>5835</td>\n",
-       "      <td>3705</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1000208</th>\n",
-       "      <td>380</td>\n",
-       "      <td>979</td>\n",
-       "      <td>3705</td>\n",
-       "    </tr>\n",
-       "  </tbody>\n",
-       "</table>\n",
-       "<p>1000209 rows × 3 columns</p>\n",
-       "</div>"
-      ],
-      "text/plain": [
-       "         timestamp  user_id  item_id\n",
-       "0               12        0        0\n",
-       "1               68        1        0\n",
-       "2               67        2        0\n",
-       "3               12        3        0\n",
-       "4              140        4        0\n",
-       "...            ...      ...      ...\n",
-       "1000204         14     4555     3705\n",
-       "1000205         90     2813     3705\n",
-       "1000206         70     2404     3705\n",
-       "1000207         25     5835     3705\n",
-       "1000208        380      979     3705\n",
-       "\n",
-       "[1000209 rows x 3 columns]"
-      ]
-     },
-     "execution_count": 4,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "from replay.preprocessing.label_encoder import LabelEncoder, LabelEncodingRule\n",
-    "\n",
-    "encoder = LabelEncoder(\n",
-    "    [\n",
-    "        LabelEncodingRule(\"user_id\", default_value=\"last\"),\n",
-    "        LabelEncodingRule(\"item_id\", default_value=\"last\"),\n",
-    "    ]\n",
-    ")\n",
-    "interactions = interactions.sort_values(by=\"item_id\", ascending=True)\n",
-    "encoded_interactions = encoder.fit_transform(interactions)\n",
-    "encoded_interactions"
-   ]
-  },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### Split interactions into the train, validation and test datasets using LastNSplitter\n",
-    "We use widespread splitting strategy Last-One-Out. We filter out cold items and users for simplicity."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 5,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from replay.splitters import LastNSplitter\n",
-    "\n",
-    "splitter = LastNSplitter(\n",
-    "    N=1,\n",
-    "    divide_column=\"user_id\",\n",
-    "    query_column=\"user_id\",\n",
-    "    strategy=\"interactions\",\n",
-    "    drop_cold_users=True,\n",
-    "    drop_cold_items=True\n",
-    ")\n",
-    "\n",
-    "test_events, test_gt = splitter.split(encoded_interactions)\n",
-    "validation_events, validation_gt = splitter.split(test_events)\n",
-    "train_events = validation_events"
-   ]
-  },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### Dataset preprocessing (\"baking\")\n",
-    "SasRec expects each user in the batch to provide their events in form of a sequence. For this reason, the event splits must be properly processed using the `groupby_sequences` function provided by RePlay."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 6,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from replay.data.nn.utils import groupby_sequences\n",
-    "\n",
-    "\n",
-    "def bake_data(full_data):\n",
-    "    grouped_interactions = groupby_sequences(events=full_data, groupby_col=\"user_id\", sort_col=\"timestamp\")\n",
-    "    return grouped_interactions"
-   ]
-  },
-  {
-   "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>timestamp</th>\n",
-       "      <th>item_id</th>\n",
-       "    </tr>\n",
-       "  </thead>\n",
-       "  <tbody>\n",
-       "    <tr>\n",
-       "      <th>0</th>\n",
-       "      <td>0</td>\n",
-       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
-       "      <td>[2426, 822, 2733, 2587, 2937, 3618, 2943, 708,...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>1</th>\n",
-       "      <td>1</td>\n",
-       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
-       "      <td>[3272, 3026, 2760, 851, 346, 3393, 1107, 515, ...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>2</th>\n",
-       "      <td>2</td>\n",
-       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
-       "      <td>[579, 1140, 1154, 2426, 1524, 1260, 2160, 2621...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>3</th>\n",
-       "      <td>3</td>\n",
-       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
-       "      <td>[1781, 2940, 2468, 890, 948, 106, 593, 309, 49...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>4</th>\n",
-       "      <td>4</td>\n",
-       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
-       "      <td>[1108, 2229, 21, 2435, 2142, 106, 1167, 593, 1...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>...</th>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6035</th>\n",
-       "      <td>6035</td>\n",
-       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
-       "      <td>[2426, 1279, 3151, 3321, 1178, 3301, 2501, 278...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6036</th>\n",
-       "      <td>6036</td>\n",
-       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
-       "      <td>[1592, 2302, 1633, 1813, 2879, 1482, 2651, 250...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6037</th>\n",
-       "      <td>6037</td>\n",
-       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
-       "      <td>[1971, 3500, 2077, 1666, 1399, 2651, 2748, 283...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6038</th>\n",
-       "      <td>6038</td>\n",
-       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
-       "      <td>[1486, 1485, 3384, 3512, 3302, 3126, 3650, 330...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6039</th>\n",
-       "      <td>6039</td>\n",
-       "      <td>[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...</td>\n",
-       "      <td>[2432, 2960, 1848, 2114, 2142, 3091, 3248, 317...</td>\n",
-       "    </tr>\n",
-       "  </tbody>\n",
-       "</table>\n",
-       "<p>6040 rows × 3 columns</p>\n",
-       "</div>"
-      ],
-      "text/plain": [
-       "      user_id                                          timestamp  \\\n",
-       "0           0  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
-       "1           1  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
-       "2           2  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
-       "3           3  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
-       "4           4  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
-       "...       ...                                                ...   \n",
-       "6035     6035  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
-       "6036     6036  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
-       "6037     6037  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
-       "6038     6038  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
-       "6039     6039  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13,...   \n",
-       "\n",
-       "                                                item_id  \n",
-       "0     [2426, 822, 2733, 2587, 2937, 3618, 2943, 708,...  \n",
-       "1     [3272, 3026, 2760, 851, 346, 3393, 1107, 515, ...  \n",
-       "2     [579, 1140, 1154, 2426, 1524, 1260, 2160, 2621...  \n",
-       "3     [1781, 2940, 2468, 890, 948, 106, 593, 309, 49...  \n",
-       "4     [1108, 2229, 21, 2435, 2142, 106, 1167, 593, 1...  \n",
-       "...                                                 ...  \n",
-       "6035  [2426, 1279, 3151, 3321, 1178, 3301, 2501, 278...  \n",
-       "6036  [1592, 2302, 1633, 1813, 2879, 1482, 2651, 250...  \n",
-       "6037  [1971, 3500, 2077, 1666, 1399, 2651, 2748, 283...  \n",
-       "6038  [1486, 1485, 3384, 3512, 3302, 3126, 3650, 330...  \n",
-       "6039  [2432, 2960, 1848, 2114, 2142, 3091, 3248, 317...  \n",
-       "\n",
-       "[6040 rows x 3 columns]"
-      ]
-     },
-     "execution_count": 7,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "train_events = bake_data(train_events)\n",
-    "\n",
-    "validation_events = bake_data(validation_events)\n",
-    "validation_gt = bake_data(validation_gt)\n",
-    "\n",
-    "test_events = bake_data(test_events)\n",
-    "test_gt = bake_data(test_gt)\n",
-    "\n",
-    "train_events"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "To ensure we don't have unknown users in ground truth, we join validation events and validation ground truth (also join test events and test ground truth correspondingly) by user ids to leave only the common ones.  "
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 8,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def add_gt_to_events(events_df, gt_df):\n",
-    "    gt_to_join = gt_df[[\"user_id\", \"item_id\"]].rename(columns={\"item_id\": \"ground_truth\"})\n",
-    "\n",
-    "    events_df = events_df.merge(gt_to_join, on=\"user_id\", how=\"inner\")\n",
-    "    return events_df\n",
-    "\n",
-    "validation_events = add_gt_to_events(validation_events, validation_gt)\n",
-    "test_events = add_gt_to_events(test_events, test_gt)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 9,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from pathlib import Path\n",
-    "\n",
-    "data_dir = Path(\"temp/data/\")\n",
-    "data_dir.mkdir(parents=True, exist_ok=True)\n",
-    "\n",
-    "TRAIN_PATH = data_dir / \"train.parquet\"\n",
-    "VAL_PATH = data_dir / \"val.parquet\"\n",
-    "PREDICT_PATH = data_dir / \"test.parquet\"\n",
-    "\n",
-    "ENCODER_PATH = data_dir / \"encoder\""
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 10,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "train_events.to_parquet(TRAIN_PATH)\n",
-    "validation_events.to_parquet(VAL_PATH)\n",
-    "test_events.to_parquet(PREDICT_PATH)\n",
-    "\n",
-    "encoder.save(ENCODER_PATH)"
-   ]
-  },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Prepare to model training\n",
-    "### 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. It is required by the SasRec model to correctly create embeddings for every source column. Note that user_id does not required in `TensorSchema`.\n",
-    "\n",
-    "Note that the **padding value** is the next value (item_id) after the last one. **Cardinality** is the number of unique values ​​given the padding value."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 11,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from replay.data import FeatureHint, FeatureType\n",
-    "from replay.data.nn import TensorFeatureInfo, TensorSchema\n",
-    "\n",
-    "\n",
-    "EMBEDDING_DIM = 64\n",
-    "\n",
-    "encoder = encoder.load(ENCODER_PATH)\n",
-    "NUM_UNIQUE_ITEMS = len(encoder.mapping[\"item_id\"])\n",
-    "\n",
-    "tensor_schema = TensorSchema(\n",
-    "    [\n",
-    "        TensorFeatureInfo(\n",
-    "            name=\"item_id\",\n",
-    "            is_seq=True,\n",
-    "            padding_value=NUM_UNIQUE_ITEMS,\n",
-    "            cardinality=NUM_UNIQUE_ITEMS + 1,  # taking into account padding\n",
-    "            embedding_dim=EMBEDDING_DIM,\n",
-    "            feature_type=FeatureType.CATEGORICAL,\n",
-    "            feature_hint=FeatureHint.ITEM_ID,\n",
-    "        )\n",
-    "    ]\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### Configure ParquetModule and transformation pipelines\n",
-    "\n",
-    "The `ParquetModule` class enables training of models on large datasets by reading data in batch-wise way. This class initialized with **paths to every data split, a metadata dict containing information about shape and padding value of every column and a dict of transforms**. `ParquetModule`'s  \"transform pipelines\" are stage-specific modules implementing additional preprocessing to be performed on batch level right before the forward pass.  \n",
-    "\n",
-    "For SasRec model, RePlay provides a function that generates a sequence of appropriate transforms for each data split named **make_default_sasrec_transforms**.\n",
-    "\n",
-    "Internally this function creates the following transforms:\n",
-    "1) Training:\n",
-    "    1. Create a target, which contains the shifted item sequence that represents the next item in the sequence (for the next item prediction task).\n",
-    "    2. Rename features to match it with expected format by the model during training.\n",
-    "    3. Unsqueeze target (*positive_labels*) and it's padding mask (*target_padding_mask*) for getting required shape of this tensors for loss computation.\n",
-    "    4. Group input features to be embed in expected format.\n",
-    "\n",
-    "2) Validation/Inference:\n",
-    "    1. Rename/group features to match it with expected format by the model during valdiation/inference.\n",
-    "\n",
-    "If a different set of transforms is required, you can create them yourself and submit them to the ParquetModule in the form of a dictionary where the key is the name of the split, and the value is the list of transforms. Available transforms are in the replay/nn/transforms/.\n",
-    "\n",
-    "**Note:** One of the transforms for the training data prepares the initial sequence for the task of Next Item Prediction so it shifts the sequence of items. For the final sequence length to be correct, you need to set shape of item_id in metadata as **model sequence length + shift**. Default shift value is 1."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 12,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import copy\n",
-    "\n",
-    "import torch\n",
-    "\n",
-    "from replay.data.nn import TensorSchema\n",
-    "from replay.nn.transform import GroupTransform, NextTokenTransform, RenameTransform, UnsqueezeTransform, ThresholdNegativeSamplingTransform"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 13,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def make_sasrec_transforms(\n",
-    "    tensor_schema: TensorSchema, query_column: str = \"query_id\", num_negative_samples: int = 128,\n",
-    ") -> dict[str, list[torch.nn.Module]]:\n",
-    "    item_column = tensor_schema.item_id_feature_name\n",
-    "    vocab_size = tensor_schema[item_column].cardinality\n",
-    "    train_transforms = [\n",
-    "        ThresholdNegativeSamplingTransform(vocab_size, num_negative_samples),\n",
-    "        NextTokenTransform(label_field=item_column, query_features=query_column, shift=1),\n",
-    "        RenameTransform(\n",
-    "            {\n",
-    "                query_column: \"query_id\",\n",
-    "                f\"{item_column}_mask\": \"padding_mask\",\n",
-    "                \"positive_labels_mask\": \"target_padding_mask\",\n",
-    "            }\n",
-    "        ),\n",
-    "        UnsqueezeTransform(\"target_padding_mask\", -1),\n",
-    "        UnsqueezeTransform(\"positive_labels\", -1),\n",
-    "        GroupTransform({\"feature_tensors\": [item_column]}),\n",
-    "    ]\n",
-    "\n",
-    "    val_transforms = [\n",
-    "        RenameTransform({query_column: \"query_id\", f\"{item_column}_mask\": \"padding_mask\"}),\n",
-    "        GroupTransform({\"feature_tensors\": [item_column]}),\n",
-    "    ]\n",
-    "    test_transforms = copy.deepcopy(val_transforms)\n",
-    "\n",
-    "    predict_transforms = copy.deepcopy(val_transforms)\n",
-    "\n",
-    "    transforms = {\n",
-    "        \"train\": train_transforms,\n",
-    "        \"validate\": val_transforms,\n",
-    "        \"test\": test_transforms,\n",
-    "        \"predict\": predict_transforms,\n",
-    "    }\n",
-    "\n",
-    "    return transforms\n",
-    "\n",
-    "transforms = make_sasrec_transforms(tensor_schema, query_column=\"user_id\")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 14,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "MAX_SEQ_LEN = 50\n",
-    "\n",
-    "def create_meta(shape: int, gt_shape: Optional[int] = None):\n",
-    "    meta = {\n",
-    "        \"user_id\": {},\n",
-    "        \"item_id\": {\"shape\": shape, \"padding\": tensor_schema[\"item_id\"].padding_value},\n",
-    "    }\n",
-    "    if gt_shape is not None:\n",
-    "        meta.update({\"ground_truth\": {\"shape\": gt_shape, \"padding\": -1}})\n",
-    "\n",
-    "    return meta\n",
-    "\n",
-    "train_metadata = {\n",
-    "    \"train\": create_meta(shape=MAX_SEQ_LEN+1),\n",
-    "    \"validate\": create_meta(shape=MAX_SEQ_LEN, gt_shape=1),\n",
-    "}"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 15,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from replay.data.nn import ParquetModule\n",
-    "\n",
-    "BATCH_SIZE = 32\n",
-    "\n",
-    "parquet_module = ParquetModule(\n",
-    "    train_path=TRAIN_PATH,\n",
-    "    validate_path=VAL_PATH,\n",
-    "    batch_size=BATCH_SIZE,\n",
-    "    metadata=train_metadata,\n",
-    "    transforms=transforms,\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Train model\n",
-    "### Create SasRec model instance and run the training stage using lightning\n",
-    "We may now train the model using the Lightning trainer class. \n",
-    "\n",
-    "RePlay's implementation of SasRec is designed in a modular, **block-based approach**. Instead of passing configuration parameters to the constructor, SasRec is now built by providing fully initialized components that makes the model more flexible and easier to extend.\n",
-    "\n",
-    "#### Default Configuration\n",
-    "\n",
-    "Default SasRec model may be created quickly via method **from_params**. Default model instance has CE loss, original SasRec transformer layes, and embeddings are aggregated via sum."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 16,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from replay.nn.sequential import SasRec\n",
-    "from typing import Literal\n",
-    "def make_sasrec(\n",
-    "    schema: TensorSchema,\n",
-    "    embedding_dim: int = 192,\n",
-    "    num_heads: int = 4,\n",
-    "    num_blocks: int = 2,\n",
-    "    max_sequence_length: int = 50,\n",
-    "    dropout: float = 0.3,\n",
-    "    excluded_features: Optional[list[str]] = None,\n",
-    "    categorical_list_feature_aggregation_method: Literal[\"sum\", \"mean\", \"max\"] = \"sum\",\n",
-    ") -> SasRec:\n",
-    "    from replay.nn.sequential.sasrec import SasRecBody, SasRecTransformerLayer\n",
-    "    from replay.nn.agg import SumAggregator\n",
-    "    from replay.nn.embedding import SequenceEmbedding\n",
-    "    from replay.nn.loss import CE, CESampled\n",
-    "    from replay.nn.mask import DefaultAttentionMask\n",
-    "    from replay.nn.sequential.sasrec.agg import PositionAwareAggregator\n",
-    "    from replay.nn.sequential.sasrec.transformer import SasRecTransformerLayer\n",
-    "    excluded_features = [\n",
-    "        schema.query_id_feature_name,\n",
-    "        schema.timestamp_feature_name,\n",
-    "        *(excluded_features or []),\n",
-    "    ]\n",
-    "    excluded_features = list(set(excluded_features))\n",
-    "    body = SasRecBody(\n",
-    "        embedder=SequenceEmbedding(\n",
-    "            schema=schema,\n",
-    "            categorical_list_feature_aggregation_method=categorical_list_feature_aggregation_method,\n",
-    "            excluded_features=excluded_features,\n",
-    "        ),\n",
-    "        embedding_aggregator=PositionAwareAggregator(\n",
-    "            embedding_aggregator=SumAggregator(embedding_dim=embedding_dim),\n",
-    "            max_sequence_length=max_sequence_length,\n",
-    "            dropout=dropout,\n",
-    "        ),\n",
-    "        attn_mask_builder=DefaultAttentionMask(\n",
-    "            reference_feature_name=schema.item_id_feature_name,\n",
-    "            num_heads=num_heads,\n",
-    "        ),\n",
-    "        encoder=SasRecTransformerLayer(\n",
-    "            embedding_dim=embedding_dim,\n",
-    "            num_heads=num_heads,\n",
-    "            num_blocks=num_blocks,\n",
-    "            dropout=dropout,\n",
-    "            activation=\"relu\",\n",
-    "        ),\n",
-    "        output_normalization=torch.nn.LayerNorm(embedding_dim),\n",
-    "    )\n",
-    "    padding_idx = schema.item_id_features.item().padding_value\n",
-    "    return SasRec(\n",
-    "        body=body,\n",
-    "        loss=CESampled(padding_idx=padding_idx),\n",
-    "    )"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 17,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "NUM_BLOCKS = 2\n",
-    "NUM_HEADS = 2\n",
-    "DROPOUT = 0.3\n",
-    "\n",
-    "sasrec = make_sasrec(\n",
-    "    schema=tensor_schema,\n",
-    "    embedding_dim=EMBEDDING_DIM,\n",
-    "    max_sequence_length=MAX_SEQ_LEN,\n",
-    "    num_heads=NUM_HEADS,\n",
-    "    num_blocks=NUM_BLOCKS,\n",
-    "    dropout=DROPOUT,\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "A universal PyTorch Lightning module is provided. It can work with any NN model."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 18,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from replay.nn.lightning.optimizer import OptimizerFactory\n",
-    "from replay.nn.lightning.scheduler import LRSchedulerFactory\n",
-    "from replay.nn.lightning import LightningModule\n",
-    "\n",
-    "model = LightningModule(\n",
-    "    sasrec,\n",
-    "    optimizer_factory=OptimizerFactory(),\n",
-    "    lr_scheduler_factory=LRSchedulerFactory(),\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "To facilitate training, we add the following callbacks:\n",
-    "1) `ModelCheckpoint` - to save the best trained model based on its Recall metric. It's a default Lightning Callback.\n",
-    "1) `ComputeMetricsCallback` - to display a detailed validation metric matrix after each epoch. It's a custom RePlay callback for computing recsys metrics on validation and test stages.\n"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 19,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "GPU available: True (cuda), used: True\n",
-      "TPU available: False, using: 0 TPU cores\n",
-      "You are using a CUDA device ('NVIDIA GeForce RTX 3060 Laptop GPU') 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 | Mode  | FLOPs\n",
-      "-------------------------------------------------\n",
-      "0 | model | SasRec | 291 K  | train | 0    \n",
-      "-------------------------------------------------\n",
-      "291 K     Trainable params\n",
-      "0         Non-trainable params\n",
-      "291 K     Total params\n",
-      "1.164     Total estimated model params size (MB)\n",
-      "39        Modules in train mode\n",
-      "0         Modules in eval mode\n",
-      "0         Total Flops\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "21ca5ea332d9418fb2481e0296b399a0",
-       "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": "dfbb5a69743640d282090979774e93a2",
-       "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": "b6b1a774cbc84c73b23e546285b0a116",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 0, global step 189: 'recall@10' reached 0.02517 (best 0.02517), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=0-step=189-v3.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.003477  0.008391  0.009940  0.006889\n",
-      "ndcg    0.003477  0.012236  0.018080  0.008533\n",
-      "recall  0.003477  0.025170  0.048684  0.013578\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f47ec635a7c143f281864d439cbbcc8e",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 1, global step 378: 'recall@10' reached 0.04090 (best 0.04090), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=1-step=378-v2.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.005133  0.013343  0.015591  0.011023\n",
-      "ndcg    0.005133  0.019693  0.028004  0.014007\n",
-      "recall  0.005133  0.040901  0.074019  0.023183\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "bbdfceaef5df4df98b29883914a7c174",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 2, global step 567: 'recall@10' reached 0.06773 (best 0.06773), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=2-step=567.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.008776  0.021569  0.024780  0.017903\n",
-      "ndcg    0.008776  0.032154  0.044198  0.022983\n",
-      "recall  0.008776  0.067726  0.116079  0.038748\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "7f43be0bcd8c49c7b29e3689b24ad818",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 3, global step 756: 'recall@10' reached 0.08197 (best 0.08197), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=3-step=756-v3.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k             1        10        20         5\n",
-      "map     0.00977  0.024920  0.029111  0.019702\n",
-      "ndcg    0.00977  0.037940  0.053492  0.025112\n",
-      "recall  0.00977  0.081967  0.144064  0.041894\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "0df4dc7b7f964a4089205386735e2968",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 4, global step 945: 'recall@10' reached 0.10051 (best 0.10051), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=4-step=945-v4.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.011757  0.032053  0.036894  0.025981\n",
-      "ndcg    0.011757  0.047823  0.065510  0.033102\n",
-      "recall  0.011757  0.100513  0.170558  0.054976\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "92b4eac8339b483684bc947c3e9f035b",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 5, global step 1134: 'recall@10' reached 0.11707 (best 0.11707), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=5-step=1134-v1.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k            1        10        20         5\n",
-      "map     0.0154  0.038303  0.043634  0.031517\n",
-      "ndcg    0.0154  0.056389  0.076198  0.039699\n",
-      "recall  0.0154  0.117072  0.196225  0.064911\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "2dde8a65268b46ce85312ed5529729a1",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 6, global step 1323: 'recall@10' reached 0.13214 (best 0.13214), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=6-step=1323-v2.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.015069  0.042212  0.047618  0.034832\n",
-      "ndcg    0.015069  0.062962  0.082865  0.044922\n",
-      "recall  0.015069  0.132141  0.211293  0.076006\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "d452da2b13ca4953b58fffaecb5143dc",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 7, global step 1512: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.015731  0.041482  0.047466  0.033985\n",
-      "ndcg    0.015731  0.061963  0.083911  0.043462\n",
-      "recall  0.015731  0.130651  0.217751  0.072694\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "0dc27c98e71f493da3ef87e551bd0c60",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 8, global step 1701: 'recall@10' reached 0.14406 (best 0.14406), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=8-step=1701-v1.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.015897  0.045616  0.051492  0.037233\n",
-      "ndcg    0.015897  0.068287  0.089852  0.047670\n",
-      "recall  0.015897  0.144064  0.229674  0.079649\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "b143a23c75a347139ec26ef54abad1b6",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 9, global step 1890: 'recall@10' reached 0.15102 (best 0.15102), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=9-step=1890-v2.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017387  0.048506  0.054854  0.040073\n",
-      "ndcg    0.017387  0.072139  0.095364  0.051380\n",
-      "recall  0.017387  0.151018  0.243087  0.086107\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "23cf891fa04440408cf156e3bb1d3245",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 10, global step 2079: 'recall@10' reached 0.15201 (best 0.15201), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=10-step=2079.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k             1        10        20         5\n",
-      "map     0.01689  0.048957  0.055559  0.040876\n",
-      "ndcg    0.01689  0.072775  0.096923  0.052810\n",
-      "recall  0.01689  0.152012  0.247723  0.089419\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "b9eee01e7c704ed9a0559f130d4ca137",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 11, global step 2268: 'recall@10' reached 0.16128 (best 0.16128), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=11-step=2268.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016228  0.049785  0.056337  0.040514\n",
-      "ndcg    0.016228  0.075524  0.099693  0.052985\n",
-      "recall  0.016228  0.161285  0.257493  0.091406\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "060953ee21ef46d49a34a75df75b57bd",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 12, global step 2457: 'recall@10' reached 0.16327 (best 0.16327), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=12-step=2457-v1.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.014738  0.048876  0.055841  0.039215\n",
-      "ndcg    0.014738  0.075207  0.100830  0.051456\n",
-      "recall  0.014738  0.163272  0.265110  0.089088\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "9db57c2e20e046e3a7635fec4ee98680",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 13, global step 2646: 'recall@10' reached 0.16393 (best 0.16393), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=13-step=2646.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016725  0.051229  0.058088  0.041919\n",
-      "ndcg    0.016725  0.077239  0.102477  0.054437\n",
-      "recall  0.016725  0.163934  0.264282  0.092896\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "4de755f88a624acd90ee5787f42d0eec",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 14, global step 2835: 'recall@10' reached 0.16791 (best 0.16791), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=14-step=2835.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.018049  0.052408  0.059556  0.042344\n",
-      "ndcg    0.018049  0.079003  0.105266  0.054474\n",
-      "recall  0.018049  0.167909  0.272231  0.091737\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "294c0a5c4b5a43d1850d722063b45faa",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 15, global step 3024: 'recall@10' reached 0.16907 (best 0.16907), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=15-step=3024.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017718  0.052563  0.059905  0.042822\n",
-      "ndcg    0.017718  0.079424  0.106455  0.055627\n",
-      "recall  0.017718  0.169068  0.276536  0.095049\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "5cec32b151e34d558aadb0a54a9fe645",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 16, global step 3213: 'recall@10' reached 0.17255 (best 0.17255), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=16-step=3213.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016725  0.053275  0.060719  0.043368\n",
-      "ndcg    0.016725  0.080840  0.108407  0.056778\n",
-      "recall  0.016725  0.172545  0.282497  0.098029\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "964972bd33324e2e88184dfa18ac414f",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 17, global step 3402: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.015069  0.052175  0.059507  0.042386\n",
-      "ndcg    0.015069  0.079773  0.106812  0.055673\n",
-      "recall  0.015069  0.171717  0.279351  0.096374\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "749e9689bfc842daa9ed09f725f4323d",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 18, global step 3591: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017387  0.052979  0.060514  0.042849\n",
-      "ndcg    0.017387  0.079924  0.107494  0.055138\n",
-      "recall  0.017387  0.169896  0.279185  0.092731\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "1226764060764a59a3f2cfb1a822f8b2",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 19, global step 3780: 'recall@10' reached 0.18066 (best 0.18066), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=19-step=3780.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016559  0.053610  0.060994  0.042965\n",
-      "ndcg    0.016559  0.082830  0.110151  0.056616\n",
-      "recall  0.016559  0.180659  0.289617  0.098692\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "3073e689c7d340b59b6989506c74c700",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 20, global step 3969: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.015069  0.052912  0.060238  0.042455\n",
-      "ndcg    0.015069  0.081482  0.108433  0.056028\n",
-      "recall  0.015069  0.176685  0.283822  0.097698\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f86fc39ef2104c69b1d1f89818719d08",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 21, global step 4158: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.015731  0.053080  0.060646  0.043241\n",
-      "ndcg    0.015731  0.080854  0.108792  0.056898\n",
-      "recall  0.015731  0.173207  0.284484  0.098857\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "356d76af5ca94c5ca0efa94d57a7f71c",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 22, global step 4347: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017056  0.055227  0.062985  0.044949\n",
-      "ndcg    0.017056  0.083916  0.112344  0.058574\n",
-      "recall  0.017056  0.179666  0.292433  0.100348\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "79f90e00872f44b7b33c0f8c77ff011a",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 23, global step 4536: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.015069  0.052468  0.060395  0.042041\n",
-      "ndcg    0.015069  0.080925  0.110088  0.055468\n",
-      "recall  0.015069  0.175857  0.291770  0.096705\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "8d04fd5e94fe4711aee0ad54b356d9eb",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 24, global step 4725: 'recall@10' reached 0.18397 (best 0.18397), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=24-step=4725.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017718  0.055757  0.063565  0.044925\n",
-      "ndcg    0.017718  0.085273  0.114202  0.058594\n",
-      "recall  0.017718  0.183971  0.299387  0.100513\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "cfd12e9d0ad847c09dcf07d9532add8f",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 25, global step 4914: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017387  0.054950  0.062547  0.044381\n",
-      "ndcg    0.017387  0.083432  0.111543  0.057612\n",
-      "recall  0.017387  0.178506  0.290611  0.098195\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "363e70c34eff435597795909950744bd",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 26, global step 5103: 'recall@10' reached 0.18463 (best 0.18463), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=26-step=5103.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.019043  0.057219  0.064977  0.046506\n",
-      "ndcg    0.019043  0.086570  0.114830  0.060294\n",
-      "recall  0.019043  0.184633  0.296407  0.102666\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "48eb130d2ce549d4a70d48979d153fd2",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 27, global step 5292: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.019208  0.057371  0.065001  0.046650\n",
-      "ndcg    0.019208  0.085884  0.114015  0.059848\n",
-      "recall  0.019208  0.180825  0.292764  0.100182\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "4f715157ede74adba3455a0ef7e5131e",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 28, global step 5481: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.018381  0.056056  0.063808  0.045615\n",
-      "ndcg    0.018381  0.085206  0.113792  0.059325\n",
-      "recall  0.018381  0.182812  0.296572  0.101507\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "32ee6a09dcc248cfa8af8e6f803c9970",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 29, global step 5670: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k             1        10        20         5\n",
-      "map     0.01689  0.054551  0.062793  0.044119\n",
-      "ndcg    0.01689  0.083030  0.113385  0.057502\n",
-      "recall  0.01689  0.178010  0.298725  0.098526\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "8b0022b7e1334ae198b407aa045798f0",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 30, global step 5859: 'recall@10' reached 0.18828 (best 0.18828), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=30-step=5859.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017387  0.057006  0.064480  0.046299\n",
-      "ndcg    0.017387  0.087281  0.114612  0.060979\n",
-      "recall  0.017387  0.188276  0.296572  0.106143\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "7e91bd3715e74bb88a5186a9bd2018b0",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 31, global step 6048: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016559  0.055557  0.063455  0.045518\n",
-      "ndcg    0.016559  0.084519  0.113808  0.059953\n",
-      "recall  0.016559  0.180825  0.297731  0.104322\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "6518c10046494dceb7539431f4342d40",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 32, global step 6237: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.015234  0.055037  0.063161  0.044312\n",
-      "ndcg    0.015234  0.085493  0.115172  0.059173\n",
-      "recall  0.015234  0.187117  0.304686  0.104984\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f6ed4206a4514b399847c7f23dc9ded5",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 33, global step 6426: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.015897  0.055037  0.063057  0.043967\n",
-      "ndcg    0.015897  0.085474  0.114917  0.058397\n",
-      "recall  0.015897  0.187117  0.304024  0.102832\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "07380901a70f4cf4bccda988256f4c63",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 34, global step 6615: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.014738  0.053996  0.061706  0.042709\n",
-      "ndcg    0.014738  0.083953  0.112215  0.056473\n",
-      "recall  0.014738  0.183971  0.296076  0.098692\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "c2f3ecd2dccc4e188f7a1e1098e35644",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 35, global step 6804: 'recall@10' reached 0.18844 (best 0.18844), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=35-step=6804-v1.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.015069  0.054859  0.063045  0.043283\n",
-      "ndcg    0.015069  0.085649  0.116068  0.057529\n",
-      "recall  0.015069  0.188442  0.309985  0.101341\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f668d2ae168b4a18aab4b55f571d65b5",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 36, global step 6993: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016725  0.057387  0.065605  0.046467\n",
-      "ndcg    0.016725  0.087272  0.117565  0.060911\n",
-      "recall  0.016725  0.186455  0.307004  0.105150\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "e99e9ff82afd4989910381ef0a5203b5",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 37, global step 7182: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.014903  0.055713  0.063883  0.044455\n",
-      "ndcg    0.014903  0.086266  0.116276  0.058719\n",
-      "recall  0.014903  0.188111  0.307336  0.102335\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "6db4c2801b96444ea9c1242071243bea",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 38, global step 7371: 'recall@10' reached 0.19208 (best 0.19208), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=38-step=7371.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.015897  0.056195  0.063902  0.044337\n",
-      "ndcg    0.015897  0.087493  0.115857  0.058652\n",
-      "recall  0.015897  0.192085  0.304852  0.102666\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "815fab87a5fe4603adf4c1db9881772e",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 39, global step 7560: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016559  0.056413  0.064644  0.045493\n",
-      "ndcg    0.016559  0.086769  0.117049  0.059937\n",
-      "recall  0.016559  0.188111  0.308495  0.104322\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "8759b410188f4ac795f88e785b27542f",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 40, global step 7749: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.014406  0.055492  0.063463  0.044392\n",
-      "ndcg    0.014406  0.086624  0.116225  0.059581\n",
-      "recall  0.014406  0.190263  0.308495  0.106309\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "7ec1302b87b3481db013eb0b41f690a8",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 41, global step 7938: 'recall@10' reached 0.19225 (best 0.19225), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=41-step=7938.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017221  0.058234  0.066382  0.047257\n",
-      "ndcg    0.017221  0.089155  0.118997  0.062111\n",
-      "recall  0.017221  0.192250  0.310647  0.107634\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "fef8f8d3fb0a4a6f8e6d3fcd7166421c",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 42, global step 8127: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.018381  0.057681  0.065612  0.046263\n",
-      "ndcg    0.018381  0.088403  0.117568  0.060508\n",
-      "recall  0.018381  0.191091  0.307004  0.104322\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "75e8ff0777b545d28653647466ff5fd3",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 43, global step 8316: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.015731  0.057188  0.065122  0.045949\n",
-      "ndcg    0.015731  0.088375  0.117464  0.060972\n",
-      "recall  0.015731  0.192250  0.307667  0.107137\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "b0dec430572c4d49a50a461c816aafe7",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 44, global step 8505: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.014903  0.055369  0.063525  0.044646\n",
-      "ndcg    0.014903  0.086007  0.115973  0.059935\n",
-      "recall  0.014903  0.187779  0.306839  0.106971\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "dfa485d5d42e49a8a9070ee85aef434f",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 45, global step 8694: 'recall@10' reached 0.19374 (best 0.19374), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=45-step=8694.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k             1        10        20         5\n",
-      "map     0.01391  0.055919  0.063806  0.044141\n",
-      "ndcg    0.01391  0.087688  0.116770  0.058898\n",
-      "recall  0.01391  0.193741  0.309488  0.104156\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "a38a593a5f2345caaf79560dc66e8d9f",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 46, global step 8883: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.018546  0.057837  0.065681  0.046448\n",
-      "ndcg    0.018546  0.088962  0.118050  0.061033\n",
-      "recall  0.018546  0.193078  0.309157  0.105978\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "6ce539a3aabb4dd695acc6c07ad72ac9",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 47, global step 9072: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016559  0.057167  0.065194  0.045987\n",
-      "ndcg    0.016559  0.088154  0.117904  0.060651\n",
-      "recall  0.016559  0.191588  0.310316  0.105647\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "30f737adc35f4b4c90a419732e8e1e91",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 48, global step 9261: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k             1        10        20         5\n",
-      "map     0.01689  0.058332  0.066214  0.047022\n",
-      "ndcg    0.01689  0.089275  0.118570  0.061644\n",
-      "recall  0.01689  0.192416  0.309488  0.106475\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "bf5a91c9655947cc91918c15b2a9462b",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 49, global step 9450: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016393  0.057926  0.066244  0.046906\n",
-      "ndcg    0.016393  0.088820  0.119375  0.062076\n",
-      "recall  0.016393  0.191422  0.312800  0.108627\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f21c5ac4b4734d74a23756a1fdf040f2",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 50, global step 9639: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017718  0.058476  0.066364  0.046785\n",
-      "ndcg    0.017718  0.089596  0.118568  0.061171\n",
-      "recall  0.017718  0.193410  0.308495  0.105315\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "7269482228054872bb3fd02ba7ed438b",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 51, global step 9828: 'recall@10' reached 0.19540 (best 0.19540), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=51-step=9828-v1.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.015731  0.057566  0.065119  0.045990\n",
-      "ndcg    0.015731  0.089326  0.117305  0.060894\n",
-      "recall  0.015731  0.195397  0.307004  0.106640\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f2e119321d274785b3001a5f44f941a4",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 52, global step 10017: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k             1        10        20         5\n",
-      "map     0.01954  0.058868  0.067041  0.048454\n",
-      "ndcg    0.01954  0.088730  0.118800  0.063159\n",
-      "recall  0.01954  0.188276  0.307832  0.108462\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "712a7c0cf2a048679ce1e68b12de03b6",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 53, global step 10206: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017718  0.058614  0.066505  0.047370\n",
-      "ndcg    0.017718  0.089564  0.118682  0.062093\n",
-      "recall  0.017718  0.192747  0.308660  0.107303\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "dfb0a6436d2147adb94f4a9c8ea16aac",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 54, global step 10395: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016725  0.057676  0.065385  0.046299\n",
-      "ndcg    0.016725  0.089061  0.117540  0.061333\n",
-      "recall  0.016725  0.193741  0.307170  0.107634\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f890babbe4e64f739ba391bce80360e9",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 55, global step 10584: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016062  0.056661  0.064644  0.045372\n",
-      "ndcg    0.016062  0.087800  0.117280  0.060313\n",
-      "recall  0.016062  0.191588  0.308992  0.106309\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f4c2e6243e484477bba7d4de22a19810",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 56, global step 10773: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.018049  0.058768  0.066948  0.047558\n",
-      "ndcg    0.018049  0.090098  0.120345  0.062504\n",
-      "recall  0.018049  0.194734  0.315284  0.108462\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "87dba884d40f4568acd8ca8bbbd64f6b",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 57, global step 10962: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.014572  0.053518  0.062111  0.042355\n",
-      "ndcg    0.014572  0.083976  0.115510  0.056682\n",
-      "recall  0.014572  0.185792  0.310979  0.100845\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "58699cb7550047feb645c6af4f054fcc",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 58, global step 11151: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.018049  0.058147  0.066599  0.047524\n",
-      "ndcg    0.018049  0.088199  0.119005  0.062165\n",
-      "recall  0.018049  0.188276  0.310151  0.107137\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "086f8398ee09499398aa7d3b70c79aa8",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 59, global step 11340: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.014406  0.055658  0.063896  0.044665\n",
-      "ndcg    0.014406  0.086842  0.117136  0.059875\n",
-      "recall  0.014406  0.190760  0.311144  0.106640\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "eb274ba8390345fc95f0e0f78a1f68e4",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 60, global step 11529: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.015731  0.056472  0.064719  0.045739\n",
-      "ndcg    0.015731  0.087657  0.117903  0.061421\n",
-      "recall  0.015731  0.191422  0.311475  0.109786\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "b4c4b8b17b194484b484a34ecf6c211b",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 61, global step 11718: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k             1        10        20         5\n",
-      "map     0.01689  0.057109  0.065266  0.045570\n",
-      "ndcg    0.01689  0.088058  0.118054  0.060104\n",
-      "recall  0.01689  0.191257  0.310482  0.104819\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "8869b85b84e141069b219175f8a5fd05",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 62, global step 11907: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.014738  0.056182  0.064093  0.044226\n",
-      "ndcg    0.014738  0.087940  0.116915  0.058911\n",
-      "recall  0.014738  0.193906  0.308826  0.103991\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "284a8281380c4b51bc3e4d1fa821ce55",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 63, global step 12096: 'recall@10' reached 0.19556 (best 0.19556), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=63-step=12096.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016228  0.058258  0.066323  0.047014\n",
-      "ndcg    0.016228  0.089951  0.119723  0.062421\n",
-      "recall  0.016228  0.195562  0.314125  0.109786\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "a172b3e4752e4ac8b1df1f435e3d0401",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 64, global step 12285: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016725  0.058308  0.066315  0.047469\n",
-      "ndcg    0.016725  0.089611  0.119061  0.063139\n",
-      "recall  0.016725  0.193741  0.310813  0.111442\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "3481bce9e0ee43c58b478c79937e831b",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 65, global step 12474: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.015069  0.055754  0.063692  0.043901\n",
-      "ndcg    0.015069  0.087118  0.116341  0.058438\n",
-      "recall  0.015069  0.191754  0.307998  0.103163\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "929ab3208337487998be56f8cc63c14a",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 66, global step 12663: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.015234  0.056930  0.065395  0.045352\n",
-      "ndcg    0.015234  0.088645  0.119658  0.060367\n",
-      "recall  0.015234  0.194403  0.317437  0.106475\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "2672359bbd564354a429be48c4ac375b",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 67, global step 12852: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.014241  0.055383  0.063480  0.044845\n",
-      "ndcg    0.014241  0.086363  0.116422  0.060699\n",
-      "recall  0.014241  0.189270  0.309323  0.109621\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "0403c1ea57694538b5bcdf72a28d3e25",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 68, global step 13041: 'recall@10' reached 0.19622 (best 0.19622), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=68-step=13041.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.018712  0.058867  0.067064  0.047028\n",
-      "ndcg    0.018712  0.090484  0.120756  0.061726\n",
-      "recall  0.018712  0.196225  0.316774  0.107137\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "a697daf59428459886d557106ec3cc86",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 69, global step 13230: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.018712  0.059296  0.067623  0.047822\n",
-      "ndcg    0.018712  0.090358  0.121063  0.062465\n",
-      "recall  0.018712  0.193906  0.316112  0.107468\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "0d87a6c21b0e4db48b4a6387e1482296",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 70, global step 13419: 'recall@10' reached 0.19672 (best 0.19672), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=70-step=13419.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017718  0.058932  0.066943  0.047301\n",
-      "ndcg    0.017718  0.090675  0.119968  0.062253\n",
-      "recall  0.017718  0.196721  0.312800  0.108296\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "bbb281775c9545eaa92d24f00b763b73",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 71, global step 13608: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017056  0.057241  0.065559  0.045708\n",
-      "ndcg    0.017056  0.088863  0.119476  0.060634\n",
-      "recall  0.017056  0.194569  0.316278  0.106640\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "285af00d69c3444a8d58e6349850a98f",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 72, global step 13797: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.015069  0.056120  0.064450  0.044969\n",
-      "ndcg    0.015069  0.087871  0.118417  0.060586\n",
-      "recall  0.015069  0.193741  0.314953  0.108793\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "021ea8662cb5411eb86c8f63337dceeb",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 73, global step 13986: 'recall@10' reached 0.19854 (best 0.19854), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=73-step=13986.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016559  0.057389  0.065353  0.045104\n",
-      "ndcg    0.016559  0.089853  0.119195  0.059893\n",
-      "recall  0.016559  0.198543  0.315284  0.105481\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "f45164cbd400421dada62034fa3d49eb",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 74, global step 14175: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.015565  0.056737  0.065008  0.045043\n",
-      "ndcg    0.015565  0.088335  0.118772  0.059753\n",
-      "recall  0.015565  0.193906  0.314953  0.104984\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "7990b92ef50c48daa04193a0749c46e1",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 75, global step 14364: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k            1        10        20         5\n",
-      "map     0.0154  0.056695  0.064954  0.044974\n",
-      "ndcg    0.0154  0.088831  0.119369  0.059974\n",
-      "recall  0.0154  0.196390  0.318099  0.106143\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "5246f04cb4324f55b47ecb33070d932c",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 76, global step 14553: 'recall@10' reached 0.19904 (best 0.19904), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=76-step=14553.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.014406  0.056638  0.064821  0.044654\n",
-      "ndcg    0.014406  0.089447  0.119451  0.060199\n",
-      "recall  0.014406  0.199040  0.318099  0.108130\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "786ee80f518d4518ae0f92245a103b52",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 77, global step 14742: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k             1        10        20         5\n",
-      "map     0.01689  0.057254  0.065636  0.045783\n",
-      "ndcg    0.01689  0.088655  0.119626  0.060673\n",
-      "recall  0.01689  0.193575  0.316940  0.106640\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "bb577bcfe6fb48b2acecd29f761a841c",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 78, global step 14931: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.013413  0.053570  0.062183  0.042154\n",
-      "ndcg    0.013413  0.084824  0.116217  0.056864\n",
-      "recall  0.013413  0.189270  0.313462  0.102169\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "56ce55e76d164531a232448502b5e89e",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 79, global step 15120: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.013744  0.056911  0.065346  0.045518\n",
-      "ndcg    0.013744  0.089081  0.119783  0.061084\n",
-      "recall  0.013744  0.196225  0.317602  0.108793\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "216cc04625a84a04ae600fcf7a58760f",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 80, global step 15309: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017056  0.059605  0.067669  0.048805\n",
-      "ndcg    0.017056  0.091420  0.121109  0.064925\n",
-      "recall  0.017056  0.197218  0.315284  0.114589\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "cd55a7a84fa444398684901f380306c7",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 81, global step 15498: 'recall@10' reached 0.19970 (best 0.19970), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=81-step=15498-v1.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.019208  0.059915  0.067934  0.047475\n",
-      "ndcg    0.019208  0.092051  0.121520  0.061684\n",
-      "recall  0.019208  0.199702  0.316774  0.105315\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "1b8397e043464992abbf5f67594004eb",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 82, global step 15687: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017056  0.059126  0.066993  0.047193\n",
-      "ndcg    0.017056  0.091040  0.119902  0.061589\n",
-      "recall  0.017056  0.197880  0.312469  0.105647\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "7d663efa0c9b418da3409d032859661c",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 83, global step 15876: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.016393  0.058910  0.067300  0.047367\n",
-      "ndcg    0.016393  0.090708  0.121604  0.062527\n",
-      "recall  0.016393  0.196556  0.319424  0.108958\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "6a583740ecc14ad9a0c10b8f102de5d9",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 84, global step 16065: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k            1        10        20         5\n",
-      "map     0.0154  0.057330  0.065707  0.046324\n",
-      "ndcg    0.0154  0.088907  0.119815  0.061892\n",
-      "recall  0.0154  0.194072  0.317105  0.109786\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "0a7154f282df47c58f709f68c7bea385",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 85, global step 16254: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017884  0.058865  0.067077  0.046901\n",
-      "ndcg    0.017884  0.090545  0.120613  0.061410\n",
-      "recall  0.017884  0.196390  0.315615  0.105978\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "638857741499485686c1b4b6733f2713",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 86, global step 16443: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.013578  0.056565  0.064611  0.044949\n",
-      "ndcg    0.013578  0.088694  0.118386  0.060223\n",
-      "recall  0.013578  0.195728  0.313959  0.106971\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "c6604bb63cc448b994d1cbc8076c0e58",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 87, global step 16632: 'recall@10' reached 0.20003 (best 0.20003), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=87-step=16632.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.014738  0.057765  0.065777  0.045747\n",
-      "ndcg    0.014738  0.090573  0.120003  0.061197\n",
-      "recall  0.014738  0.200033  0.316940  0.108627\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "62a8bd2f89314bb3a6a6a3a45e6c0877",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 88, global step 16821: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.014406  0.056543  0.064898  0.045063\n",
-      "ndcg    0.014406  0.088396  0.118998  0.060280\n",
-      "recall  0.014406  0.194569  0.315946  0.106971\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "1e01bdb063ae49369e758c0a4da17182",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 89, global step 17010: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.015731  0.057330  0.065518  0.045824\n",
-      "ndcg    0.015731  0.088940  0.118949  0.060987\n",
-      "recall  0.015731  0.194237  0.313297  0.107634\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "bbc483c2ccad4f0da0c4311a56f1a4b5",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 90, global step 17199: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.015897  0.058550  0.066817  0.046854\n",
-      "ndcg    0.015897  0.090655  0.121206  0.062134\n",
-      "recall  0.015897  0.197549  0.319258  0.108958\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "64e9ddfeb5424717a97741b58dbd6f0e",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 91, global step 17388: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.014075  0.057416  0.065911  0.046128\n",
-      "ndcg    0.014075  0.089700  0.120977  0.062278\n",
-      "recall  0.014075  0.196887  0.321245  0.111939\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "e3d7f5856e1f4fedb35c0f77a956fab0",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 92, global step 17577: 'recall@10' reached 0.20036 (best 0.20036), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=92-step=17577-v2.ckpt' as top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.014241  0.057366  0.065310  0.044878\n",
-      "ndcg    0.014241  0.090287  0.119605  0.059684\n",
-      "recall  0.014241  0.200364  0.317105  0.104984\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "c6e17c5b3e6342c79419d1b83e848f30",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 93, global step 17766: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017221  0.058259  0.066380  0.047455\n",
-      "ndcg    0.017221  0.089409  0.119171  0.063038\n",
-      "recall  0.017221  0.193078  0.311144  0.111111\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "1dd6b48245004c5d994afe4da809e807",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 94, global step 17955: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.015897  0.058542  0.066742  0.046857\n",
-      "ndcg    0.015897  0.090490  0.120543  0.062166\n",
-      "recall  0.015897  0.196721  0.315946  0.109124\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "626ec34631fd432f82b6ab0bcb51a20c",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 95, global step 18144: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017221  0.059014  0.067464  0.047632\n",
-      "ndcg    0.017221  0.090545  0.121542  0.062936\n",
-      "recall  0.017221  0.195397  0.318430  0.109952\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "e1ad3bc2d94c407b995ac6696468895f",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 96, global step 18333: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.017056  0.059395  0.067978  0.047806\n",
-      "ndcg    0.017056  0.091108  0.122518  0.062947\n",
-      "recall  0.017056  0.196721  0.321245  0.109455\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "d4b13e0d6c38444fa1533a6ca3c51133",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 97, global step 18522: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.015731  0.057509  0.065688  0.046131\n",
-      "ndcg    0.015731  0.089842  0.119812  0.062111\n",
-      "recall  0.015731  0.197549  0.316443  0.111442\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "271f58b7cce546aba843acb2cf87ebe5",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 98, global step 18711: 'recall@10' was not in top 1\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k              1        10        20         5\n",
-      "map     0.014241  0.057645  0.065858  0.045747\n",
-      "ndcg    0.014241  0.090083  0.120261  0.061015\n",
-      "recall  0.014241  0.198212  0.318099  0.107799\n",
-      "\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "33fc04054c0942dba61679e0f8bda4b2",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Validation: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "Epoch 99, global step 18900: 'recall@10' reached 0.20136 (best 0.20136), saving model to '/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=99-step=18900.ckpt' as top 1\n",
-      "`Trainer.fit` stopped: `max_epochs=100` reached.\n"
-     ]
-    },
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "k             1        10        20         5\n",
-      "map     0.01391  0.056347  0.064440  0.044130\n",
-      "ndcg    0.01391  0.089694  0.119503  0.059647\n",
-      "recall  0.01391  0.201358  0.319921  0.107468\n",
-      "\n"
-     ]
-    }
-   ],
-   "source": [
-    "from lightning.pytorch.callbacks import ModelCheckpoint\n",
-    "from lightning.pytorch.loggers import CSVLogger\n",
-    "\n",
-    "from replay.nn.lightning.callback import ComputeMetricsCallback\n",
-    "\n",
-    "\n",
-    "checkpoint_callback = ModelCheckpoint(\n",
-    "    dirpath=\"sasrec/checkpoints\",\n",
-    "    save_top_k=1,\n",
-    "    verbose=True,\n",
-    "    monitor=\"recall@10\",\n",
-    "    mode=\"max\",\n",
-    ")\n",
-    "\n",
-    "validation_metrics_callback = ComputeMetricsCallback(\n",
-    "    metrics=[\"map\", \"ndcg\", \"recall\"],\n",
-    "    ks=[1, 5, 10, 20],\n",
-    "    item_count=NUM_UNIQUE_ITEMS,\n",
-    ")\n",
-    "\n",
-    "csv_logger = CSVLogger(save_dir=\"sasrec/logs/train\", name=\"SasRec-example\")\n",
-    "\n",
-    "trainer = L.Trainer(\n",
-    "    max_epochs=100,\n",
-    "    callbacks=[checkpoint_callback, validation_metrics_callback],\n",
-    "    logger=csv_logger,\n",
-    ")\n",
-    "\n",
-    "trainer.fit(model, datamodule=parquet_module)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 21,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAiEAAAGnCAYAAABhFoffAAAAOnRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjEwLjgsIGh0dHBzOi8vbWF0cGxvdGxpYi5vcmcvwVt1zgAAAAlwSFlzAAASdAAAEnQB3mYfeAAAKj5JREFUeJzt3XtU1XW+//EXqNwFJUVQseyoo4MameDIcSRLPUzemJI0x6bSdI5aap3K0OUq82QWWmtaxtHmdNLRpouXdHRKU0edTpk7b53OiYZDjgoIomOGG4FUPr8//LGPtEEufvGziedjLRbD57Y/3/es5fe19vezd37GGCMAAIDrzN/2BgAAQPNECAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWtLS9AV9w9uxZ7dmzR7GxsQoMDLS9HQAAmozy8nLl5uYqOTlZbdq0qddcQoikPXv2KDU11fY2AABosjZu3KgxY8bUaw4hRFJsbKykywXs1q2bI2u63W65XC4lJiYqLCzMkTWbM+rpPGrqPGrqLOrpvMaoaU5OjlJTUz330voghEieRzDdunVTXFycI2sWFxersLBQvXr1Unh4uCNrNmfU03nU1HnU1FnU03mNWdOGHGfgYCoAALCCEAIAAKwghAAAACsIIQAAwApCCAAAsIIQAgAArCCEAAAAKwghAADACkIIAACwghACAACsIIQAAAArCCEAAMAKQggAALCCEAIAAKxoaXsDAADAGTc9/aer9kcHG6XHX5+91AXvhAAAACsIIQAAwApCCAAAsIIQAgAArCCEAAAAKwghAADACkIIAACwghACAACsIIQAAAArCCEAAMAKQggAALCCEAIAAKwghAAAACsIIQAAwApCCAAAsIIQAgAArCCEAAAAKwghAADACkIIAACwghACAACsIIQAAAArCCEAAMAKQggAALCCEAIAAKwghAAAACsIIQAAwApCCAAAsIIQAgAArCCEAAAAKwghAADACkIIAACwokEh5PPPP9cjjzyiuLg4hYaGqkuXLrr33nuVnZ3tNTYrK0spKSkKCwtTZGSk7r//fp06dcprXEVFhV566SV17dpVQUFB6tu3r95+++1qX7+uawIAAN/VsiGTXnzxRX3yySdKS0tT3759VVhYqGXLlqlfv3767LPP1Lt3b0lSXl6eBg8erIiICC1atEhut1tLlizRl19+KZfLpYCAAM+a8+bN0+LFizVlyhQlJCRo06ZNmjBhgvz8/DR+/HjPuPqsCQAAfFeDQsjjjz+uP/zhD1Vu+OPGjVOfPn20ePFirVmzRpK0aNEilZSU6MCBA+rSpYskKTExUcOGDdPKlSs1depUSVJ+fr6WLl2qGTNmaNmyZZKkhx9+WMnJyXryySeVlpamFi1a1GtNAADg2xr0OCYpKcnrHYfu3bsrLi5OWVlZnrb169dr5MiRnrAgSUOHDlWPHj303nvvedo2bdqkCxcuaPr06Z42Pz8/TZs2TXl5edq7d2+91wQAAL6tQe+EVMcYo5MnTyouLk7S5Xc3ioqK1L9/f6+xiYmJ+uCDDzx/Hzp0SKGhoerVq5fXuMr+QYMG1WvNmhQVFXmdH8nJyZEkud1uFRcX17pGXZSUlFT5jWtDPZ1HTZ1HTZ1FPesvOthctb9d0OV+J2vqdrsbPNexEPLWW28pPz9fzz33nCSpoKBAkhQTE+M1NiYmRmfOnFF5ebkCAwNVUFCgDh06yM/Pz2ucJJ04caLea9YkMzNTCxYsqLbP5XKpsLCwtkutF5fL5eh6zR31dB41dR41dRb1rLv0+LqNc7Kmx48fb/BcR0LI119/rRkzZmjgwIF64IEHJEmlpaWSVG0gCAoK8owJDAz0/L7auPquWZPp06crLS2tSltOTo5SU1OVmJjo9W5MQ5WUlMjlcikxMVGhoaGOrNmcUU/nUVPnUVNnUc/6G/7KX67a3y7IaErPCkdreuUxjPq65hBSWFioESNGKCIiQuvWrfMcIA0ODpYklZeXe80pKyurMiY4OLjO4+q6Zk2ioqIUFRVVbV9YWJjCw8OvOr++QkNDHV+zOaOezqOmzqOmzqKedVdY6lf7IDlb07CwsAbPvaYvK/vuu+/0i1/8QmfPntXWrVvVsWNHT1/lI5PKRyhXKigoUGRkpOcdi5iYGBUWFsoY4zVOkmfd+qwJAAB8W4NDSFlZmUaNGqXs7Gxt2bJFP/3pT6v0d+rUSe3bt9f+/fu95rpcLsXHx3v+jo+P1/nz573e0tm3b5+nv75rAgAA39agEHLp0iWNGzdOe/fu1dq1azVw4MBqx91zzz3asmWLcnNzPW07d+5UdnZ2lXMZY8aMUatWrZSZmelpM8Zo+fLl6tSpk5KSkuq9JgAA8G0NOhPyL//yL/rjH/+oUaNG6cyZM54vJ6s0ceJESdLcuXO1du1aDRkyRLNmzZLb7VZGRob69Omjhx56yDO+c+fOmj17tjIyMnThwgUlJCRo48aN+vjjj/XWW295zpnUZ00AAODbGhRCDh8+LEnavHmzNm/e7NVfGUJiY2O1Z88ePf7443r66acVEBCgESNGaOnSpV5nNxYvXqy2bdtqxYoVWrlypbp37641a9ZowoQJVcbVZ00AAOC7GhRCdu/eXeexcXFx2rZtW63j/P39lZ6ervT0dMfWBAAAvuuaPh0DAADQUIQQAABgBSEEAABYQQgBAABWEEIAAIAVhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWEEIAAIAVhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWEEIAAIAVhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWEEIAAIAVhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWEEIAAIAVhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWNDiEuN1uPfPMM0pJSVFkZKT8/Py0cuVKr3EPPvig/Pz8vH569uzpNbaiokIvvfSSunbtqqCgIPXt21dvv/12ta+flZWllJQUhYWFKTIyUvfff79OnTrV0MsBAADXWcuGTjx9+rSee+45denSRbfccot2795d49jAwED9+7//e5W2iIgIr3Hz5s3T4sWLNWXKFCUkJGjTpk2aMGGC/Pz8NH78eM+4vLw8DR48WBEREVq0aJHcbreWLFmiL7/8Ui6XSwEBAQ29LAAAcJ00OITExMSooKBA0dHR2r9/vxISEmp+kZYtNXHixKuul5+fr6VLl2rGjBlatmyZJOnhhx9WcnKynnzySaWlpalFixaSpEWLFqmkpEQHDhxQly5dJEmJiYkaNmyYVq5cqalTpzb0sgAAwHXS4McxgYGBio6OrvP4S5cuqbi4uMb+TZs26cKFC5o+fbqnzc/PT9OmTVNeXp727t3raV+/fr1GjhzpCSCSNHToUPXo0UPvvfdePa8EAADYcF0Opp4/f17h4eGKiIhQZGSkZsyYIbfbXWXMoUOHFBoaql69elVpT0xM9PRLl98xKSoqUv/+/b1eJzEx0TMOAAD4tgY/jqmrmJgYPfXUU+rXr58qKiq0detWZWZm6osvvtDu3bvVsuXlLRQUFKhDhw7y8/Pzmi9JJ06c8Iy7sv2HY8+cOaPy8nIFBgZWu5+ioiKvA6w5OTmSLh+2vdq7NfVRUlJS5TeuDfV0HjV1HjV1FvWsv+hgc9X+dkGX+52s6Q/fVKiPRg8hL7zwQpW/x48frx49emjevHlat26d58BpaWlptcEhKCjI03/l79rG1hRCMjMztWDBgmr7XC6XCgsL63JZdeZyuRxdr7mjns6jps6jps6innWXHl+3cU7W9Pjx4w2e2+ghpDqPPfaY5s+frx07dnhCSHBwsMrLy73GlpWVefqv/F2XsdWZPn260tLSqrTl5OQoNTVViYmJXo+DGqqkpEQul0uJiYkKDQ11ZM3mjHo6j5o6j5o6i3rW3/BX/nLV/nZBRlN6Vjha06ysrAbPtRJCgoODdcMNN+jMmTOetpiYGO3atUvGmCqPZCofv3Ts2NEz7sr2KxUUFCgyMrLGd0EkKSoqSlFRUdX2hYWFKTw8vP4XdBWhoaGOr9mcUU/nUVPnUVNnUc+6Kyz1q32QnK1pWFhYg+da+cbUc+fO6fTp02rfvr2nLT4+XufPn/dKVPv27fP0S1KnTp3Uvn177d+/32tdl8vlGQcAAHxbo4aQsrIynTt3zqt94cKFMsYoJSXF0zZmzBi1atVKmZmZnjZjjJYvX65OnTopKSnJ037PPfdoy5Ytys3N9bTt3LlT2dnZXo9aAACAb7qmxzHLli3T2bNnPZ9c2bx5s/Ly8iRJjz76qL799lvdeuutuu+++zxf075t2zZ98MEHSklJ0ZgxYzxrde7cWbNnz1ZGRoYuXLighIQEbdy4UR9//LHeeustzxeVSdLcuXO1du1aDRkyRLNmzZLb7VZGRob69Omjhx566FouyXHDX/lLjW+PHV084jrvBgAA33FNIWTJkiU6duyY5+8NGzZow4YNkqSJEyeqTZs2GjlypLZv365Vq1bp0qVL6tatmxYtWqQnnnhC/v5V34hZvHix2rZtqxUrVmjlypXq3r271qxZowkTJlQZFxsbqz179ujxxx/X008/rYCAAI0YMUJLly696nkQAADgO64phBw9erTWMatXr67zev7+/kpPT1d6enqtY+Pi4rRt27Y6rw0AAHyLlYOpAAAAhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWEEIAAIAVhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWEEIAAIAVhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWEEIAAIAVhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWEEIAAIAVhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWEEIAAIAVhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBUNDiFut1vPPPOMUlJSFBkZKT8/P61cubLasVlZWUpJSVFYWJgiIyN1//3369SpU17jKioq9NJLL6lr164KCgpS37599fbbb1/TmgAAwDe1bOjE06dP67nnnlOXLl10yy23aPfu3dWOy8vL0+DBgxUREaFFixbJ7XZryZIl+vLLL+VyuRQQEOAZO2/ePC1evFhTpkxRQkKCNm3apAkTJsjPz0/jx49v0JoAAMA3NTiExMTEqKCgQNHR0dq/f78SEhKqHbdo0SKVlJTowIED6tKliyQpMTFRw4YN08qVKzV16lRJUn5+vpYuXaoZM2Zo2bJlkqSHH35YycnJevLJJ5WWlqYWLVrUa00AAOC7Gvw4JjAwUNHR0bWOW79+vUaOHOkJC5I0dOhQ9ejRQ++9956nbdOmTbpw4YKmT5/uafPz89O0adOUl5envXv31ntNAADguxr1YGp+fr6KiorUv39/r77ExEQdOnTI8/ehQ4cUGhqqXr16eY2r7K/vmgAAwHc1+HFMXRQUFEi6/Ojmh2JiYnTmzBmVl5crMDBQBQUF6tChg/z8/LzGSdKJEyfqvWZ1ioqKvA6w5uTkSLp82La4uLg+l1ijkpISSVK7IFPjGKdeqzmorGflb1w7auo8auos6ll/0cE133Ok/7snOVlTt9vd4LmNGkJKS0slqdpAEBQU5BkTGBjo+X21cfVdszqZmZlasGBBtX0ul0uFhYVXvab6mtKzosa+Xbt2OfpazYHL5bK9hR8dauo8auos6ll36fF1G+dkTY8fP97guY0aQoKDgyVJ5eXlXn1lZWVVxgQHB9d5XF3XrM706dOVlpZWpS0nJ0epqalKTEz0ehzUUCUlJXK5XPrd1/46XeZX7ZiPHhvsyGs1B5X1TExMVGhoqO3t/ChQU+dRU2dRz/ob/spfrtrfLshoSs8KR2ualZXV4LmNGkIqH5lUPkK5UkFBgSIjIz3vWMTExGjXrl0yxlR5JFM5t2PHjvVeszpRUVGKioqqti8sLEzh4eF1ubQ6O13mp8LS6kOI06/VHISGhlI3h1FT51FTZ1HPuqvpfvNDTtY0LCyswXMb9WBqp06d1L59e+3fv9+rz+VyKT4+3vN3fHy8zp8/75Wo9u3b5+mv75oAAMB3NfrXtt9zzz3asmWLcnNzPW07d+5UdnZ2lcciY8aMUatWrZSZmelpM8Zo+fLl6tSpk5KSkuq9JgAA8F3X9Dhm2bJlOnv2rOeTK5s3b1ZeXp4k6dFHH1VERITmzp2rtWvXasiQIZo1a5bcbrcyMjLUp08fPfTQQ561OnfurNmzZysjI0MXLlxQQkKCNm7cqI8//lhvvfWW54vKJNV5TQAA4LuuKYQsWbJEx44d8/y9YcMGbdiwQZI0ceJERUREKDY2Vnv27NHjjz+up59+WgEBARoxYoSWLl3qdXZj8eLFatu2rVasWKGVK1eqe/fuWrNmjSZMmFBlXH3WBAAAvumaQsjRo0frNC4uLk7btm2rdZy/v7/S09OVnp7u2JoAAMA3NfqZEAAAgOoQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWEEIAAIAVhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWEEIAAIAVhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWEEIAAIAVhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWEEIAAIAVhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWEEIAAIAVhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVjR6CNm9e7f8/Pyq/fnss8+qjP300081aNAghYSEKDo6WjNnzpTb7fZas7y8XHPmzFHHjh0VHBysAQMGaPv27Y19KQAAwEEtr9cLzZw5UwkJCVXaunXr5vnfhw8f1p133qlevXrp5ZdfVl5enpYsWaL//d//1Ycfflhl3oMPPqh169Zp9uzZ6t69u1auXKm77rpLu3bt0qBBg67L9QAAgGtz3ULIz3/+c40dO7bG/rlz56pt27bavXu3wsPDJUk33XSTpkyZoo8++kjDhw+XJLlcLr3zzjvKyMjQE088IUn69a9/rd69e+upp57Sp59+2vgXAwAArtl1PRNy7tw5Xbx40au9uLhY27dv18SJEz0BRLocLsLCwvTee+952tatW6cWLVpo6tSpnragoCBNnjxZe/fuVW5ubuNeBAAAcMR1CyEPPfSQwsPDFRQUpCFDhmj//v2evi+//FIXL15U//79q8wJCAhQfHy8Dh065Gk7dOiQevToUSWsSFJiYqKky491AACA72v0xzEBAQG65557dNddd6ldu3b66quvtGTJEv385z/Xp59+qltvvVUFBQWSpJiYGK/5MTEx+vjjjz1/FxQU1DhOkk6cOHHV/RQVFenUqVNV2nJyciRJbrdbxcXF9bvAGpSUlEiS2gWZGsc49VrNQWU9K3/j2lFT51FTZ1HP+osOrvmeI/3fPcnJmlb3AZK6avQQkpSUpKSkJM/fo0eP1tixY9W3b1+lp6dr69atKi0tlSQFBgZ6zQ8KCvL0S1JpaWmN4yr7ryYzM1MLFiyots/lcqmwsLD2i6qHKT0rauzbtWuXo6/VHLhcLttb+NGhps6jps6innWXHl+3cU7W9Pjx4w2ee90Opl6pW7duGjNmjDZs2KBLly4pODhY0uWP3v5QWVmZp1+SgoODaxxX2X8106dPV1paWpW2nJwcpaamKjExUb169ar39VSnpKRELpdLv/vaX6fL/Kod89Fjgx15reagsp6JiYkKDQ21vZ0fBWrqPGrqLOpZf8Nf+ctV+9sFGU3pWeFoTbOysho810oIkaTY2Fh9//33Kikp8TxKqXwsc6WCggJ17NjR83dMTIzy8/OrHSepytjqREVFKSoqqtq+sLAwr7Mm1+p0mZ8KS6sPIU6/VnMQGhpK3RxGTZ1HTZ1FPeuupvvNDzlZ07CwsAbPtfaNqUeOHFFQUJDCwsLUu3dvtWzZssphVUn6/vvvdfjwYcXHx3va4uPjlZ2d7XWeYt++fZ5+AADg+xo9hPzwEKgkffHFF/rjH/+o4cOHy9/fXxERERo6dKjWrFmjc+fOecatXr1abre7yuOTsWPH6tKlS3r99dc9beXl5XrzzTc1YMAAxcbGNu4FAQAARzT645hx48YpODhYSUlJioqK0ldffaXXX39dISEhWrx4sWfc888/r6SkJCUnJ2vq1KnKy8vT0qVLNXz4cKWkpHjGDRgwQGlpaUpPT1dRUZG6deumVatW6ejRo3rjjTca+3IAAIBDGv2dkNTUVJ0+fVovv/yypk+frnfffVd333239u/fX+UQaL9+/bRjxw4FBwfrscce0+uvv67Jkydr3bp1Xmv+/ve/1+zZs7V69WrNnDlTFy5c0JYtWzR4MAc9AQBoKhr9nZCZM2dq5syZdRo7aNAgffLJJ7WOCwoKUkZGhjIyMq51ewAAwBJrB1MBAEDzRggBAABWEEIAAIAVhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWEEIAAIAVhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWEEIAAIAVhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWEEIAAIAVhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWEEIAAIAVhBAAAGAFIQQAAFhBCAEAAFYQQgAAgBWEEAAAYAUhBAAAWEEIAQAAVhBCAACAFYQQAABgBSEEAABYQQgBAABWEEIAAIAVhBAAAGAFIQQAAFjRZENIeXm55syZo44dOyo4OFgDBgzQ9u3bbW8LAADUUZMNIQ8++KBefvll/epXv9Jvf/tbtWjRQnfddZf+8z//0/bWAABAHbS0vYGGcLlceuedd5SRkaEnnnhCkvTrX/9avXv31lNPPaVPP/3U8g4BAEBtmuQ7IevWrVOLFi00depUT1tQUJAmT56svXv3Kjc31+LuAABAXTTJd0IOHTqkHj16KDw8vEp7YmKiJOnw4cOKjY2tdm5RUZFOnTpVpe2rr76SJP3Xf/2X3G63I3ssLS3V8ePHFfidv8LLqx+zb98+R16rOais58GDBxUcHGx7Oz8K1NR51NRZ1LP+ws8dvWp/4PfS8eMVjtb0yJEjki6f1ayvJhlCCgoKFBMT49Ve2XbixIka52ZmZmrBggXV9k2YMMGZDdbRzzKv68sBAKC/NNK6ubm56tevX73mNMkQUlpaqsDAQK/2oKAgT39Npk+frrS0tCptxcXFys7OVp8+fapdtyFycnKUmpqqjRs3qlu3bo6s2ZxRT+dRU+dRU2dRT+c1Rk3Ly8uVm5ur5OTkes9tkiEkODi42rd9ysrKPP01iYqKUlRUlFf7wIEDndvgFbp166a4uLhGWbs5op7Oo6bOo6bOop7Oc7qm9X0HpFKTPJgaExOjgoICr/bKto4dO17vLQEAgHpqkiEkPj5e2dnZKi4urtJeedAzPj7ewq4AAEB9NMkQMnbsWF26dEmvv/66p628vFxvvvmmBgwYUOMnYwAAgO9okmdCBgwYoLS0NKWnp6uoqEjdunXTqlWrdPToUb3xxhu2tydJat++vZ555hm1b9/e9lZ+FKin86ip86ips6in83ytpn7GGGN7Ew1RVlam+fPna82aNfr222/Vt29fLVy4UP/0T/9ke2sAAKAOmmwIAQAATVuTPBMCAACaPkIIAACwghACAACsIIQAAAArCCH/38GDBzV69GhFRkYqJCREvXv31quvvurpv/322+Xn5+f1k5KSUmWdzz//XI888oji4uIUGhqqLl266N5771V2dnad93L27FlNnTpV7du3V2hoqIYMGaKDBw86dq3Xi6/UdOfOnZo0aZJ69OihkJAQ3XzzzXr44Yer/dZdX+Yr9fyhKVOmyM/PTyNHjrym67PB12q6Y8cO3XHHHYqIiFDr1q1122236d1333XkWq8XX6rpgQMHNHLkSEVHRyssLEx9+/bVq6++qkuXLjl2vY3NqXr+z//8j9LS0nTzzTcrJCRE7dq10+DBg7V58+Y676Ux7k1N8ntCnPbRRx9p1KhRuvXWWzV//nyFhYXpm2++UV5eXpVxnTt31gsvvFCl7YdfEf/iiy/qk08+UVpamvr27avCwkItW7ZM/fr102effabevXtfdS8VFRUaMWKEvvjiCz355JNq166dMjMzdfvtt+vAgQPq3r27MxfdyHyppnPmzNGZM2eUlpam7t2768iRI1q2bJm2bNmiw4cPKzo62pmLbkS+VM8r7d+/XytXrvT8xyObEl+r6ZtvvqnJkydr2LBhWrRokVq0aKG//vWvys3NvfaLvU58qaYHDhxQUlKSunfvrjlz5igkJEQffvihZs2apW+++Ua//e1vnbnoRuRkPY8dO6Zz587pgQceUMeOHXX+/HmtX79eo0eP1ooVKzR16tSr7qXR7k2mmfvuu+9Mhw4dzC9/+Utz6dKlGsclJyebuLi4Wtf75JNPTHl5eZW27OxsExgYaH71q1/VOv/dd981kszatWs9bUVFRaZNmzbmvvvuq3W+L/C1mu7Zs8drH3v27DGSzLx582qdb5uv1bNSRUWFGThwoJk0aZK58cYbzYgRI+o81zZfq+nf/vY3ExwcbGbOnFn75n2Ur9V0ypQpJiAgwPz973+v0j548GATHh5e63zbnK5ndS5evGhuueUW85Of/KTWsY11b2r2j2P+8Ic/6OTJk3r++efl7++vkpISVVRU1Dj+4sWLcrvdNfYnJSUpICCgSlv37t0VFxenrKysWvezbt06dejQQXfffbenrX379rr33nu1adOmav/rwb7G12o6ePBg+fv7e7VFRkbWab5tvlbPSqtXr9Z///d/6/nnn6/zHF/hazVdvny5Ll26pOeee06S5Ha7ZZrYVzj5Wk2Li4sVFBSkNm3aVGmPiYm56n9p3Vc4Xc/qtGjRQrGxsTp79mytYxvr3tTsQ8iOHTsUHh6u/Px8/eQnP1FYWJjCw8M1bdo0lZWVVRmbnZ2t0NBQtW7dWtHR0Zo/f74uXLhQ62sYY3Ty5Em1a9eu1rGHDh1Sv379vG6aiYmJOn/+fIOf219PvlbT6rjdbrnd7gbPv558sZ7nzp3TnDlzNHfu3CbxOOuHfK2mO3bsUM+ePfXBBx+oc+fOat26tW644QbNnz//qjceX+JrNb399ttVXFys3/zmN8rKytKxY8e0fPlybdiwQenp6Q2+zuulsepZUlKi06dP65tvvtErr7yiDz/8UHfeeWet+2m0e1OD30P5kejbt68JCQkxISEh5tFHHzXr1683jz76qJFkxo8f7xk3adIk8+yzz5r169eb3//+92b06NFGkrn33ntrfY3Vq1cbSeaNN96odWxoaKiZNGmSV/uf/vQnI8ls3bq1fhdoga/VtDoLFy40kszOnTsbNP968sV6PvHEE6Zr166mrKzMGGOa3OMYX6tpeHi4adu2rQkMDDTz588369atMxMmTDCSzNNPP31N13q9+FpNL168aB555BHTqlUrI8lIMi1atDD/9m//dk3Xeb00Vj1/85vfeOrh7+9vxo4da86cOVPrfhrr3tTsQ8jNN99sJJl//ud/rtJe+X9UdnZ2jXOnTJliJJm9e/fWOCYrK8uEh4ebgQMHmosXL9a6H39/fzNt2jSv9p07dxpJ5v333691Ddt8raY/tGfPHtOyZcs6/aPnC3ytnn/9619Nq1atzLp16zxtTS2E+FpN/f39jSSzePHiKu0pKSkmODjYFBcX17qGbb5WU2OMeeWVV8zIkSPNqlWrzLvvvmtSU1NNy5Ytm/W/o1lZWWb79u1m1apVZsSIEeaXv/ylKSwsrHU/jXVvavYhJC4uzkgye/bsqdJeeXBx1apVNc79+uuvjSSzcOHCavsLCgrMzTffbGJjY01+fn6d9vNjeCfE12p6paysLBMZGWni4+ObxD/sxvhePVNSUkxycnKVtqYWQnytpqGhoUaSOXbsWJX2VatWVbtPX+RrNX3hhRdMdHS0OXfuXJX222+/3XTs2NFcuHChTuvY0pj1vNKwYcNMQkKCqaiouOq4xro3NfszIZUfY+rQoUOV9qioKEnSt99+W+Pc2NhYSdKZM2e8+r777jv94he/0NmzZ7V161avj0vVJCYmptrvr6hsq+s6NvlaTSvl5uZq+PDhioiI0AcffKDWrVvXa74tvlTPP//5z9q6datmzZqlo0ePen4uXryo0tJSHT16VMXFxXW+Nlt8qabXuh9f4Ws1zczM1B133KGwsLAq7aNHj9aJEyd09OjROq1jS2PV84fGjh2rzz//vNYzHY11b2r2IeS2226TJOXn51dpP3HihKTLp39rcuTIkWrHlJWVadSoUcrOztaWLVv005/+tM77iY+P18GDB70Oo+3bt08hISHq0aNHndeyxddqKkl///vfNXz4cJWXl2vbtm2KiYmp13ybfKmex48flyTdfffd6tq1q+cnPz9ff/7zn9W1a1f9x3/8R90uzCJfqum17sdX+FpNT548We2XklUe2Lx48WKd17KhMepZndLSUkmXw97VNNq9qUHvn/yIHDx40EgyEyZMqNJ+3333mZYtW5r8/Hzz3XffeQ7gVaqoqDDjxo0zksyBAwc87RcvXjSjR482LVu2NH/605+u+tonTpwwWVlZ5vvvv/e0vfPOO16fxT516pRp06aNGTdu3LVc6nXjazV1u90mMTHRtG7d2uzfv9+BK7y+fKmex44dM++//77XT/v27U3//v3N+++/b3Jychy68sbjSzU1xpj333/fSDJz5871tF26dMkMGjTIREZGeu3DF/laTXv37m0iIyPN6dOnq6x52223mdatW1cZ64ucrufJkye9XuP77783/fr1M8HBwVUeW13Pe1OzDyHGXD5drP9/mvi1114zaWlpRpJJT083xhiza9cuEx0dbR577DHz2muvmSVLlph//Md/NJLM1KlTq6w1a9YsI8mMGjXKrF692uvnSg888ICRZP72t7952i5evGh+9rOfmbCwMLNgwQLz2muvmbi4ONO6dWvz9ddfN3otnOJLNR0zZoyRZCZNmuQ1tykcUDPGt+pZnaZ2JsQY36ppRUWFufPOO42fn5+ZOnWqee2118ywYcOMJLNixYpGr4VTfKmma9asMZLMP/zDP5gXX3zRvPrqq2bgwIFGkvnXf/3XRq+FE5ysZ2pqqrnjjjvMs88+a373u9+ZhQsXmp49expJZunSpVXGXs97EyHEXE6Dzz77rLnxxhtNq1atTLdu3cwrr7zi6T9y5IhJS0szN910kwkKCjIhISHmtttuM8uXL/c6zJOcnOz5+FN1P1eq6R/4M2fOmMmTJ5sbbrjBhISEmOTkZPP555831uU3Cl+q6Y033ljj3BtvvLERq+AcX6pndZpiCPG1mp47d87MmjXLREdHm4CAANOnTx+zZs2axrr8RuFrNd26datJTk427dq189R0+fLljXX5jnOynm+//bYZOnSo6dChg2nZsqVp27atGTp0qNm0aZPX617Pe5OfMU3sa/kAAMCPQrM/mAoAAOwghAAAACsIIQAAwApCCAAAsIIQAgAArCCEAAAAKwghAADACkIIAACwghACAACsIIQAAAArCCEAAMAKQggAALCCEAIAAKwghAAAACsIIQAAwIr/B8ppMEdqEMoaAAAAAElFTkSuQmCC",
-      "text/plain": [
-       "<Figure size 600x480 with 1 Axes>"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    }
-   ],
-   "source": [
-    "import matplotlib.pyplot as plt\n",
-    "\n",
-    "plt.figure(figsize = (5, 4), dpi = 120)\n",
-    "plt.hist(transforms[\"train\"][0].frequencies.cpu().numpy(), bins = 50)\n",
-    "plt.grid()\n"
-   ]
-  },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Now we can get the best model path stored in the checkpoint callback."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 22,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "'/home/nkulikov/RePlay/examples/sasrec/checkpoints/epoch=99-step=18900.ckpt'"
-      ]
-     },
-     "execution_count": 22,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "best_model_path = checkpoint_callback.best_model_path\n",
-    "best_model_path"
-   ]
-  },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Inference\n",
-    "\n",
-    "To obtain model scores, we will load the weights from the best checkpoint. To do this, we use the `LightningModule`, provide there the path to the checkpoint and the model instance."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 23,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import replay\n",
-    "torch.serialization.add_safe_globals([\n",
-    "    replay.nn.lightning.optimizer.OptimizerFactory,\n",
-    "    replay.nn.lightning.scheduler.LRSchedulerFactory\n",
-    "])"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 24,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "sasrec = make_sasrec(\n",
-    "    schema=tensor_schema,\n",
-    "    embedding_dim=EMBEDDING_DIM,\n",
-    "    max_sequence_length=MAX_SEQ_LEN,\n",
-    "    num_heads=NUM_HEADS,\n",
-    "    num_blocks=NUM_BLOCKS,\n",
-    "    dropout=DROPOUT,\n",
-    ")\n",
-    "\n",
-    "best_model = LightningModule.load_from_checkpoint(best_model_path, model=sasrec)\n",
-    "best_model.eval();"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Configure `ParquetModule` for inference"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 25,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "inference_metadata = {\"predict\": create_meta(shape=MAX_SEQ_LEN)}\n",
-    "\n",
-    "parquet_module = ParquetModule(\n",
-    "    predict_path=PREDICT_PATH,\n",
-    "    batch_size=BATCH_SIZE,\n",
-    "    metadata=inference_metadata,\n",
-    "    transforms=transforms,\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "During inference, we can use `TopItemsCallback`. Such callback allows you to get scores for each user throughout the entire catalog and get recommendations in the form of ids of items with the highest score values.\n",
-    "\n",
-    "\n",
-    "Recommendations can be fetched in four formats: PySpark DataFrame, Pandas DataFrame, Polars DataFrame or raw PyTorch tensors. Each of the types corresponds a callback. In this example, we'll be using the `PandasTopItemsCallback`."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 26,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "💡 Tip: For seamless cloud uploads and versioning, try installing [litmodels](https://pypi.org/project/litmodels/) to enable LitModelCheckpoint, which syncs automatically with the Lightning model registry.\n",
-      "GPU available: True (cuda), used: True\n",
-      "TPU available: False, using: 0 TPU cores\n",
-      "LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]\n"
-     ]
-    },
-    {
-     "data": {
-      "application/vnd.jupyter.widget-view+json": {
-       "model_id": "94a84cd4fca248e7b340909e2048450f",
-       "version_major": 2,
-       "version_minor": 0
-      },
-      "text/plain": [
-       "Predicting: |          | 0/? [00:00<?, ?it/s]"
-      ]
-     },
-     "metadata": {},
-     "output_type": "display_data"
-    }
-   ],
-   "source": [
-    "from replay.nn.lightning.callback import PandasTopItemsCallback\n",
-    "\n",
-    "csv_logger = CSVLogger(save_dir=\"sasrec/logs/test\", name=\"SasRec-example\")\n",
-    "\n",
-    "TOPK = [1, 5, 10, 20]\n",
-    "\n",
-    "pandas_prediction_callback = PandasTopItemsCallback(\n",
-    "    top_k=max(TOPK),\n",
-    "    query_column=\"user_id\",\n",
-    "    item_column=\"item_id\",\n",
-    "    rating_column=\"score\",\n",
-    ")\n",
-    "\n",
-    "trainer = L.Trainer(callbacks=[pandas_prediction_callback], logger=csv_logger, inference_mode=True)\n",
-    "trainer.predict(best_model, datamodule=parquet_module, return_predictions=False)\n",
-    "\n",
-    "pandas_res = pandas_prediction_callback.get_result()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 27,
-   "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>486</td>\n",
-       "      <td>7.196251</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>0</th>\n",
-       "      <td>0</td>\n",
-       "      <td>101</td>\n",
-       "      <td>7.086551</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>0</th>\n",
-       "      <td>0</td>\n",
-       "      <td>210</td>\n",
-       "      <td>6.872698</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>0</th>\n",
-       "      <td>0</td>\n",
-       "      <td>1613</td>\n",
-       "      <td>6.834656</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>0</th>\n",
-       "      <td>0</td>\n",
-       "      <td>357</td>\n",
-       "      <td>6.555361</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>...</th>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6037</th>\n",
-       "      <td>6039</td>\n",
-       "      <td>84</td>\n",
-       "      <td>11.528885</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6037</th>\n",
-       "      <td>6039</td>\n",
-       "      <td>2700</td>\n",
-       "      <td>11.400557</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6037</th>\n",
-       "      <td>6039</td>\n",
-       "      <td>2557</td>\n",
-       "      <td>11.273298</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6037</th>\n",
-       "      <td>6039</td>\n",
-       "      <td>2634</td>\n",
-       "      <td>11.121119</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6037</th>\n",
-       "      <td>6039</td>\n",
-       "      <td>2633</td>\n",
-       "      <td>11.119842</td>\n",
-       "    </tr>\n",
-       "  </tbody>\n",
-       "</table>\n",
-       "<p>120760 rows × 3 columns</p>\n",
-       "</div>"
-      ],
-      "text/plain": [
-       "      user_id item_id      score\n",
-       "0           0     486   7.196251\n",
-       "0           0     101   7.086551\n",
-       "0           0     210   6.872698\n",
-       "0           0    1613   6.834656\n",
-       "0           0     357   6.555361\n",
-       "...       ...     ...        ...\n",
-       "6037     6039      84  11.528885\n",
-       "6037     6039    2700  11.400557\n",
-       "6037     6039    2557  11.273298\n",
-       "6037     6039    2634  11.121119\n",
-       "6037     6039    2633  11.119842\n",
-       "\n",
-       "[120760 rows x 3 columns]"
-      ]
-     },
-     "execution_count": 27,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "pandas_res"
-   ]
-  },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### Calculating metrics\n",
-    "\n",
-    "*test_gt* is already encoded, so we can use it for computing metrics."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 28,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "from replay.metrics import MAP, OfflineMetrics, Precision, Recall\n",
-    "from replay.metrics.torch_metrics_builder import metrics_to_df"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 29,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "result_metrics = OfflineMetrics(\n",
-    "    [Recall(TOPK), Precision(TOPK), MAP(TOPK)],\n",
-    "    query_column=\"user_id\",\n",
-    "    rating_column=\"score\",\n",
-    ")(pandas_res, test_gt.explode(\"item_id\"))"
-   ]
-  },
-  {
-   "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>5</th>\n",
-       "    </tr>\n",
-       "  </thead>\n",
-       "  <tbody>\n",
-       "    <tr>\n",
-       "      <th>MAP</th>\n",
-       "      <td>0.016065</td>\n",
-       "      <td>0.056637</td>\n",
-       "      <td>0.064067</td>\n",
-       "      <td>0.045948</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>Precision</th>\n",
-       "      <td>0.016065</td>\n",
-       "      <td>0.018665</td>\n",
-       "      <td>0.014798</td>\n",
-       "      <td>0.021133</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>Recall</th>\n",
-       "      <td>0.016065</td>\n",
-       "      <td>0.186651</td>\n",
-       "      <td>0.295959</td>\n",
-       "      <td>0.105664</td>\n",
-       "    </tr>\n",
-       "  </tbody>\n",
-       "</table>\n",
-       "</div>"
-      ],
-      "text/plain": [
-       "k                 1        10        20         5\n",
-       "MAP        0.016065  0.056637  0.064067  0.045948\n",
-       "Precision  0.016065  0.018665  0.014798  0.021133\n",
-       "Recall     0.016065  0.186651  0.295959  0.105664"
-      ]
-     },
-     "execution_count": 30,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "metrics_to_df(result_metrics)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Let's call the `inverse_transform` encoder's function to get the final dataframe with recommendations"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 31,
-   "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>2012</td>\n",
-       "      <td>500</td>\n",
-       "      <td>7.196251</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>0</th>\n",
-       "      <td>2012</td>\n",
-       "      <td>104</td>\n",
-       "      <td>7.086551</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>0</th>\n",
-       "      <td>2012</td>\n",
-       "      <td>216</td>\n",
-       "      <td>6.872698</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>0</th>\n",
-       "      <td>2012</td>\n",
-       "      <td>1777</td>\n",
-       "      <td>6.834656</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>0</th>\n",
-       "      <td>2012</td>\n",
-       "      <td>367</td>\n",
-       "      <td>6.555361</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>...</th>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "      <td>...</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6037</th>\n",
-       "      <td>5727</td>\n",
-       "      <td>86</td>\n",
-       "      <td>11.528885</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6037</th>\n",
-       "      <td>5727</td>\n",
-       "      <td>2907</td>\n",
-       "      <td>11.400557</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6037</th>\n",
-       "      <td>5727</td>\n",
-       "      <td>2762</td>\n",
-       "      <td>11.273298</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6037</th>\n",
-       "      <td>5727</td>\n",
-       "      <td>2841</td>\n",
-       "      <td>11.121119</td>\n",
-       "    </tr>\n",
-       "    <tr>\n",
-       "      <th>6037</th>\n",
-       "      <td>5727</td>\n",
-       "      <td>2840</td>\n",
-       "      <td>11.119842</td>\n",
-       "    </tr>\n",
-       "  </tbody>\n",
-       "</table>\n",
-       "<p>120760 rows × 3 columns</p>\n",
-       "</div>"
-      ],
-      "text/plain": [
-       "      user_id  item_id      score\n",
-       "0        2012      500   7.196251\n",
-       "0        2012      104   7.086551\n",
-       "0        2012      216   6.872698\n",
-       "0        2012     1777   6.834656\n",
-       "0        2012      367   6.555361\n",
-       "...       ...      ...        ...\n",
-       "6037     5727       86  11.528885\n",
-       "6037     5727     2907  11.400557\n",
-       "6037     5727     2762  11.273298\n",
-       "6037     5727     2841  11.121119\n",
-       "6037     5727     2840  11.119842\n",
-       "\n",
-       "[120760 rows x 3 columns]"
-      ]
-     },
-     "execution_count": 31,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "encoder.inverse_transform(pandas_res)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": []
-  }
- ],
- "metadata": {
-  "kernelspec": {
-   "display_name": "new_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.12.3"
-  },
-  "orig_nbformat": 4
- },
- "nbformat": 4,
- "nbformat_minor": 2
-}

From 5e98b75dbe4d7bc1e3cb1f247ad14f2c0156d1ac Mon Sep 17 00:00:00 2001
From: Nikita Kulikov <neilkulikov@gmail.com>
Date: Mon, 2 Feb 2026 07:24:43 +0300
Subject: [PATCH 3/5] Documentation updated

---
 replay/nn/transform/__init__.py          |  5 +----
 replay/nn/transform/negative_sampling.py | 24 ++++++++++++++----------
 2 files changed, 15 insertions(+), 14 deletions(-)

diff --git a/replay/nn/transform/__init__.py b/replay/nn/transform/__init__.py
index 40c2cdec1..1baa08b32 100644
--- a/replay/nn/transform/__init__.py
+++ b/replay/nn/transform/__init__.py
@@ -1,14 +1,11 @@
 from .copy import CopyTransform
 from .grouping import GroupTransform
-<<<<<<< HEAD
-from .negative_sampling import MultiClassNegativeSamplingTransform, UniformNegativeSamplingTransform
-=======
 from .negative_sampling import (
+    MultiClassNegativeSamplingTransform,
     UniformNegativeSamplingTransform, 
     FrequencyNegativeSamplingTransform,
     ThresholdNegativeSamplingTransform,
 )
->>>>>>> 36b25601 (Gitignore updated to recent examples)
 from .next_token import NextTokenTransform
 from .rename import RenameTransform
 from .reshape import UnsqueezeTransform
diff --git a/replay/nn/transform/negative_sampling.py b/replay/nn/transform/negative_sampling.py
index 34debab8e..76aa815f3 100644
--- a/replay/nn/transform/negative_sampling.py
+++ b/replay/nn/transform/negative_sampling.py
@@ -101,8 +101,10 @@ class FrequencyNegativeSamplingTransform(torch.nn.Module):
     Transform for global negative sampling.
 
     For every batch, transform generates a vector of size ``(num_negative_samples)``
-    consisting of random indices sampeled from a range of ``cardinality``. Unless a custom sample
-    distribution is provided, the indices are weighted equally.
+    consisting of random indices sampeled from a range of ``cardinality``.
+
+    Indices frequency will be computed and their sampling will be done
+    according to their respective frequencies.
 
     Example:
 
@@ -110,7 +112,7 @@ class FrequencyNegativeSamplingTransform(torch.nn.Module):
 
         >>> _ = torch.manual_seed(0)
         >>> input_batch = {"item_id": torch.LongTensor([[1, 0, 4]])}
-        >>> transform = UniformNegativeSamplingTransform(cardinality=4, num_negative_samples=2)
+        >>> transform = FrequencyNegativeSamplingTransform(cardinality=4, num_negative_samples=2)
         >>> output_batch = transform(input_batch)
         >>> output_batch
         {'item_id': tensor([[1, 0, 4]]), 'negative_labels': tensor([2, 1])}
@@ -130,10 +132,10 @@ def __init__(
         :param cardinality: The size of sample vocabulary.
         :param num_negative_samples: The size of negatives vector to generate.
         :param out_feature_name: The name of result feature in batch.
-        :param sample_distribution: The weighs of indices in the vocabulary. If specified, must
-                match the ``cardinality``. Default: ``None``.
         :param generator: Random number generator to be used for sampling
                 from the distribution. Default: ``None``.
+        :param mode: Mode of frequency-based samping for undersampled items.
+            Default: ``softmax``.
         """
         assert num_negative_samples < cardinality
 
@@ -187,8 +189,10 @@ class ThresholdNegativeSamplingTransform(torch.nn.Module):
     Transform for global negative sampling.
 
     For every batch, transform generates a vector of size ``(num_negative_samples)``
-    consisting of random indices sampeled from a range of ``cardinality``. Unless a custom sample
-    distribution is provided, the indices are weighted equally.
+    consisting of random indices sampeled from a range of ``cardinality``.
+
+    Indices that are oversampled at this point will be ignored, while
+    other samples will be chosen according to their respective frequency.
 
     Example:
 
@@ -196,7 +200,7 @@ class ThresholdNegativeSamplingTransform(torch.nn.Module):
 
         >>> _ = torch.manual_seed(0)
         >>> input_batch = {"item_id": torch.LongTensor([[1, 0, 4]])}
-        >>> transform = UniformNegativeSamplingTransform(cardinality=4, num_negative_samples=2)
+        >>> transform = ThresholdNegativeSamplingTransform(cardinality=4, num_negative_samples=2)
         >>> output_batch = transform(input_batch)
         >>> output_batch
         {'item_id': tensor([[1, 0, 4]]), 'negative_labels': tensor([2, 1])}
@@ -216,10 +220,10 @@ def __init__(
         :param cardinality: The size of sample vocabulary.
         :param num_negative_samples: The size of negatives vector to generate.
         :param out_feature_name: The name of result feature in batch.
-        :param sample_distribution: The weighs of indices in the vocabulary. If specified, must
-                match the ``cardinality``. Default: ``None``.
         :param generator: Random number generator to be used for sampling
                 from the distribution. Default: ``None``.
+        :param mode: Mode of frequency-based samping for undersampled items.
+            Default: ``softmax``.
         """
         assert num_negative_samples < cardinality
 

From bd614aaa61cbba62606b7dc54bee9a3f0222cd66 Mon Sep 17 00:00:00 2001
From: Nikita Kulikov <neilkulikov@gmail.com>
Date: Mon, 2 Feb 2026 07:30:20 +0300
Subject: [PATCH 4/5] Formatted and checked

---
 CONTRIBUTING.md                          |  3 +-
 replay/nn/transform/__init__.py          |  6 +--
 replay/nn/transform/negative_sampling.py | 51 ++++++++++--------------
 3 files changed, 26 insertions(+), 34 deletions(-)

diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index 638adc03f..236d90953 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -136,8 +136,7 @@ You can just get pyproject.toml file from step 6, to start using linters and for
 
 In order to automate checking of the code quality, please run:
 ```bash
-poetry run ruff check .
-poetry run black --check --diff -- .
+./poetry_wrapper.sh run ruff check .
 ./poetry_wrapper.sh check
 ./poetry_wrapper.sh --experimental check
 ```
diff --git a/replay/nn/transform/__init__.py b/replay/nn/transform/__init__.py
index 1baa08b32..4f47dba4b 100644
--- a/replay/nn/transform/__init__.py
+++ b/replay/nn/transform/__init__.py
@@ -2,9 +2,9 @@
 from .grouping import GroupTransform
 from .negative_sampling import (
     MultiClassNegativeSamplingTransform,
-    UniformNegativeSamplingTransform, 
     FrequencyNegativeSamplingTransform,
     ThresholdNegativeSamplingTransform,
+    UniformNegativeSamplingTransform,
 )
 from .next_token import NextTokenTransform
 from .rename import RenameTransform
@@ -15,15 +15,15 @@
 
 __all__ = [
     "CopyTransform",
+    "FrequencyNegativeSamplingTransform",
     "GroupTransform",
     "MultiClassNegativeSamplingTransform",
     "NextTokenTransform",
     "RenameTransform",
     "SequenceRollTransform",
+    "ThresholdNegativeSamplingTransform",
     "TokenMaskTransform",
     "TrimTransform",
     "UniformNegativeSamplingTransform",
-    "FrequencyNegativeSamplingTransform",
-    "ThresholdNegativeSamplingTransform",
     "UnsqueezeTransform",
 ]
diff --git a/replay/nn/transform/negative_sampling.py b/replay/nn/transform/negative_sampling.py
index 76aa815f3..c1e9bd5b2 100644
--- a/replay/nn/transform/negative_sampling.py
+++ b/replay/nn/transform/negative_sampling.py
@@ -1,10 +1,9 @@
-from typing import Optional, Literal, cast
+import warnings
+from typing import Literal, Optional, cast
 
 import torch
 import torch.nn.functional as func
 
-import warnings
-
 
 class UniformNegativeSamplingTransform(torch.nn.Module):
     """
@@ -73,11 +72,7 @@ def __init__(
         self.out_feature_name = out_feature_name
         self.num_negative_samples = num_negative_samples
         self.generator = generator
-        if sample_distribution is not None:
-            sample_distribution = sample_distribution
-        else:
-            sample_distribution = torch.ones(cardinality)
-        
+        sample_distribution = sample_distribution if sample_distribution is not None else torch.ones(cardinality)
         self.sample_distribution = torch.nn.Buffer(
             cast(torch.Tensor, sample_distribution)
         )
@@ -153,16 +148,15 @@ def __init__(
 
     def get_probas(self) -> torch.Tensor:
         raw: torch.Tensor = 1.0 / (1.0 + self.frequencies)
-        match self.mode:
-            case "softsum":
-                result: torch.Tensor = raw / torch.sum(raw)
-            case "softmax":
+        if self.mode == "softsum":
+            result: torch.Tensor = raw / torch.sum(raw)
+        elif self.mode == "softmax":
                 result: torch.Tensor = func.softmax(raw, dim = -1)
-            case _:
-                msg: str = f"Unsupported mode: {self.mode}."
-                raise TypeError(msg)
+        else:
+            msg: str = f"Unsupported mode: {self.mode}."
+            raise TypeError(msg)
         return result
-    
+
     def update_probas(self, selected: torch.Tensor) -> None:
         device = self.frequencies.device
         one = torch.ones(1, dtype = torch.int64, device = device)
@@ -243,20 +237,19 @@ def get_probas(self) -> torch.Tensor:
         raw: torch.Tensor = 1.0 / (1.0 + self.frequencies)
         thr: torch.Tensor = torch.max(self.frequencies)
         mask: torch.Tensor = thr != self.frequencies
-        match self.mode:
-            case "softsum":
-                eps = torch.finfo(raw.dtype).eps
-                raw = torch.where(mask, raw, eps)
-                result: torch.Tensor = raw / torch.sum(raw)
-            case "softmax":
-                inf = torch.finfo(raw.dtype).min
-                raw = torch.where(mask, raw, inf)
-                result: torch.Tensor = func.softmax(raw, dim = -1)
-            case _:
-                msg: str = f"Unsupported mode: {self.mode}."
-                raise TypeError(msg)
+        if self.mode == "softsum":
+            eps = torch.finfo(raw.dtype).eps
+            raw = torch.where(mask, raw, eps)
+            result: torch.Tensor = raw / torch.sum(raw)
+        elif self.mode == "softmax":
+            inf = torch.finfo(raw.dtype).min
+            raw = torch.where(mask, raw, inf)
+            result: torch.Tensor = func.softmax(raw, dim = -1)
+        else:
+            msg: str = f"Unsupported mode: {self.mode}."
+            raise TypeError(msg)
         return result
-    
+
     def update_probas(self, selected: torch.Tensor) -> None:
         device = self.frequencies.device
         one = torch.ones(1, dtype = torch.int64, device = device)

From 3e03d8e206364b3d7e619ecbf29bd8e63a858222 Mon Sep 17 00:00:00 2001
From: Nikita Kulikov <neilkulikov@gmail.com>
Date: Mon, 2 Feb 2026 08:13:09 +0300
Subject: [PATCH 5/5] Ruff format applied

---
 replay/nn/transform/__init__.py          |  2 +-
 replay/nn/transform/negative_sampling.py | 33 ++++++++++--------------
 2 files changed, 14 insertions(+), 21 deletions(-)

diff --git a/replay/nn/transform/__init__.py b/replay/nn/transform/__init__.py
index 4f47dba4b..0ae787bf3 100644
--- a/replay/nn/transform/__init__.py
+++ b/replay/nn/transform/__init__.py
@@ -1,8 +1,8 @@
 from .copy import CopyTransform
 from .grouping import GroupTransform
 from .negative_sampling import (
-    MultiClassNegativeSamplingTransform,
     FrequencyNegativeSamplingTransform,
+    MultiClassNegativeSamplingTransform,
     ThresholdNegativeSamplingTransform,
     UniformNegativeSamplingTransform,
 )
diff --git a/replay/nn/transform/negative_sampling.py b/replay/nn/transform/negative_sampling.py
index c1e9bd5b2..ef68936a3 100644
--- a/replay/nn/transform/negative_sampling.py
+++ b/replay/nn/transform/negative_sampling.py
@@ -48,8 +48,7 @@ def __init__(
         if sample_distribution is not None:
             if sample_distribution.ndim != 1:
                 msg: str = (
-                    "The `sample_distribution` parameter must be 1D."
-                    f"Got {sample_distribution.ndim}, will be flattened."
+                    f"The `sample_distribution` parameter must be 1D.Got {sample_distribution.ndim}, will be flattened."
                 )
                 warnings.warn(msg)
                 sample_distribution = sample_distribution.flatten()
@@ -73,9 +72,7 @@ def __init__(
         self.num_negative_samples = num_negative_samples
         self.generator = generator
         sample_distribution = sample_distribution if sample_distribution is not None else torch.ones(cardinality)
-        self.sample_distribution = torch.nn.Buffer(
-            cast(torch.Tensor, sample_distribution)
-        )
+        self.sample_distribution = torch.nn.Buffer(cast(torch.Tensor, sample_distribution))
 
     def forward(self, batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
         output_batch = dict(batch.items())
@@ -91,6 +88,7 @@ def forward(self, batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
         output_batch[self.out_feature_name] = negatives.to(device)
         return output_batch
 
+
 class FrequencyNegativeSamplingTransform(torch.nn.Module):
     """
     Transform for global negative sampling.
@@ -142,16 +140,14 @@ def __init__(
         self.generator = generator
         self.mode = mode
 
-        self.frequencies = torch.nn.Buffer(
-            torch.zeros(cardinality, dtype = torch.int64)
-        )
+        self.frequencies = torch.nn.Buffer(torch.zeros(cardinality, dtype=torch.int64))
 
     def get_probas(self) -> torch.Tensor:
         raw: torch.Tensor = 1.0 / (1.0 + self.frequencies)
         if self.mode == "softsum":
             result: torch.Tensor = raw / torch.sum(raw)
         elif self.mode == "softmax":
-                result: torch.Tensor = func.softmax(raw, dim = -1)
+            result: torch.Tensor = func.softmax(raw, dim=-1)
         else:
             msg: str = f"Unsupported mode: {self.mode}."
             raise TypeError(msg)
@@ -159,14 +155,14 @@ def get_probas(self) -> torch.Tensor:
 
     def update_probas(self, selected: torch.Tensor) -> None:
         device = self.frequencies.device
-        one = torch.ones(1, dtype = torch.int64, device = device)
+        one = torch.ones(1, dtype=torch.int64, device=device)
         self.frequencies.index_add_(-1, selected, one.expand(selected.numel()))
 
     def forward(self, batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
         output_batch = dict(batch.items())
 
         negatives = torch.multinomial(
-            input = self.get_probas(),
+            input=self.get_probas(),
             num_samples=self.num_negative_samples,
             replacement=False,
             generator=self.generator,
@@ -178,6 +174,7 @@ def forward(self, batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
         output_batch[self.out_feature_name] = negatives.to(device)
         return output_batch
 
+
 class ThresholdNegativeSamplingTransform(torch.nn.Module):
     """
     Transform for global negative sampling.
@@ -229,9 +226,7 @@ def __init__(
         self.generator = generator
         self.mode = mode
 
-        self.frequencies = torch.nn.Buffer(
-            torch.zeros(cardinality, dtype = torch.int64)
-        )
+        self.frequencies = torch.nn.Buffer(torch.zeros(cardinality, dtype=torch.int64))
 
     def get_probas(self) -> torch.Tensor:
         raw: torch.Tensor = 1.0 / (1.0 + self.frequencies)
@@ -244,7 +239,7 @@ def get_probas(self) -> torch.Tensor:
         elif self.mode == "softmax":
             inf = torch.finfo(raw.dtype).min
             raw = torch.where(mask, raw, inf)
-            result: torch.Tensor = func.softmax(raw, dim = -1)
+            result: torch.Tensor = func.softmax(raw, dim=-1)
         else:
             msg: str = f"Unsupported mode: {self.mode}."
             raise TypeError(msg)
@@ -252,14 +247,14 @@ def get_probas(self) -> torch.Tensor:
 
     def update_probas(self, selected: torch.Tensor) -> None:
         device = self.frequencies.device
-        one = torch.ones(1, dtype = torch.int64, device = device)
+        one = torch.ones(1, dtype=torch.int64, device=device)
         self.frequencies.index_add_(-1, selected, one.expand(selected.numel()))
 
     def forward(self, batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
         output_batch = dict(batch.items())
 
         negatives = torch.multinomial(
-            input = self.get_probas(),
+            input=self.get_probas(),
             num_samples=self.num_negative_samples,
             replacement=False,
             generator=self.generator,
@@ -344,9 +339,7 @@ def __init__(
 
         super().__init__()
 
-        self.sample_mask = torch.nn.Buffer(
-            sample_mask.float()
-        )
+        self.sample_mask = torch.nn.Buffer(sample_mask.float())
 
         self.num_negative_samples = num_negative_samples
         self.negative_selector_name = negative_selector_name