examples/05-Retrieval-Model.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "63959d5a",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Copyright 2021 NVIDIA Corporation. All Rights Reserved.\n",
    "#\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "#     http://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License.\n",
    "# ================================\n",
    "\n",
    "# Each user is responsible for checking the content of datasets and the\n",
    "# applicable licenses and determining if suitable for the intended use."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "906fc19a",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "<img src=\"https://developer.download.nvidia.com/notebooks/dlsw-notebooks/merlin_models_05-retrieval-model/nvidia_logo.png\" style=\"width: 90px; float: right;\">\n",
    "\n",
    "# Building a Retrieval Model with Merlin Models\n",
    "\n",
    "This notebook is created using the latest stable [merlin-tensorflow](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/merlin/containers/merlin-tensorflow/tags) container. \n",
    "\n",
    "In large scale recommender systems pipelines, the size of the item catalog (number of unique items) might be in the order of millions. At such scale, a typical setup is having two-stage pipeline, where a faster candidate retrieval model quickly extracts thousands of relevant items and a then a more powerful ranking model (i.e. with more features and more powerful architecture) ranks the top-k items that are going to be displayed to the user. For ML-based candidate retrieval model, as it needs to quickly score millions of items for a given user, a popular choices are models that can produce recommendation scores by just computing the dot product the user embeddings and item embeddings. Popular choices of such models are **Matrix Factorization**, which learns low-rank user and item embeddings, and the **Two-Tower architecture**, which is a neural network with two MLP towers where both user and item features are fed to generate user and item embeddings in the output."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "401daa6d",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "\n",
    "### Dataset\n",
    "\n",
    "In this notebook, we are building a Two-Tower model for Item Retrieval task using synthetic datasets that are mimicking the real [Ali-CCP: Alibaba Click and Conversion Prediction](https://tianchi.aliyun.com/dataset/dataDetail?dataId=408#1) dataset.\n",
    "### Learning objectives\n",
    "- Preparing the data with NVTabular\n",
    "- Training and evaluating Two-Tower model with Merlin Models\n",
    "- Exporting the model for deployment\n",
    "- Generating the top K recommendations from the trained model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cdac1fea",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Importing Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "92aa8daa",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/lib/python3/dist-packages/requests/__init__.py:89: RequestsDependencyWarning: urllib3 (1.26.12) or chardet (3.0.4) doesn't match a supported version!\n",
      "  warnings.warn(\"urllib3 ({}) or chardet ({}) doesn't match a supported \"\n",
      "2022-11-30 18:38:41.379796: I tensorflow/core/platform/cpu_feature_guard.cc:194] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  SSE3 SSE4.1 SSE4.2 AVX\n",
      "To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.\n",
      "2022-11-30 18:38:42.590379: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1532] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 16254 MB memory:  -> device: 0, name: Quadro GV100, pci bus id: 0000:15:00.0, compute capability: 7.0\n",
      "/usr/local/lib/python3.8/dist-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "\n",
    "import nvtabular as nvt\n",
    "from nvtabular.ops import *\n",
    "from merlin.models.utils.example_utils import workflow_fit_transform\n",
    "\n",
    "from merlin.schema.tags import Tags\n",
    "\n",
    "import merlin.models.tf as mm\n",
    "from merlin.io.dataset import Dataset\n",
    "from merlin.models.utils.dataset import unique_rows_by_features\n",
    "\n",
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "126c8bb8",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# disable INFO and DEBUG logging everywhere\n",
    "import logging\n",
    "\n",
    "logging.disable(logging.WARNING)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ac0eef4",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### Feature Engineering with NVTabular"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d21e559b",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Let's generate synthetic train and validation dataset objects."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "870acd78",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.8/dist-packages/merlin/schema/tags.py:148: UserWarning: Compound tags like Tags.USER_ID have been deprecated and will be removed in a future version. Please use the atomic versions of these tags, like [<Tags.USER: 'user'>, <Tags.ID: 'id'>].\n",
      "  warnings.warn(\n",
      "/usr/local/lib/python3.8/dist-packages/merlin/schema/tags.py:148: UserWarning: Compound tags like Tags.ITEM_ID have been deprecated and will be removed in a future version. Please use the atomic versions of these tags, like [<Tags.ITEM: 'item'>, <Tags.ID: 'id'>].\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "from merlin.datasets.synthetic import generate_data\n",
    "\n",
    "DATA_FOLDER = os.environ.get(\"DATA_FOLDER\", \"/workspace/data/\")\n",
    "NUM_ROWS = os.environ.get(\"NUM_ROWS\", 1000000)\n",
    "SYNTHETIC_DATA = eval(os.environ.get(\"SYNTHETIC_DATA\", \"True\"))\n",
    "\n",
    "if SYNTHETIC_DATA:\n",
    "    train, valid = generate_data(\"aliccp-raw\", int(NUM_ROWS), set_sizes=(0.7, 0.3))\n",
    "else:\n",
    "    train = nvt.Dataset(DATA_FOLDER + \"/train/*.parquet\")\n",
    "    valid = nvt.Dataset(DATA_FOLDER + \"/valid/*.parquet\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "1dd757a3-2f93-495f-b2c5-5098c997d46e",
   "metadata": {},
   "outputs": [],
   "source": [
    "train = train.to_ddf().compute()\n",
    "valid = valid.to_ddf().compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f7b7387-7c6c-48c0-82d8-d39685f99e4c",
   "metadata": {},
   "source": [
    "We keep only positive interactions where clicks==1 in the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4a0454c8-8eec-47cc-82e8-b45637b76d35",
   "metadata": {},
   "outputs": [],
   "source": [
    "train = train.loc[train['click']==1].reset_index(drop=True)\n",
    "valid = valid.loc[valid['click']==1].reset_index(drop=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1787e97-e7c2-4f2b-a728-6329b7d6a8c2",
   "metadata": {},
   "source": [
    "We can drop the target column since in this example we will only use positive interactions and then generate negative samples via negative sampling technique."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "202d150f-1a82-4001-b36b-8ef226222c52",
   "metadata": {},
   "outputs": [],
   "source": [
    "train = train.drop(['click', 'conversion'], axis=1)\n",
    "valid = valid.drop(['click', 'conversion'], axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "901b8e6d-0db4-4010-ac98-fbd2ec6a3b4b",
   "metadata": {},
   "source": [
    "Create Dataset objects"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "c6bd9c08-42c6-4aaa-9de5-494e0d5e8f0b",
   "metadata": {},
   "outputs": [],
   "source": [
    "train = Dataset(train)\n",
    "valid = Dataset(valid)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce4ef6b3-cb4c-4f8a-ad7d-b1a81ea4e1ec",
   "metadata": {},
   "source": [
    "Define output path for the processed parquet files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "f23232bf-7822-498f-a036-23e44ef48e50",
   "metadata": {},
   "outputs": [],
   "source": [
    "output_path = os.path.join(DATA_FOLDER, \"processed\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "078593b4",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "category_temp_directory = os.path.join(DATA_FOLDER, \"categories\")\n",
    "user_id = [\"user_id\"] >> Categorify(out_path=category_temp_directory) >> TagAsUserID()\n",
    "item_id = [\"item_id\"] >> Categorify(out_path=category_temp_directory) >> TagAsItemID()\n",
    "\n",
    "item_features = [\"item_category\", \"item_shop\", \"item_brand\"] >> Categorify(out_path=category_temp_directory) >> TagAsItemFeatures()\n",
    "\n",
    "user_features = (\n",
    "    [\n",
    "        \"user_shops\",\n",
    "        \"user_profile\",\n",
    "        \"user_group\",\n",
    "        \"user_gender\",\n",
    "        \"user_age\",\n",
    "        \"user_consumption_2\",\n",
    "        \"user_is_occupied\",\n",
    "        \"user_geography\",\n",
    "        \"user_intentions\",\n",
    "        \"user_brands\",\n",
    "        \"user_categories\",\n",
    "    ]\n",
    "    >> Categorify(out_path=category_temp_directory)\n",
    "    >> TagAsUserFeatures()\n",
    ")\n",
    "\n",
    "outputs = user_id + item_id + item_features + user_features"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2aa99f0b",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "With `transform_aliccp` function, we can execute fit() and transform() on the raw dataset applying the operators defined in the NVTabular workflow pipeline above. The processed parquet files are saved to output_path."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "07ab4edc",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from merlin.datasets.ecommerce import transform_aliccp\n",
    "\n",
    "transform_aliccp((train, valid), output_path, nvt_workflow=outputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df7be4a6",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Building a Two-Tower Model with Merlin Models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7f7e4a6b",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "We will use Two-Tower Model for item retrieval task. Real-world large scale recommender systems have hundreds of millions of items (products) and users. Thus, these systems often composed of two stages: candidate generation (retrieval) and ranking (scoring the retrieved items). At candidate generation step, a subset of relevant items from large item corpus is retrieved. You can read more about two stage Recommender Systems here. In this example, we're going to focus on the retrieval stage.\n",
    "\n",
    "A Two-Tower Model consists of item (candidate) and user (query) encoder towers. With two towers, the model can learn representations (embeddings) for queries and candidates separately. \n",
    "\n",
    "<img src=\"./images/TwoTower.png\"  width=\"30%\">\n",
    "\n",
    "Image Adapted from: [Off-policy Learning in Two-stage Recommender Systems](https://dl.acm.org/doi/abs/10.1145/3366423.3380130)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb5f26a8",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "We use the `schema` object to define our model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "a7e0c257",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "train = Dataset(os.path.join(output_path, \"train\", \"*.parquet\"))\n",
    "valid = Dataset(os.path.join(output_path, \"valid\", \"*.parquet\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fcdb2dbc-9bbd-46ad-aede-44c3ffbaac9d",
   "metadata": {},
   "source": [
    "Select features with user and item tags, and be sure to exclude target column."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "94eb4048",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "schema = train.schema.select_by_tag([Tags.ITEM_ID, Tags.USER_ID, Tags.ITEM, Tags.USER])\n",
    "train.schema = schema\n",
    "valid.schema = schema"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3e3469c",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "We can print out the feature column names."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "8b18bcf3-65c8-418a-aa43-55abb73a4542",
   "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>name</th>\n",
       "      <th>tags</th>\n",
       "      <th>dtype</th>\n",
       "      <th>is_list</th>\n",
       "      <th>is_ragged</th>\n",
       "      <th>properties.num_buckets</th>\n",
       "      <th>properties.freq_threshold</th>\n",
       "      <th>properties.max_size</th>\n",
       "      <th>properties.start_index</th>\n",
       "      <th>properties.cat_path</th>\n",
       "      <th>properties.embedding_sizes.cardinality</th>\n",
       "      <th>properties.embedding_sizes.dimension</th>\n",
       "      <th>properties.domain.min</th>\n",
       "      <th>properties.domain.max</th>\n",
       "      <th>properties.domain.name</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>user_id</td>\n",
       "      <td>(Tags.USER, Tags.ID, Tags.CATEGORICAL, Tags.US...</td>\n",
       "      <td>int64</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>None</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>/workspace/data/categories/categories/unique.u...</td>\n",
       "      <td>663.0</td>\n",
       "      <td>61.0</td>\n",
       "      <td>0</td>\n",
       "      <td>662</td>\n",
       "      <td>user_id</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>item_id</td>\n",
       "      <td>(Tags.ID, Tags.ITEM, Tags.CATEGORICAL, Tags.IT...</td>\n",
       "      <td>int64</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>None</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>/workspace/data/categories/categories/unique.i...</td>\n",
       "      <td>662.0</td>\n",
       "      <td>61.0</td>\n",
       "      <td>0</td>\n",
       "      <td>661</td>\n",
       "      <td>item_id</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>item_category</td>\n",
       "      <td>(Tags.CATEGORICAL, Tags.ITEM)</td>\n",
       "      <td>int64</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>None</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>/workspace/data/categories/categories/unique.i...</td>\n",
       "      <td>662.0</td>\n",
       "      <td>61.0</td>\n",
       "      <td>0</td>\n",
       "      <td>661</td>\n",
       "      <td>item_category</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>item_shop</td>\n",
       "      <td>(Tags.CATEGORICAL, Tags.ITEM)</td>\n",
       "      <td>int64</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>None</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>/workspace/data/categories/categories/unique.i...</td>\n",
       "      <td>662.0</td>\n",
       "      <td>61.0</td>\n",
       "      <td>0</td>\n",
       "      <td>661</td>\n",
       "      <td>item_shop</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>item_brand</td>\n",
       "      <td>(Tags.CATEGORICAL, Tags.ITEM)</td>\n",
       "      <td>int64</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>None</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>/workspace/data/categories/categories/unique.i...</td>\n",
       "      <td>662.0</td>\n",
       "      <td>61.0</td>\n",
       "      <td>0</td>\n",
       "      <td>661</td>\n",
       "      <td>item_brand</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>user_shops</td>\n",
       "      <td>(Tags.USER, Tags.CATEGORICAL)</td>\n",
       "      <td>int64</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>None</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>/workspace/data/categories/categories/unique.u...</td>\n",
       "      <td>663.0</td>\n",
       "      <td>61.0</td>\n",
       "      <td>0</td>\n",
       "      <td>662</td>\n",
       "      <td>user_shops</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>user_profile</td>\n",
       "      <td>(Tags.USER, Tags.CATEGORICAL)</td>\n",
       "      <td>int64</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>None</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>/workspace/data/categories/categories/unique.u...</td>\n",
       "      <td>48.0</td>\n",
       "      <td>16.0</td>\n",
       "      <td>0</td>\n",
       "      <td>47</td>\n",
       "      <td>user_profile</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>user_group</td>\n",
       "      <td>(Tags.USER, Tags.CATEGORICAL)</td>\n",
       "      <td>int64</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>None</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>/workspace/data/categories/categories/unique.u...</td>\n",
       "      <td>10.0</td>\n",
       "      <td>16.0</td>\n",
       "      <td>0</td>\n",
       "      <td>9</td>\n",
       "      <td>user_group</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>user_gender</td>\n",
       "      <td>(Tags.USER, Tags.CATEGORICAL)</td>\n",
       "      <td>int64</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>None</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>/workspace/data/categories/categories/unique.u...</td>\n",
       "      <td>3.0</td>\n",
       "      <td>16.0</td>\n",
       "      <td>0</td>\n",
       "      <td>2</td>\n",
       "      <td>user_gender</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>user_age</td>\n",
       "      <td>(Tags.USER, Tags.CATEGORICAL)</td>\n",
       "      <td>int64</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>None</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>/workspace/data/categories/categories/unique.u...</td>\n",
       "      <td>7.0</td>\n",
       "      <td>16.0</td>\n",
       "      <td>0</td>\n",
       "      <td>6</td>\n",
       "      <td>user_age</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>user_consumption_2</td>\n",
       "      <td>(Tags.USER, Tags.CATEGORICAL)</td>\n",
       "      <td>int64</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>None</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>/workspace/data/categories/categories/unique.u...</td>\n",
       "      <td>4.0</td>\n",
       "      <td>16.0</td>\n",
       "      <td>0</td>\n",
       "      <td>3</td>\n",
       "      <td>user_consumption_2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>user_is_occupied</td>\n",
       "      <td>(Tags.USER, Tags.CATEGORICAL)</td>\n",
       "      <td>int64</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>None</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>/workspace/data/categories/categories/unique.u...</td>\n",
       "      <td>3.0</td>\n",
       "      <td>16.0</td>\n",
       "      <td>0</td>\n",
       "      <td>2</td>\n",
       "      <td>user_is_occupied</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>user_geography</td>\n",
       "      <td>(Tags.USER, Tags.CATEGORICAL)</td>\n",
       "      <td>int64</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>None</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>/workspace/data/categories/categories/unique.u...</td>\n",
       "      <td>5.0</td>\n",
       "      <td>16.0</td>\n",
       "      <td>0</td>\n",
       "      <td>4</td>\n",
       "      <td>user_geography</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>user_intentions</td>\n",
       "      <td>(Tags.USER, Tags.CATEGORICAL)</td>\n",
       "      <td>int64</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>None</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>/workspace/data/categories/categories/unique.u...</td>\n",
       "      <td>663.0</td>\n",
       "      <td>61.0</td>\n",
       "      <td>0</td>\n",
       "      <td>662</td>\n",
       "      <td>user_intentions</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14</th>\n",
       "      <td>user_brands</td>\n",
       "      <td>(Tags.USER, Tags.CATEGORICAL)</td>\n",
       "      <td>int64</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>None</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>/workspace/data/categories/categories/unique.u...</td>\n",
       "      <td>663.0</td>\n",
       "      <td>61.0</td>\n",
       "      <td>0</td>\n",
       "      <td>662</td>\n",
       "      <td>user_brands</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15</th>\n",
       "      <td>user_categories</td>\n",
       "      <td>(Tags.USER, Tags.CATEGORICAL)</td>\n",
       "      <td>int64</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>None</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>/workspace/data/categories/categories/unique.u...</td>\n",
       "      <td>663.0</td>\n",
       "      <td>61.0</td>\n",
       "      <td>0</td>\n",
       "      <td>662</td>\n",
       "      <td>user_categories</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "[{'name': 'user_id', 'tags': {<Tags.USER: 'user'>, <Tags.ID: 'id'>, <Tags.CATEGORICAL: 'categorical'>, <Tags.USER_ID: 'user_id'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/workspace/data/categories/categories/unique.user_id.parquet', 'embedding_sizes': {'cardinality': 663.0, 'dimension': 61.0}, 'domain': {'min': 0, 'max': 662, 'name': 'user_id'}}, 'dtype': dtype('int64'), 'is_list': False, 'is_ragged': False}, {'name': 'item_id', 'tags': {<Tags.ID: 'id'>, <Tags.ITEM: 'item'>, <Tags.CATEGORICAL: 'categorical'>, <Tags.ITEM_ID: 'item_id'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/workspace/data/categories/categories/unique.item_id.parquet', 'embedding_sizes': {'cardinality': 662.0, 'dimension': 61.0}, 'domain': {'min': 0, 'max': 661, 'name': 'item_id'}}, 'dtype': dtype('int64'), 'is_list': False, 'is_ragged': False}, {'name': 'item_category', 'tags': {<Tags.CATEGORICAL: 'categorical'>, <Tags.ITEM: 'item'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/workspace/data/categories/categories/unique.item_category.parquet', 'embedding_sizes': {'cardinality': 662.0, 'dimension': 61.0}, 'domain': {'min': 0, 'max': 661, 'name': 'item_category'}}, 'dtype': dtype('int64'), 'is_list': False, 'is_ragged': False}, {'name': 'item_shop', 'tags': {<Tags.CATEGORICAL: 'categorical'>, <Tags.ITEM: 'item'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/workspace/data/categories/categories/unique.item_shop.parquet', 'embedding_sizes': {'cardinality': 662.0, 'dimension': 61.0}, 'domain': {'min': 0, 'max': 661, 'name': 'item_shop'}}, 'dtype': dtype('int64'), 'is_list': False, 'is_ragged': False}, {'name': 'item_brand', 'tags': {<Tags.CATEGORICAL: 'categorical'>, <Tags.ITEM: 'item'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/workspace/data/categories/categories/unique.item_brand.parquet', 'embedding_sizes': {'cardinality': 662.0, 'dimension': 61.0}, 'domain': {'min': 0, 'max': 661, 'name': 'item_brand'}}, 'dtype': dtype('int64'), 'is_list': False, 'is_ragged': False}, {'name': 'user_shops', 'tags': {<Tags.USER: 'user'>, <Tags.CATEGORICAL: 'categorical'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/workspace/data/categories/categories/unique.user_shops.parquet', 'embedding_sizes': {'cardinality': 663.0, 'dimension': 61.0}, 'domain': {'min': 0, 'max': 662, 'name': 'user_shops'}}, 'dtype': dtype('int64'), 'is_list': False, 'is_ragged': False}, {'name': 'user_profile', 'tags': {<Tags.USER: 'user'>, <Tags.CATEGORICAL: 'categorical'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/workspace/data/categories/categories/unique.user_profile.parquet', 'embedding_sizes': {'cardinality': 48.0, 'dimension': 16.0}, 'domain': {'min': 0, 'max': 47, 'name': 'user_profile'}}, 'dtype': dtype('int64'), 'is_list': False, 'is_ragged': False}, {'name': 'user_group', 'tags': {<Tags.USER: 'user'>, <Tags.CATEGORICAL: 'categorical'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/workspace/data/categories/categories/unique.user_group.parquet', 'embedding_sizes': {'cardinality': 10.0, 'dimension': 16.0}, 'domain': {'min': 0, 'max': 9, 'name': 'user_group'}}, 'dtype': dtype('int64'), 'is_list': False, 'is_ragged': False}, {'name': 'user_gender', 'tags': {<Tags.USER: 'user'>, <Tags.CATEGORICAL: 'categorical'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/workspace/data/categories/categories/unique.user_gender.parquet', 'embedding_sizes': {'cardinality': 3.0, 'dimension': 16.0}, 'domain': {'min': 0, 'max': 2, 'name': 'user_gender'}}, 'dtype': dtype('int64'), 'is_list': False, 'is_ragged': False}, {'name': 'user_age', 'tags': {<Tags.USER: 'user'>, <Tags.CATEGORICAL: 'categorical'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/workspace/data/categories/categories/unique.user_age.parquet', 'embedding_sizes': {'cardinality': 7.0, 'dimension': 16.0}, 'domain': {'min': 0, 'max': 6, 'name': 'user_age'}}, 'dtype': dtype('int64'), 'is_list': False, 'is_ragged': False}, {'name': 'user_consumption_2', 'tags': {<Tags.USER: 'user'>, <Tags.CATEGORICAL: 'categorical'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/workspace/data/categories/categories/unique.user_consumption_2.parquet', 'embedding_sizes': {'cardinality': 4.0, 'dimension': 16.0}, 'domain': {'min': 0, 'max': 3, 'name': 'user_consumption_2'}}, 'dtype': dtype('int64'), 'is_list': False, 'is_ragged': False}, {'name': 'user_is_occupied', 'tags': {<Tags.USER: 'user'>, <Tags.CATEGORICAL: 'categorical'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/workspace/data/categories/categories/unique.user_is_occupied.parquet', 'embedding_sizes': {'cardinality': 3.0, 'dimension': 16.0}, 'domain': {'min': 0, 'max': 2, 'name': 'user_is_occupied'}}, 'dtype': dtype('int64'), 'is_list': False, 'is_ragged': False}, {'name': 'user_geography', 'tags': {<Tags.USER: 'user'>, <Tags.CATEGORICAL: 'categorical'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/workspace/data/categories/categories/unique.user_geography.parquet', 'embedding_sizes': {'cardinality': 5.0, 'dimension': 16.0}, 'domain': {'min': 0, 'max': 4, 'name': 'user_geography'}}, 'dtype': dtype('int64'), 'is_list': False, 'is_ragged': False}, {'name': 'user_intentions', 'tags': {<Tags.USER: 'user'>, <Tags.CATEGORICAL: 'categorical'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/workspace/data/categories/categories/unique.user_intentions.parquet', 'embedding_sizes': {'cardinality': 663.0, 'dimension': 61.0}, 'domain': {'min': 0, 'max': 662, 'name': 'user_intentions'}}, 'dtype': dtype('int64'), 'is_list': False, 'is_ragged': False}, {'name': 'user_brands', 'tags': {<Tags.USER: 'user'>, <Tags.CATEGORICAL: 'categorical'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/workspace/data/categories/categories/unique.user_brands.parquet', 'embedding_sizes': {'cardinality': 663.0, 'dimension': 61.0}, 'domain': {'min': 0, 'max': 662, 'name': 'user_brands'}}, 'dtype': dtype('int64'), 'is_list': False, 'is_ragged': False}, {'name': 'user_categories', 'tags': {<Tags.USER: 'user'>, <Tags.CATEGORICAL: 'categorical'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/workspace/data/categories/categories/unique.user_categories.parquet', 'embedding_sizes': {'cardinality': 663.0, 'dimension': 61.0}, 'domain': {'min': 0, 'max': 662, 'name': 'user_categories'}}, 'dtype': dtype('int64'), 'is_list': False, 'is_ragged': False}]"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "schema"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3095bd9b",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "We expect the label names to be empty."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "68bf449f",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[]"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "label_names = schema.select_by_tag(Tags.TARGET).column_names\n",
    "label_names"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd210772-2930-433c-83c3-a1e2dad7d0f6",
   "metadata": {},
   "source": [
    "### Negative sampling\n",
    "\n",
    "Many datasets for recommender systems contain implicit feedback with logs of user interactions like clicks, add-to-cart, purchases, music listening events, rather than explicit ratings that reflects user preferences over items. To be able to learn from implicit feedback, we use the general (and naive) assumption that the interacted items are more relevant for the user than the non-interacted ones.\n",
    "In Merlin Models we provide some scalable negative sampling algorithms for the Item Retrieval Task. In particular, in this example, we use the `in-batch` sampling algorithm which uses the items interacted by other users as negatives within the same mini-batch."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "994ddb84-fd90-4547-ae58-c8f678d3dcd7",
   "metadata": {},
   "source": [
    "### Building the Model\n",
    "\n",
    "Now, let's build our Two-Tower model. In a nutshell, we aggregate all user features to feed in user tower and feed the item features to the item tower. Then we compute the positive score by multiplying the user embedding with the item embedding and sample negative items (read more about negative sampling [here](https://openreview.net/pdf?id=824xC-SgWgU) and [here](https://medium.com/mlearning-ai/overview-negative-sampling-on-recommendation-systems-230a051c6cd7)), whose item embeddings are also multiplied by the user embedding. Then we apply the loss function on top of the positive and negative scores."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64745656-a678-48a3-a431-24b45f9119d9",
   "metadata": {},
   "source": [
    "We make sure that the mlp blocks used for the user and query towers have the same last dimension. This is needed because we will compute the dot product between the two towers' outputs to get the similarity scores."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "30446249-70e8-46c2-98eb-40a5e358d86e",
   "metadata": {},
   "outputs": [],
   "source": [
    "tower_dim = 64 \n",
    "\n",
    "# create user schema using USER tag\n",
    "user_schema = schema.select_by_tag(Tags.USER)\n",
    "# create user (query) tower input block\n",
    "user_inputs = mm.InputBlockV2(user_schema)\n",
    "# create user (query) encoder block\n",
    "query = mm.Encoder(user_inputs, mm.MLPBlock([128, tower_dim], no_activation_last_layer=True))\n",
    "\n",
    "# create item schema using ITEM tag\n",
    "item_schema = schema.select_by_tag(Tags.ITEM)\n",
    "# create item (candidate) tower input block\n",
    "item_inputs = mm.InputBlockV2(item_schema)\n",
    "# create item (candidate) encoder block\n",
    "candidate = mm.Encoder(item_inputs, mm.MLPBlock([128, tower_dim], no_activation_last_layer=True))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25407870-5e59-4562-8c39-2869a13ac8bd",
   "metadata": {},
   "source": [
    "`no_activation_last_layer:` when set True, no activation is used for top hidden layer. Learn more [here](https://storage.googleapis.com/pub-tools-public-publication-data/pdf/b9f4e78a8830fe5afcf2f0452862fb3c0d6584ea.pdf)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "261f2f38-65f7-4473-aa42-1237b613cb84",
   "metadata": {},
   "source": [
    "Build the model class."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "64e02cc5-d966-400f-9b40-4ed18cbb6e49",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = mm.TwoTowerModelV2(query, candidate)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00108875",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Note that in the `TwoTowerModelV2` function we did not set `negative_samplers` arg, that means it is set to None. In that case, Two-tower model is trained with contrastive learning and `in-batch` negative sampling strategy.\n",
    "\n",
    "**Metrics:**\n",
    "\n",
    "The following information retrieval metrics are used to compute the Top-10 accuracy of recommendation lists containing all items:\n",
    "\n",
    "- **Normalized Discounted Cumulative Gain (NDCG@10)**: NDCG accounts for rank of the relevant item in the recommendation list and is a more fine-grained metric than HR, which only verifies whether the relevant item is among the top-k items.\n",
    "\n",
    "- **Recall@10**: Also known as HitRate@n when there is only one relevant item in the recommendation list. Recall just verifies whether the relevant item is among the top-n items."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c7ec768",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "We need to initialize the dataloaders."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "1280bb4f",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/2\n",
      "86/86 [==============================] - 12s 48ms/step - loss: 8.2954 - recall_at_10: 0.0246 - ndcg_at_10: 0.0220 - regularization_loss: 0.0000e+00 - loss_batch: 8.2847 - val_loss: 8.2915 - val_recall_at_10: 0.0308 - val_ndcg_at_10: 0.0308 - val_regularization_loss: 0.0000e+00 - val_loss_batch: 7.8375\n",
      "Epoch 2/2\n",
      "86/86 [==============================] - 3s 34ms/step - loss: 8.2950 - recall_at_10: 0.0300 - ndcg_at_10: 0.0298 - regularization_loss: 0.0000e+00 - loss_batch: 8.2842 - val_loss: 8.2922 - val_recall_at_10: 0.0217 - val_ndcg_at_10: 0.0198 - val_regularization_loss: 0.0000e+00 - val_loss_batch: 7.8383\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<keras.callbacks.History at 0x7f8af54c8730>"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.compile(optimizer=\"adam\", run_eagerly=False, metrics=[mm.RecallAt(10), mm.NDCGAt(10)])\n",
    "model.fit(train, validation_data=valid, batch_size=4096, epochs=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a417782-8c3f-4128-a888-f2f9d9974515",
   "metadata": {},
   "source": [
    "The validation metric values are calculated given the positive and negative scores in each batch, and then averaged over batches per epoch. That means validation metrics are not computed using the entire item catalog."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "423c10c2-ae46-45bf-9623-6214c7570ac8",
   "metadata": {},
   "source": [
    "### Evaluate the model accuracy"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "155af447-97c4-4875-97ad-84e678fd7b40",
   "metadata": {},
   "source": [
    "Note that above when we  set `validation_data=valid` in the `model.fit()`, we compute evaluation metrics on validation set using the negative sampling strategy used for training. To determine the exact accuracy of our trained retrieval model, we need to compute the similarity score between a given query and all possible candidates. The higher the score of the positive candidate (the one that is already interacted with, i.e. target item_id returned by dataloader), the more accurate the model is. We can do this using the `topk_model` model that we create below via `to_top_k_encoder` method, and the following section shows how to instantiate it. The `to_top_k_encoder()` is a method of the [RetrievalModelV2](https://github.com/NVIDIA-Merlin/models/blob/stable/merlin/models/tf/models/base.py) class. \n",
    "\n",
    "`unique_rows_by_features` : A utility function allows extracting both unique user and item features tables as Merlin Dataset object that can easily be converted to a cuDF data frame. The function extracts unique rows from a specified dataset (transformed train set) based on a specified id-column tags (`ITEM` and `ITEM_ID`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "d4a15e73-4835-4720-bb64-1d745b3b3f61",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>item_id</th>\n",
       "      <th>item_category</th>\n",
       "      <th>item_shop</th>\n",
       "      <th>item_brand</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>93</th>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>2</td>\n",
       "      <td>2</td>\n",
       "      <td>2</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>3</td>\n",
       "      <td>3</td>\n",
       "      <td>3</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>39</th>\n",
       "      <td>4</td>\n",
       "      <td>4</td>\n",
       "      <td>4</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>22</th>\n",
       "      <td>5</td>\n",
       "      <td>5</td>\n",
       "      <td>5</td>\n",
       "      <td>5</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    item_id  item_category  item_shop  item_brand\n",
       "93        1              1          1           1\n",
       "6         2              2          2           2\n",
       "5         3              3          3           3\n",
       "39        4              4          4           4\n",
       "22        5              5          5           5"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Top-K evaluation\n",
    "candidate_features = unique_rows_by_features(train, Tags.ITEM, Tags.ITEM_ID)\n",
    "candidate_features.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9cb010d9",
   "metadata": {},
   "source": [
    "Below, by using the `topk_model` we can evaluate the trained retrieval model using the entire item catalog. This is applying dot product for entire catalog, and by default it is brute force."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "288bbf45-2afa-4b1a-ae96-a3a4e525de80",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.8/dist-packages/merlin/schema/tags.py:148: UserWarning: Compound tags like Tags.ITEM_ID have been deprecated and will be removed in a future version. Please use the atomic versions of these tags, like [<Tags.ITEM: 'item'>, <Tags.ID: 'id'>].\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "topk = 20\n",
    "topk_model = model.to_top_k_encoder(candidate_features, k=topk, batch_size=128)\n",
    "\n",
    "# we can set `metrics` param in the `compile(), if we want\n",
    "topk_model.compile(run_eagerly=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "eedfe1c5-f706-4713-ad6e-1cf0e92a3e8c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "147/147 [==============================] - 4s 15ms/step - loss: 0.4212 - recall_at_10: 0.0934 - mrr_at_10: 0.0341 - ndcg_at_10: 0.0478 - map_at_10: 0.0341 - precision_at_10: 0.0093 - regularization_loss: 0.0000e+00 - loss_batch: 0.4209\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'loss': 0.4212242662906647,\n",
       " 'recall_at_10': 0.09342904388904572,\n",
       " 'mrr_at_10': 0.03406976908445358,\n",
       " 'ndcg_at_10': 0.047799430787563324,\n",
       " 'map_at_10': 0.03406976908445358,\n",
       " 'precision_at_10': 0.009342900477349758,\n",
       " 'regularization_loss': 0.0,\n",
       " 'loss_batch': 0.3857615888118744}"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "eval_loader = mm.Loader(valid, batch_size=1024).map(mm.ToTarget(schema, \"item_id\"))\n",
    "\n",
    "metrics = topk_model.evaluate(eval_loader, return_dict=True)\n",
    "metrics"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f49037a0-a1c2-4a6f-a89e-e67e0e696445",
   "metadata": {},
   "source": [
    "### Generate top-K recommendations"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "454a1acf-5baf-4f38-8ce4-2c05c6e685cb",
   "metadata": {},
   "source": [
    "We trained a model, now we can generate recommendations offline using `to_top_k_encoder` method. The `to_top_k_encoder()` uses the pre-trained candidate and query encoders to initialize a top-k encoder model, called as `topk_model` in this example. Practically, this method applies the candidate_encoder on the provided candidate_features dataset to set the top-k index of the `topk_model`. Therefore, topk_model object is the one responsible of generating the top-k predictions.\n",
    "\n",
    "Let's generate top-K (k=20 in our example) recommendations for a given batch of 8 samples. The `to_top_k_encoder()` method uses the candidate (item) features dataset as the identifiers, i.e., we extract the`  candidate_id` arg using `Tags.ITEM_ID` tag by default and set it as `index` when calculating the candidate embeddings. The forward method of `topk_model` takes as the query features as input, and computes the dot product scores between the given query embeddings and all the candidates of the top-k index. Then, it returns the top-k (k=20) item ids with the highest scores. Note that instead of calculating the candidate (item) tower embeddings for each user query, we compute the output of the item tower once and store it in the `TopKEncoder` class  to use for the Top-k index. This is computationally more efficient."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "e9e3c032-0084-46c8-b61c-eb7b8595d71b",
   "metadata": {},
   "outputs": [],
   "source": [
    "eval_loader = mm.Loader(valid, batch_size=8, shuffle=False)\n",
    "batch =next(iter(eval_loader))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7198e12-1d29-49ac-bbc3-b12935ba9d08",
   "metadata": {},
   "source": [
    "Let's check the `user_id` column in a given batch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "09f3ada5-eaa2-43cf-810c-97ccfe78f17d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(8, 1), dtype=int64, numpy=\n",
       "array([[ 48],\n",
       "       [  6],\n",
       "       [  2],\n",
       "       [184],\n",
       "       [ 22],\n",
       "       [  6],\n",
       "       [ 30],\n",
       "       [ 36]])>"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch[0]['user_id']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b7b8048-1a90-4311-8d19-be6e8715fc7f",
   "metadata": {},
   "source": [
    "The recommended top 20 item ids are returned below. The output of the method is a named tuple `TopKPrediction`, where the first element is the dot product scores and the second element is the encoded item ids (not the original ids)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "0e720af4-70cd-40c8-857b-75e127e7d867",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "TopKPrediction(scores=<tf.Tensor: shape=(8, 20), dtype=float32, numpy=\n",
       "array([[0.07914167, 0.06689327, 0.06665718, 0.0558873 , 0.05419895,\n",
       "        0.05162647, 0.05141887, 0.04779754, 0.04400383, 0.04338987,\n",
       "        0.04314638, 0.04218695, 0.04206986, 0.03964205, 0.03870573,\n",
       "        0.03738754, 0.03681045, 0.03679563, 0.03595291, 0.03488279],\n",
       "       [0.11142829, 0.1087153 , 0.09285022, 0.09184089, 0.09017296,\n",
       "        0.0819862 , 0.07968194, 0.07311831, 0.07131816, 0.06624649,\n",
       "        0.0660262 , 0.06577689, 0.06359623, 0.06272355, 0.06217033,\n",
       "        0.05999621, 0.05986039, 0.05904868, 0.05829312, 0.05809326],\n",
       "       [0.04108726, 0.03786915, 0.03679999, 0.03606556, 0.0341182 ,\n",
       "        0.03333249, 0.03283013, 0.03179491, 0.03027055, 0.03001958,\n",
       "        0.02996242, 0.02868335, 0.02752452, 0.02751509, 0.02742849,\n",
       "        0.02658201, 0.02591877, 0.02579155, 0.02546304, 0.02458519],\n",
       "       [0.03614997, 0.03377194, 0.03346077, 0.03037825, 0.02977451,\n",
       "        0.02895313, 0.02835435, 0.02678858, 0.02556792, 0.02539463,\n",
       "        0.02518579, 0.0251436 , 0.02497348, 0.02443891, 0.02432298,\n",
       "        0.02418283, 0.02345275, 0.02344878, 0.02332238, 0.0229562 ],\n",
       "       [0.08556695, 0.08040293, 0.08004925, 0.07905072, 0.07071923,\n",
       "        0.06996118, 0.06975519, 0.06796281, 0.06512717, 0.06391992,\n",
       "        0.06375775, 0.06356508, 0.06197023, 0.06085409, 0.05857427,\n",
       "        0.05686603, 0.05492353, 0.05487213, 0.05287076, 0.05180455],\n",
       "       [0.11142829, 0.1087153 , 0.09285022, 0.09184089, 0.09017296,\n",
       "        0.0819862 , 0.07968194, 0.07311831, 0.07131816, 0.06624649,\n",
       "        0.0660262 , 0.06577689, 0.06359623, 0.06272355, 0.06217033,\n",
       "        0.05999621, 0.05986039, 0.05904868, 0.05829312, 0.05809326],\n",
       "       [0.10821203, 0.09383179, 0.06738893, 0.06365363, 0.05353299,\n",
       "        0.04497425, 0.04346779, 0.04234669, 0.04175502, 0.04119734,\n",
       "        0.0407769 , 0.04059385, 0.04049742, 0.04042677, 0.03921955,\n",
       "        0.03844847, 0.03626748, 0.03532993, 0.03404151, 0.03349668],\n",
       "       [0.11652211, 0.11571186, 0.09148473, 0.08903328, 0.08306861,\n",
       "        0.08037847, 0.08032752, 0.0744494 , 0.07389667, 0.07335471,\n",
       "        0.07194579, 0.07136803, 0.06793442, 0.06658586, 0.06494202,\n",
       "        0.06410206, 0.06317658, 0.06307507, 0.05969208, 0.05949583]],\n",
       "      dtype=float32)>, identifiers=<tf.Tensor: shape=(8, 20), dtype=int32, numpy=\n",
       "array([[ 27, 119, 145,  17, 130, 112,  28, 192, 222, 122, 307,  25, 125,\n",
       "        247, 118, 335, 601, 584, 398,  41],\n",
       "       [  2,   3,   5,   9, 220, 301, 305, 308,  36, 277, 559, 545, 230,\n",
       "         13,  10, 615, 518, 132, 185, 468],\n",
       "       [144, 303,  18,  22, 209, 117,  14,  51,  37, 101, 331,  23, 404,\n",
       "         21,  26, 388, 156, 307,  15, 164],\n",
       "       [164, 121,  14, 117, 175, 118,  70, 100, 567, 368,   1,   4, 330,\n",
       "        554,  25, 207,  22, 206,  37, 248],\n",
       "       [ 26,  76,  35, 404,  59,  23, 117, 127, 430,  33,  87, 217, 158,\n",
       "        328,  96, 295, 167, 203, 169, 289],\n",
       "       [  2,   3,   5,   9, 220, 301, 305, 308,  36, 277, 559, 545, 230,\n",
       "         13,  10, 615, 518, 132, 185, 468],\n",
       "       [  4,  16,  68, 195, 324,  30, 334, 231, 128, 459, 114, 238, 120,\n",
       "          2, 660,  86,  89, 192, 255,  83],\n",
       "       [ 27,  21, 307,  11, 207,  76, 178, 138, 139,  48,  26,  78, 246,\n",
       "        146, 479,  67, 590, 188,   7, 169]], dtype=int32)>)"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "topk_model(batch[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b19062ef",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Exporting Retrieval Models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "abf8cb35",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "So far we have trained and evaluated our Retrieval model. Now, the next step is to deploy our model and generate top-K recommendations given a user (query). We can efficiently serve our model by indexing the trained item embeddings into an **Approximate Nearest Neighbors (ANN)** engine. Basically, for a given user query vector, that is generated passing the user features into user tower of retrieval model, we do an ANN search query to find the ids of nearby item vectors, and at serve time, we score user embeddings over all indexed top-K item embeddings within the ANN engine.\n",
    "\n",
    "In doing so, we need to export\n",
    " \n",
    "- user (query) tower\n",
    "- item and user features\n",
    "- item embeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a01df11e",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "#### Save and Load User (query) tower"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4dc5e2c",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "We are able to save the user tower model as a TF model to disk. The user tower model is needed to generate a user embedding vector when a user feature vector <i>x</i> is fed into that model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "07c68eb4",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "query_tower = model.query_encoder\n",
    "query_tower.save(os.path.join(DATA_FOLDER, \"query_tower\"))\n",
    "\n",
    "## we can load back the saved model via the following script.\n",
    "#query_tower_loaded = tf.keras.models.load_model(os.path.join(DATA_FOLDER, 'query_tower'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4c031f8",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "#### Extract and save User features"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9cdd88a9",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "With `unique_rows_by_features` utility function we can easily extract both unique user and item features tables as cuDF dataframes. Note that for user features table, we use `USER` and `USER_ID` tags."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "1f861ac4",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "user_features = (\n",
    "    unique_rows_by_features(train, Tags.USER, Tags.USER_ID).compute().reset_index(drop=True)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "f88fcb3a",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "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>user_shops</th>\n",
       "      <th>user_profile</th>\n",
       "      <th>user_group</th>\n",
       "      <th>user_gender</th>\n",
       "      <th>user_age</th>\n",
       "      <th>user_consumption_2</th>\n",
       "      <th>user_is_occupied</th>\n",
       "      <th>user_geography</th>\n",
       "      <th>user_intentions</th>\n",
       "      <th>user_brands</th>\n",
       "      <th>user_categories</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2</td>\n",
       "      <td>2</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>2</td>\n",
       "      <td>2</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3</td>\n",
       "      <td>3</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>3</td>\n",
       "      <td>3</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>4</td>\n",
       "      <td>4</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>4</td>\n",
       "      <td>4</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>5</td>\n",
       "      <td>5</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>5</td>\n",
       "      <td>5</td>\n",
       "      <td>5</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   user_id  user_shops  user_profile  user_group  user_gender  user_age  \\\n",
       "0        1           1             1           1            1         1   \n",
       "1        2           2             1           1            1         1   \n",
       "2        3           3             1           1            1         1   \n",
       "3        4           4             1           1            1         1   \n",
       "4        5           5             1           1            1         1   \n",
       "\n",
       "   user_consumption_2  user_is_occupied  user_geography  user_intentions  \\\n",
       "0                   1                 1               1                1   \n",
       "1                   1                 1               1                2   \n",
       "2                   1                 1               1                3   \n",
       "3                   1                 1               1                4   \n",
       "4                   1                 1               1                5   \n",
       "\n",
       "   user_brands  user_categories  \n",
       "0            1                1  \n",
       "1            2                2  \n",
       "2            3                3  \n",
       "3            4                4  \n",
       "4            5                5  "
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "user_features.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "33758417",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# save to disk\n",
    "user_features.to_parquet(os.path.join(DATA_FOLDER, \"user_features.parquet\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "184f6cf6-cc5e-459e-9d87-5fdcd01d8c85",
   "metadata": {},
   "source": [
    "#### Generate Query embeddings for entire user catalog"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "5d31a12c-ca41-4b02-9288-7d6b31d54e39",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.8/dist-packages/merlin/schema/tags.py:148: UserWarning: Compound tags like Tags.USER_ID have been deprecated and will be removed in a future version. Please use the atomic versions of these tags, like [<Tags.USER: 'user'>, <Tags.ID: 'id'>].\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "queries = model.query_embeddings(Dataset(user_features, schema=schema.select_by_tag(Tags.USER)), \n",
    "                                 batch_size=1024, index=Tags.USER_ID)\n",
    "query_embs_df = queries.compute(scheduler=\"synchronous\").reset_index()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "42b45f67-b3b3-42c3-8d76-b7ac0e391d7d",
   "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>0</th>\n",
       "      <th>1</th>\n",
       "      <th>2</th>\n",
       "      <th>3</th>\n",
       "      <th>4</th>\n",
       "      <th>5</th>\n",
       "      <th>6</th>\n",
       "      <th>7</th>\n",
       "      <th>8</th>\n",
       "      <th>...</th>\n",
       "      <th>54</th>\n",
       "      <th>55</th>\n",
       "      <th>56</th>\n",
       "      <th>57</th>\n",
       "      <th>58</th>\n",
       "      <th>59</th>\n",
       "      <th>60</th>\n",
       "      <th>61</th>\n",
       "      <th>62</th>\n",
       "      <th>63</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>0.041951</td>\n",
       "      <td>-0.020223</td>\n",
       "      <td>0.004236</td>\n",
       "      <td>-0.019033</td>\n",
       "      <td>0.002884</td>\n",
       "      <td>-0.080953</td>\n",
       "      <td>0.060585</td>\n",
       "      <td>0.052524</td>\n",
       "      <td>-0.029782</td>\n",
       "      <td>...</td>\n",
       "      <td>0.099736</td>\n",
       "      <td>-0.028312</td>\n",
       "      <td>0.038981</td>\n",
       "      <td>-0.098385</td>\n",
       "      <td>-0.101973</td>\n",
       "      <td>-0.051495</td>\n",
       "      <td>0.024279</td>\n",
       "      <td>-0.053948</td>\n",
       "      <td>-0.080442</td>\n",
       "      <td>0.004703</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2</td>\n",
       "      <td>0.053807</td>\n",
       "      <td>-0.033114</td>\n",
       "      <td>-0.038405</td>\n",
       "      <td>-0.050862</td>\n",
       "      <td>0.040319</td>\n",
       "      <td>-0.060739</td>\n",
       "      <td>0.003956</td>\n",
       "      <td>0.012470</td>\n",
       "      <td>-0.063603</td>\n",
       "      <td>...</td>\n",
       "      <td>0.056157</td>\n",
       "      <td>0.034962</td>\n",
       "      <td>-0.003549</td>\n",
       "      <td>-0.070555</td>\n",
       "      <td>-0.023475</td>\n",
       "      <td>-0.058738</td>\n",
       "      <td>0.057749</td>\n",
       "      <td>-0.013966</td>\n",
       "      <td>-0.001620</td>\n",
       "      <td>-0.030841</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3</td>\n",
       "      <td>-0.009804</td>\n",
       "      <td>0.025582</td>\n",
       "      <td>-0.070978</td>\n",
       "      <td>-0.062639</td>\n",
       "      <td>0.034251</td>\n",
       "      <td>-0.105107</td>\n",
       "      <td>0.072854</td>\n",
       "      <td>0.016531</td>\n",
       "      <td>-0.029134</td>\n",
       "      <td>...</td>\n",
       "      <td>0.107710</td>\n",
       "      <td>-0.068276</td>\n",
       "      <td>0.050339</td>\n",
       "      <td>-0.135261</td>\n",
       "      <td>-0.115445</td>\n",
       "      <td>-0.076873</td>\n",
       "      <td>0.058914</td>\n",
       "      <td>-0.050154</td>\n",
       "      <td>-0.071277</td>\n",
       "      <td>-0.031820</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>4</td>\n",
       "      <td>-0.054897</td>\n",
       "      <td>-0.090031</td>\n",
       "      <td>0.010316</td>\n",
       "      <td>0.045851</td>\n",
       "      <td>-0.019615</td>\n",
       "      <td>0.002119</td>\n",
       "      <td>0.034532</td>\n",
       "      <td>0.125433</td>\n",
       "      <td>-0.134117</td>\n",
       "      <td>...</td>\n",
       "      <td>0.073807</td>\n",
       "      <td>-0.021637</td>\n",
       "      <td>0.013071</td>\n",
       "      <td>-0.145496</td>\n",
       "      <td>-0.091502</td>\n",
       "      <td>0.029946</td>\n",
       "      <td>0.058051</td>\n",
       "      <td>-0.018283</td>\n",
       "      <td>0.038066</td>\n",
       "      <td>-0.105121</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>5</td>\n",
       "      <td>0.066658</td>\n",
       "      <td>0.028629</td>\n",
       "      <td>0.027072</td>\n",
       "      <td>-0.024320</td>\n",
       "      <td>0.035782</td>\n",
       "      <td>-0.049752</td>\n",
       "      <td>0.026108</td>\n",
       "      <td>0.030448</td>\n",
       "      <td>-0.002598</td>\n",
       "      <td>...</td>\n",
       "      <td>0.067228</td>\n",
       "      <td>0.029789</td>\n",
       "      <td>0.096082</td>\n",
       "      <td>-0.076030</td>\n",
       "      <td>-0.118867</td>\n",
       "      <td>-0.053601</td>\n",
       "      <td>0.025621</td>\n",
       "      <td>-0.001777</td>\n",
       "      <td>-0.025906</td>\n",
       "      <td>0.021156</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5 rows × 65 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "   user_id         0         1         2         3         4         5  \\\n",
       "0        1  0.041951 -0.020223  0.004236 -0.019033  0.002884 -0.080953   \n",
       "1        2  0.053807 -0.033114 -0.038405 -0.050862  0.040319 -0.060739   \n",
       "2        3 -0.009804  0.025582 -0.070978 -0.062639  0.034251 -0.105107   \n",
       "3        4 -0.054897 -0.090031  0.010316  0.045851 -0.019615  0.002119   \n",
       "4        5  0.066658  0.028629  0.027072 -0.024320  0.035782 -0.049752   \n",
       "\n",
       "          6         7         8  ...        54        55        56        57  \\\n",
       "0  0.060585  0.052524 -0.029782  ...  0.099736 -0.028312  0.038981 -0.098385   \n",
       "1  0.003956  0.012470 -0.063603  ...  0.056157  0.034962 -0.003549 -0.070555   \n",
       "2  0.072854  0.016531 -0.029134  ...  0.107710 -0.068276  0.050339 -0.135261   \n",
       "3  0.034532  0.125433 -0.134117  ...  0.073807 -0.021637  0.013071 -0.145496   \n",
       "4  0.026108  0.030448 -0.002598  ...  0.067228  0.029789  0.096082 -0.076030   \n",
       "\n",
       "         58        59        60        61        62        63  \n",
       "0 -0.101973 -0.051495  0.024279 -0.053948 -0.080442  0.004703  \n",
       "1 -0.023475 -0.058738  0.057749 -0.013966 -0.001620 -0.030841  \n",
       "2 -0.115445 -0.076873  0.058914 -0.050154 -0.071277 -0.031820  \n",
       "3 -0.091502  0.029946  0.058051 -0.018283  0.038066 -0.105121  \n",
       "4 -0.118867 -0.053601  0.025621 -0.001777 -0.025906  0.021156  \n",
       "\n",
       "[5 rows x 65 columns]"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query_embs_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "595b1578",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "#### Extract and save Item features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "3ae320ec",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "item_features = (\n",
    "    unique_rows_by_features(train, Tags.ITEM, Tags.ITEM_ID).compute().reset_index(drop=True)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "8a24dcb8",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>item_id</th>\n",
       "      <th>item_category</th>\n",
       "      <th>item_shop</th>\n",
       "      <th>item_brand</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2</td>\n",
       "      <td>2</td>\n",
       "      <td>2</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3</td>\n",
       "      <td>3</td>\n",
       "      <td>3</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>4</td>\n",
       "      <td>4</td>\n",
       "      <td>4</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>5</td>\n",
       "      <td>5</td>\n",
       "      <td>5</td>\n",
       "      <td>5</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   item_id  item_category  item_shop  item_brand\n",
       "0        1              1          1           1\n",
       "1        2              2          2           2\n",
       "2        3              3          3           3\n",
       "3        4              4          4           4\n",
       "4        5              5          5           5"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "item_features.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "40102380",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# save to disk\n",
    "item_features.to_parquet(os.path.join(DATA_FOLDER, \"item_features.parquet\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f57f5a3e",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "#### Extract and save Item embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "30f50d2a-dbf3-4459-af46-1a8194730896",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.8/dist-packages/merlin/schema/tags.py:148: UserWarning: Compound tags like Tags.ITEM_ID have been deprecated and will be removed in a future version. Please use the atomic versions of these tags, like [<Tags.ITEM: 'item'>, <Tags.ID: 'id'>].\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "item_embs = model.candidate_embeddings(Dataset(item_features, schema=schema.select_by_tag(Tags.ITEM)), \n",
    "                                       batch_size=1024, index=Tags.ITEM_ID)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "138678ae",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "item_embs_df = item_embs.compute(scheduler=\"synchronous\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "a69694b0",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "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>0</th>\n",
       "      <th>1</th>\n",
       "      <th>2</th>\n",
       "      <th>3</th>\n",
       "      <th>4</th>\n",
       "      <th>5</th>\n",
       "      <th>6</th>\n",
       "      <th>7</th>\n",
       "      <th>8</th>\n",
       "      <th>9</th>\n",
       "      <th>...</th>\n",
       "      <th>54</th>\n",
       "      <th>55</th>\n",
       "      <th>56</th>\n",
       "      <th>57</th>\n",
       "      <th>58</th>\n",
       "      <th>59</th>\n",
       "      <th>60</th>\n",
       "      <th>61</th>\n",
       "      <th>62</th>\n",
       "      <th>63</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>item_id</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>-0.021552</td>\n",
       "      <td>0.009041</td>\n",
       "      <td>0.039659</td>\n",
       "      <td>-0.115629</td>\n",
       "      <td>0.060703</td>\n",
       "      <td>-0.025404</td>\n",
       "      <td>0.001678</td>\n",
       "      <td>-0.022630</td>\n",
       "      <td>-0.043372</td>\n",
       "      <td>-0.043020</td>\n",
       "      <td>...</td>\n",
       "      <td>0.055401</td>\n",
       "      <td>-0.020197</td>\n",
       "      <td>0.091445</td>\n",
       "      <td>-0.031282</td>\n",
       "      <td>-0.038851</td>\n",
       "      <td>0.003069</td>\n",
       "      <td>0.018751</td>\n",
       "      <td>0.007528</td>\n",
       "      <td>-0.035626</td>\n",
       "      <td>-0.057195</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>-0.083753</td>\n",
       "      <td>-0.116476</td>\n",
       "      <td>0.073937</td>\n",
       "      <td>-0.021895</td>\n",
       "      <td>-0.037875</td>\n",
       "      <td>0.079780</td>\n",
       "      <td>0.004227</td>\n",
       "      <td>-0.062871</td>\n",
       "      <td>0.067589</td>\n",
       "      <td>-0.113532</td>\n",
       "      <td>...</td>\n",
       "      <td>0.016558</td>\n",
       "      <td>-0.062858</td>\n",
       "      <td>0.007280</td>\n",
       "      <td>-0.145148</td>\n",
       "      <td>-0.073590</td>\n",
       "      <td>0.060061</td>\n",
       "      <td>0.042919</td>\n",
       "      <td>-0.088881</td>\n",
       "      <td>0.045371</td>\n",
       "      <td>-0.048104</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>-0.022572</td>\n",
       "      <td>-0.143781</td>\n",
       "      <td>0.095565</td>\n",
       "      <td>-0.072755</td>\n",
       "      <td>-0.024208</td>\n",
       "      <td>0.016993</td>\n",
       "      <td>-0.063992</td>\n",
       "      <td>-0.016109</td>\n",
       "      <td>-0.046727</td>\n",
       "      <td>-0.076944</td>\n",
       "      <td>...</td>\n",
       "      <td>-0.016716</td>\n",
       "      <td>-0.120303</td>\n",
       "      <td>-0.027779</td>\n",
       "      <td>-0.030434</td>\n",
       "      <td>0.003762</td>\n",
       "      <td>0.044557</td>\n",
       "      <td>0.028406</td>\n",
       "      <td>-0.097061</td>\n",
       "      <td>0.047091</td>\n",
       "      <td>-0.070876</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>-0.034698</td>\n",
       "      <td>0.046747</td>\n",
       "      <td>0.050764</td>\n",
       "      <td>-0.054780</td>\n",
       "      <td>0.049690</td>\n",
       "      <td>-0.054132</td>\n",
       "      <td>0.044862</td>\n",
       "      <td>-0.051085</td>\n",
       "      <td>0.025827</td>\n",
       "      <td>-0.008969</td>\n",
       "      <td>...</td>\n",
       "      <td>0.034949</td>\n",
       "      <td>-0.009016</td>\n",
       "      <td>-0.010063</td>\n",
       "      <td>-0.007224</td>\n",
       "      <td>0.004805</td>\n",
       "      <td>-0.033879</td>\n",
       "      <td>0.026008</td>\n",
       "      <td>0.066739</td>\n",
       "      <td>-0.081608</td>\n",
       "      <td>0.002007</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>0.017667</td>\n",
       "      <td>-0.093071</td>\n",
       "      <td>0.119211</td>\n",
       "      <td>-0.039053</td>\n",
       "      <td>-0.044223</td>\n",
       "      <td>-0.002900</td>\n",
       "      <td>-0.011234</td>\n",
       "      <td>0.021149</td>\n",
       "      <td>0.006894</td>\n",
       "      <td>-0.073159</td>\n",
       "      <td>...</td>\n",
       "      <td>0.041408</td>\n",
       "      <td>-0.068992</td>\n",
       "      <td>0.081714</td>\n",
       "      <td>0.003833</td>\n",
       "      <td>-0.015269</td>\n",
       "      <td>-0.011885</td>\n",
       "      <td>0.041777</td>\n",
       "      <td>-0.039943</td>\n",
       "      <td>0.100604</td>\n",
       "      <td>-0.012898</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>657</th>\n",
       "      <td>-0.008736</td>\n",
       "      <td>-0.025814</td>\n",
       "      <td>0.021996</td>\n",
       "      <td>-0.067662</td>\n",
       "      <td>0.038304</td>\n",
       "      <td>-0.039822</td>\n",
       "      <td>-0.014580</td>\n",
       "      <td>-0.014729</td>\n",
       "      <td>-0.051024</td>\n",
       "      <td>-0.017028</td>\n",
       "      <td>...</td>\n",
       "      <td>0.036276</td>\n",
       "      <td>-0.017140</td>\n",
       "      <td>0.026932</td>\n",
       "      <td>-0.039948</td>\n",
       "      <td>-0.042317</td>\n",
       "      <td>0.036567</td>\n",
       "      <td>0.014907</td>\n",
       "      <td>-0.069974</td>\n",
       "      <td>0.014789</td>\n",
       "      <td>-0.000786</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>658</th>\n",
       "      <td>0.027235</td>\n",
       "      <td>-0.055484</td>\n",
       "      <td>0.100215</td>\n",
       "      <td>-0.022288</td>\n",
       "      <td>-0.015791</td>\n",
       "      <td>-0.013113</td>\n",
       "      <td>-0.011211</td>\n",
       "      <td>-0.060337</td>\n",
       "      <td>-0.005251</td>\n",
       "      <td>-0.066255</td>\n",
       "      <td>...</td>\n",
       "      <td>-0.007403</td>\n",
       "      <td>0.035502</td>\n",
       "      <td>0.051252</td>\n",
       "      <td>0.036262</td>\n",
       "      <td>-0.033077</td>\n",
       "      <td>0.054807</td>\n",
       "      <td>-0.015973</td>\n",
       "      <td>-0.037896</td>\n",
       "      <td>0.036095</td>\n",
       "      <td>-0.012522</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>659</th>\n",
       "      <td>0.000050</td>\n",
       "      <td>-0.083014</td>\n",
       "      <td>-0.003769</td>\n",
       "      <td>-0.016448</td>\n",
       "      <td>0.058482</td>\n",
       "      <td>0.003810</td>\n",
       "      <td>0.002106</td>\n",
       "      <td>-0.021589</td>\n",
       "      <td>-0.029082</td>\n",
       "      <td>-0.005173</td>\n",
       "      <td>...</td>\n",
       "      <td>-0.014442</td>\n",
       "      <td>-0.045111</td>\n",
       "      <td>0.025007</td>\n",
       "      <td>0.011599</td>\n",
       "      <td>-0.026002</td>\n",
       "      <td>0.031850</td>\n",
       "      <td>0.043968</td>\n",
       "      <td>-0.053928</td>\n",
       "      <td>0.040387</td>\n",
       "      <td>0.037862</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>660</th>\n",
       "      <td>0.000492</td>\n",
       "      <td>-0.041255</td>\n",
       "      <td>0.081524</td>\n",
       "      <td>-0.083455</td>\n",
       "      <td>-0.020575</td>\n",
       "      <td>-0.030294</td>\n",
       "      <td>-0.026690</td>\n",
       "      <td>-0.076162</td>\n",
       "      <td>0.006011</td>\n",
       "      <td>-0.035908</td>\n",
       "      <td>...</td>\n",
       "      <td>-0.015531</td>\n",
       "      <td>0.006367</td>\n",
       "      <td>0.016362</td>\n",
       "      <td>0.038916</td>\n",
       "      <td>-0.027108</td>\n",
       "      <td>-0.009466</td>\n",
       "      <td>0.007149</td>\n",
       "      <td>0.016705</td>\n",
       "      <td>0.010482</td>\n",
       "      <td>0.012888</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>661</th>\n",
       "      <td>-0.022624</td>\n",
       "      <td>-0.034587</td>\n",
       "      <td>0.018065</td>\n",
       "      <td>-0.007772</td>\n",
       "      <td>-0.001248</td>\n",
       "      <td>0.064035</td>\n",
       "      <td>0.059197</td>\n",
       "      <td>0.004046</td>\n",
       "      <td>-0.047599</td>\n",
       "      <td>-0.050420</td>\n",
       "      <td>...</td>\n",
       "      <td>0.067813</td>\n",
       "      <td>0.057818</td>\n",
       "      <td>0.074110</td>\n",
       "      <td>-0.052934</td>\n",
       "      <td>-0.015619</td>\n",
       "      <td>0.083385</td>\n",
       "      <td>-0.017303</td>\n",
       "      <td>0.056356</td>\n",
       "      <td>0.062265</td>\n",
       "      <td>-0.011084</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>661 rows × 64 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                0         1         2         3         4         5         6  \\\n",
       "item_id                                                                         \n",
       "1       -0.021552  0.009041  0.039659 -0.115629  0.060703 -0.025404  0.001678   \n",
       "2       -0.083753 -0.116476  0.073937 -0.021895 -0.037875  0.079780  0.004227   \n",
       "3       -0.022572 -0.143781  0.095565 -0.072755 -0.024208  0.016993 -0.063992   \n",
       "4       -0.034698  0.046747  0.050764 -0.054780  0.049690 -0.054132  0.044862   \n",
       "5        0.017667 -0.093071  0.119211 -0.039053 -0.044223 -0.002900 -0.011234   \n",
       "...           ...       ...       ...       ...       ...       ...       ...   \n",
       "657     -0.008736 -0.025814  0.021996 -0.067662  0.038304 -0.039822 -0.014580   \n",
       "658      0.027235 -0.055484  0.100215 -0.022288 -0.015791 -0.013113 -0.011211   \n",
       "659      0.000050 -0.083014 -0.003769 -0.016448  0.058482  0.003810  0.002106   \n",
       "660      0.000492 -0.041255  0.081524 -0.083455 -0.020575 -0.030294 -0.026690   \n",
       "661     -0.022624 -0.034587  0.018065 -0.007772 -0.001248  0.064035  0.059197   \n",
       "\n",
       "                7         8         9  ...        54        55        56  \\\n",
       "item_id                                ...                                 \n",
       "1       -0.022630 -0.043372 -0.043020  ...  0.055401 -0.020197  0.091445   \n",
       "2       -0.062871  0.067589 -0.113532  ...  0.016558 -0.062858  0.007280   \n",
       "3       -0.016109 -0.046727 -0.076944  ... -0.016716 -0.120303 -0.027779   \n",
       "4       -0.051085  0.025827 -0.008969  ...  0.034949 -0.009016 -0.010063   \n",
       "5        0.021149  0.006894 -0.073159  ...  0.041408 -0.068992  0.081714   \n",
       "...           ...       ...       ...  ...       ...       ...       ...   \n",
       "657     -0.014729 -0.051024 -0.017028  ...  0.036276 -0.017140  0.026932   \n",
       "658     -0.060337 -0.005251 -0.066255  ... -0.007403  0.035502  0.051252   \n",
       "659     -0.021589 -0.029082 -0.005173  ... -0.014442 -0.045111  0.025007   \n",
       "660     -0.076162  0.006011 -0.035908  ... -0.015531  0.006367  0.016362   \n",
       "661      0.004046 -0.047599 -0.050420  ...  0.067813  0.057818  0.074110   \n",
       "\n",
       "               57        58        59        60        61        62        63  \n",
       "item_id                                                                        \n",
       "1       -0.031282 -0.038851  0.003069  0.018751  0.007528 -0.035626 -0.057195  \n",
       "2       -0.145148 -0.073590  0.060061  0.042919 -0.088881  0.045371 -0.048104  \n",
       "3       -0.030434  0.003762  0.044557  0.028406 -0.097061  0.047091 -0.070876  \n",
       "4       -0.007224  0.004805 -0.033879  0.026008  0.066739 -0.081608  0.002007  \n",
       "5        0.003833 -0.015269 -0.011885  0.041777 -0.039943  0.100604 -0.012898  \n",
       "...           ...       ...       ...       ...       ...       ...       ...  \n",
       "657     -0.039948 -0.042317  0.036567  0.014907 -0.069974  0.014789 -0.000786  \n",
       "658      0.036262 -0.033077  0.054807 -0.015973 -0.037896  0.036095 -0.012522  \n",
       "659      0.011599 -0.026002  0.031850  0.043968 -0.053928  0.040387  0.037862  \n",
       "660      0.038916 -0.027108 -0.009466  0.007149  0.016705  0.010482  0.012888  \n",
       "661     -0.052934 -0.015619  0.083385 -0.017303  0.056356  0.062265 -0.011084  \n",
       "\n",
       "[661 rows x 64 columns]"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "item_embs_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "5d152e82",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# save to disk\n",
    "item_embs_df.to_parquet(os.path.join(DATA_FOLDER, \"item_embeddings.parquet\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1eed43c3",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "That's it. You have learned how to train and evaluate your Two-Tower retrieval model, and then how to export the required components to be able to deploy this model to generate recommendations. In order to learn more on serving a model to [Triton Inference Server](https://github.com/triton-inference-server/server), please explore the examples in the [Merlin](https://github.com/NVIDIA-Merlin/Merlin) and [Merlin Systems](https://github.com/NVIDIA-Merlin/systems) repos."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.10"
  },
  "merlin": {
   "containers": [
    "nvcr.io/nvidia/merlin/merlin-tensorflow:latest"
   ]
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  },
  "vscode": {
   "interpreter": {
    "hash": "a398807c5c2ed8e5ff9d9890488d007fa99cbabcec733962e21659a28c5da99b"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}