offline-batch-recsys/Offline-Batch-Recommender-System.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "ca23d058-eefd-45e5-8769-1f277acc032b",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Offline Batch Recommender System\n",
    "\n",
    "In this notebook, we will build a system to generate recommendations \"offline\". This means that at some interval (like a cron job) the system writes recommendations to a database (Redis) for later retrieval. The architecture diagram below shows how it comes together from a bird's eye view.\n",
    "\n",
    "![img](https://raw.githubusercontent.com/RedisVentures/Redis-Recsys/master/assets/OfflineBatchRecsys.png)\n",
    "\n",
    "## 1.0 - Architecture\n",
    "\n",
    "A multi-stage pipeline for a recommender system is a common approach to efficiently retrieve relevant items from a large catalog. The pipeline includes several stages:\n",
    "\n",
    "1) A fast **Candidate Retrieval Model** quickly truncates the large item catalog to a relevant set of hundreds (or thousands) of items.\n",
    "2) Filtering is performed to remove undesirable or already-seen items.\n",
    "3) A finely-tuned deep learning **Ranking Model** (i.e. more powerful) ranks the most likely items that are going to interacted with.\n",
    "4) Results are ordered and returned to the user.\n",
    "\n",
    "Writing the recommendations to a key-value store like Redis allows developers to access the recommendations in near real-time at a later point without the complexity of hosting a live multi-stage recommendation system. This can be especially useful for developers who can't afford the complexity of hosting a live multi-stage recsys or need to get something up and running quickly. Additionally, by saving the recommendations in a key-value store rather than serving the models directly, this system can also save on the cost of hosting the model. The next notebook in this series will dive into when \"online\" recommendation systems are valuable using the same models created below.\n",
    "\n",
    "In this notebook, we will:\n",
    "\n",
    "1) Prepare the [**Dataset**](#2.0---Dataset-Preparation)\n",
    "2) Build a [**Candidate Retrieval Model**](#3.0---Candidate-Retrieval-Model)\n",
    "3) Build a [**Ranking Model**](#4.0---Ranking-Model)\n",
    "4) [**Write Recommendations**](#5.0---Write-Recommendations-to-Redis) to Redis (offline)\n",
    "5) [**Fetch Recommendations**](#6.0---Fetch-Recommendations-from-Redis) from Redis\n",
    "6) [**Export Models**](#7.0---Export-Models) for later use"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bbf37d1b-6aa8-4df8-9ac2-db968b595078",
   "metadata": {},
   "source": [
    "*This notebook was created using the latest stable [merlin-tensorflow](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/merlin/containers/merlin-tensorflow/tags) container and was heavily based on the work done by the NVIDIA Merlin team [here](https://github.com/NVIDIA-Merlin/models/blob/main/examples/05-Retrieval-Model.ipynb)*"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "401daa6d",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## 2.0 - Dataset Preparation\n",
    "\n",
    "We will use a synthetic dataset that mimicks the [Ali-CCP: Alibaba Click and Conversion Prediction](https://tianchi.aliyun.com/dataset/dataDetail?dataId=408#1) dataset. This allows us to tune it to our exact needs for demonstration/learning purposes.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cdac1fea",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### 2.1 - Importing Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "92aa8daa",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2023-01-21 02:31:29.856298: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:991] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2023-01-21 02:31:29.856734: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:991] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2023-01-21 02:31:29.856847: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:991] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2023-01-21 02:31:30.142625: 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",
      "2023-01-21 02:31:30.143343: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:991] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2023-01-21 02:31:30.143513: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:991] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2023-01-21 02:31:30.143624: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:991] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2023-01-21 02:31:31.187117: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:991] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2023-01-21 02:31:31.187302: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:991] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2023-01-21 02:31:31.187420: I tensorflow/stream_executor/cuda/cuda_gpu_executor.cc:991] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero\n",
      "2023-01-21 02:31:31.187521: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1532] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 5632 MB memory:  -> device: 0, name: NVIDIA GeForce RTX 2080 Ti, pci bus id: 0000:01:00.0, compute capability: 7.5\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import logging\n",
    "import time\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "\n",
    "import nvtabular as nvt\n",
    "import merlin.models.tf as mm\n",
    "import tensorflow as tf\n",
    "\n",
    "from nvtabular.ops import *\n",
    "\n",
    "from merlin.datasets.synthetic import generate_data\n",
    "from merlin.datasets.ecommerce import transform_aliccp\n",
    "from merlin.models.utils.example_utils import workflow_fit_transform\n",
    "from merlin.models.utils.dataset import unique_rows_by_features\n",
    "from merlin.schema.tags import Tags\n",
    "from merlin.io.dataset import Dataset\n",
    "\n",
    "\n",
    "# disable INFO and DEBUG logging everywhere\n",
    "logging.disable(logging.WARNING)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ac0eef4",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### 2.2 Generate and Process Synthetic Ali-CCP Dataset\n",
    "\n",
    "The Merlin ecosystem built by NVIDIA, provides a number of pre-built datasets that come with convience functions for generating and processing them for use in recommendation systems.\n",
    "\n",
    "With `transform_aliccp` function, we transform the raw dataset... applying the operators defined in the NVTabular workflow pipeline above. The processed parquet files are saved to the output path.\n",
    "\n",
    "This pipeline is used to process the data and prepare it for use in a recommendation system by converting columns to appropriate types, categorifying features, tagging features and targets, and removing null values."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1dd757a3-2f93-495f-b2c5-5098c997d46e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Generate the synthetic data\n",
    "NUM_ROWS = 1000000\n",
    "TRAIN_SIZE = 0.7\n",
    "VALID_SIZE = 0.3\n",
    "\n",
    "train, valid = generate_data(\"aliccp-raw\", NUM_ROWS, set_sizes=(TRAIN_SIZE, VALID_SIZE))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f23232bf-7822-498f-a036-23e44ef48e50",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define output path for data\n",
    "DATA_DIR = \"/model-data/aliccp\"\n",
    "OUTPUT_DATA_DIR = os.path.join(DATA_DIR, \"processed\")\n",
    "OUTPUT_RETRIEVAL_DATA_DIR = os.path.join(OUTPUT_DATA_DIR, \"retrieval\")\n",
    "CATEGORY_TEMP_DIR = os.path.join(DATA_DIR, \"categories\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "3d1d91d2-3b6a-457b-8b18-2589ddc5fc99",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define NVTabular Feature Transformation Pipeline\n",
    "\n",
    "user_id_raw = [\"user_id\"] >> Rename(postfix='_raw') >> LambdaOp(lambda col: col.astype(\"int32\")) >> TagAsUserFeatures()\n",
    "item_id_raw = [\"item_id\"] >> Rename(postfix='_raw') >> LambdaOp(lambda col: col.astype(\"int32\")) >> TagAsItemFeatures()\n",
    "\n",
    "user_id = [\"user_id\"] >> Categorify(dtype=\"int32\", out_path=CATEGORY_TEMP_DIR) >> TagAsUserID()\n",
    "item_id = [\"item_id\"] >> Categorify(dtype=\"int32\", out_path=CATEGORY_TEMP_DIR) >> TagAsItemID()\n",
    "\n",
    "item_features = (\n",
    "    [\"item_category\", \"item_shop\", \"item_brand\"] >> Categorify(dtype=\"int32\", out_path=CATEGORY_TEMP_DIR) >> TagAsItemFeatures()\n",
    ")\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(dtype=\"int32\", out_path=CATEGORY_TEMP_DIR)\n",
    "    >> TagAsUserFeatures()\n",
    ")\n",
    "\n",
    "targets = [\"click\"] >> AddMetadata(tags=[Tags.BINARY_CLASSIFICATION, \"target\"])\n",
    "\n",
    "outputs = user_id + item_id + item_features + user_features +  user_id_raw + item_id_raw + targets\n",
    "\n",
    "# add dropna op to filter rows with nulls\n",
    "outputs = outputs >> Dropna()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "07ab4edc",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Transform data and create files\n",
    "transform_aliccp((train, valid), OUTPUT_DATA_DIR, nvt_workflow=outputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b97bdb6e-c22e-4679-bf5c-a5213eca444f",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## 3.0 - Candidate Retrieval Model\n",
    "\n",
    "We will use a **Two-Tower** model to infer a subset of relevant items from large item corpus for a given user. \n",
    "\n",
    "A Two-Tower recommendation system is a type of recommendation system that uses two neural network architectures, or \"towers,\" to generate recommendations. One tower, called the \"user tower,\" is used to model the user's preferences, while the other tower, called the \"item tower,\" is used to model the characteristics of the items being recommended.\n",
    "\n",
    "The two towers are typically trained separately, but their outputs are combined to generate recommendations. The user tower produces a user representation vector, while the item tower produces an item representation vector. These vectors are then used to compute the similarity between the user and the items, which is used to rank the items and generate recommendations.\n",
    "\n",
    "\n",
    "<img src=\"https://d3i71xaburhd42.cloudfront.net/8c32706b6af49db5d9cd9217e5196f701e473537/2-Figure1-1.png\"  width=\"30%\">\n",
    "\n",
    "Image from: [Mixed Negative Sampling for Learning Two-tower Neural Networks in Recommendations](https://www.semanticscholar.org/paper/Mixed-Negative-Sampling-for-Learning-Two-tower-in-Yang-Yi/29f080a1bb6df6f45afd82c443f72da745983bee)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1bbbb1c4-ae30-44bf-888c-d0d9500a23ea",
   "metadata": {},
   "source": [
    "### 3.1 - Feature Engineering for Candidate Retrieval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a7e0c257",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Load Datasets from generated files\n",
    "train_retrieval = Dataset(os.path.join(OUTPUT_DATA_DIR, \"train\", \"*.parquet\"))\n",
    "valid_retrieval = Dataset(os.path.join(OUTPUT_DATA_DIR, \"valid\", \"*.parquet\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "f328b83d-c39d-4dd5-a8b4-c09cb8b5820f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define NVTabular Feature Transformation Pipeline\n",
    "\n",
    "inputs = train_retrieval.schema.column_names\n",
    "\n",
    "# Select only positive interaction rows where click==1 in the dataset with Filter() operator\n",
    "outputs = inputs >> Filter(f=lambda df: df[\"click\"] == 1)\n",
    "\n",
    "# Execute the transformation workflow for both datasets\n",
    "workflow = nvt.Workflow(outputs)\n",
    "workflow.fit(train_retrieval)\n",
    "workflow.transform(train_retrieval).to_parquet(\n",
    "    output_path=os.path.join(OUTPUT_RETRIEVAL_DATA_DIR, \"train\")\n",
    ")\n",
    "workflow.transform(valid_retrieval).to_parquet(\n",
    "    output_path=os.path.join(OUTPUT_RETRIEVAL_DATA_DIR, \"valid\")\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "affeb1d9-a8d9-4bb4-b3b3-dec99b0825e3",
   "metadata": {},
   "source": [
    "*NVTabular exported the schema file, schema.pbtxt a protobuf text file, of our processed dataset. To learn more about the schema object and schema file you can explore [this notebook](https://github.com/NVIDIA-Merlin/models/blob/main/examples/02-Merlin-Models-and-NVTabular-integration.ipynb).*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "19f7ee51-ddc1-4bfc-b330-fe5c992b60a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read transformed parquet files as Dataset objects\n",
    "train_retrieval = Dataset(os.path.join(OUTPUT_RETRIEVAL_DATA_DIR, \"train\", \"*.parquet\"), part_size=\"500MB\")\n",
    "valid_retrieval = Dataset(os.path.join(OUTPUT_RETRIEVAL_DATA_DIR, \"valid\", \"*.parquet\"), part_size=\"500MB\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fcdb2dbc-9bbd-46ad-aede-44c3ffbaac9d",
   "metadata": {},
   "source": [
    "Now we can use the `schema` object to define the model inputs. We select features with user and item tags, and exclude raw IDs and target column."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "4a0454c8-8eec-47cc-82e8-b45637b76d35",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create model input schema\n",
    "schema = train_retrieval.schema.select_by_tag([Tags.ITEM_ID, Tags.USER_ID, Tags.ITEM, Tags.USER]).without(['user_id_raw', 'item_id_raw', 'click'])\n",
    "train_retrieval.schema = schema\n",
    "valid_retrieval.schema = schema"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "7d9cffc1-21f8-4574-82bd-89ebcd47713f",
   "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.CATEGORICAL, Tags.USER, Tags.USER_ID, Ta...</td>\n",
       "      <td>int32</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>/model-data/aliccp/categories/categories/uniqu...</td>\n",
       "      <td>753.0</td>\n",
       "      <td>65.0</td>\n",
       "      <td>0</td>\n",
       "      <td>752</td>\n",
       "      <td>user_id</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>item_id</td>\n",
       "      <td>(Tags.ITEM, Tags.CATEGORICAL, Tags.ID, Tags.IT...</td>\n",
       "      <td>int32</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>/model-data/aliccp/categories/categories/uniqu...</td>\n",
       "      <td>754.0</td>\n",
       "      <td>65.0</td>\n",
       "      <td>0</td>\n",
       "      <td>753</td>\n",
       "      <td>item_id</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>item_category</td>\n",
       "      <td>(Tags.ITEM, Tags.CATEGORICAL)</td>\n",
       "      <td>int32</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>/model-data/aliccp/categories/categories/uniqu...</td>\n",
       "      <td>754.0</td>\n",
       "      <td>65.0</td>\n",
       "      <td>0</td>\n",
       "      <td>753</td>\n",
       "      <td>item_category</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>item_shop</td>\n",
       "      <td>(Tags.ITEM, Tags.CATEGORICAL)</td>\n",
       "      <td>int32</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>/model-data/aliccp/categories/categories/uniqu...</td>\n",
       "      <td>754.0</td>\n",
       "      <td>65.0</td>\n",
       "      <td>0</td>\n",
       "      <td>753</td>\n",
       "      <td>item_shop</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>item_brand</td>\n",
       "      <td>(Tags.ITEM, Tags.CATEGORICAL)</td>\n",
       "      <td>int32</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>/model-data/aliccp/categories/categories/uniqu...</td>\n",
       "      <td>754.0</td>\n",
       "      <td>65.0</td>\n",
       "      <td>0</td>\n",
       "      <td>753</td>\n",
       "      <td>item_brand</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>user_shops</td>\n",
       "      <td>(Tags.CATEGORICAL, Tags.USER)</td>\n",
       "      <td>int32</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>/model-data/aliccp/categories/categories/uniqu...</td>\n",
       "      <td>753.0</td>\n",
       "      <td>65.0</td>\n",
       "      <td>0</td>\n",
       "      <td>752</td>\n",
       "      <td>user_shops</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>user_profile</td>\n",
       "      <td>(Tags.CATEGORICAL, Tags.USER)</td>\n",
       "      <td>int32</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>/model-data/aliccp/categories/categories/uniqu...</td>\n",
       "      <td>51.0</td>\n",
       "      <td>16.0</td>\n",
       "      <td>0</td>\n",
       "      <td>50</td>\n",
       "      <td>user_profile</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>user_group</td>\n",
       "      <td>(Tags.CATEGORICAL, Tags.USER)</td>\n",
       "      <td>int32</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>/model-data/aliccp/categories/categories/uniqu...</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.CATEGORICAL, Tags.USER)</td>\n",
       "      <td>int32</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>/model-data/aliccp/categories/categories/uniqu...</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.CATEGORICAL, Tags.USER)</td>\n",
       "      <td>int32</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>/model-data/aliccp/categories/categories/uniqu...</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.CATEGORICAL, Tags.USER)</td>\n",
       "      <td>int32</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>/model-data/aliccp/categories/categories/uniqu...</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.CATEGORICAL, Tags.USER)</td>\n",
       "      <td>int32</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>/model-data/aliccp/categories/categories/uniqu...</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.CATEGORICAL, Tags.USER)</td>\n",
       "      <td>int32</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>/model-data/aliccp/categories/categories/uniqu...</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.CATEGORICAL, Tags.USER)</td>\n",
       "      <td>int32</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>/model-data/aliccp/categories/categories/uniqu...</td>\n",
       "      <td>753.0</td>\n",
       "      <td>65.0</td>\n",
       "      <td>0</td>\n",
       "      <td>752</td>\n",
       "      <td>user_intentions</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14</th>\n",
       "      <td>user_brands</td>\n",
       "      <td>(Tags.CATEGORICAL, Tags.USER)</td>\n",
       "      <td>int32</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>/model-data/aliccp/categories/categories/uniqu...</td>\n",
       "      <td>753.0</td>\n",
       "      <td>65.0</td>\n",
       "      <td>0</td>\n",
       "      <td>752</td>\n",
       "      <td>user_brands</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15</th>\n",
       "      <td>user_categories</td>\n",
       "      <td>(Tags.CATEGORICAL, Tags.USER)</td>\n",
       "      <td>int32</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>/model-data/aliccp/categories/categories/uniqu...</td>\n",
       "      <td>753.0</td>\n",
       "      <td>65.0</td>\n",
       "      <td>0</td>\n",
       "      <td>752</td>\n",
       "      <td>user_categories</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "[{'name': 'user_id', 'tags': {<Tags.CATEGORICAL: 'categorical'>, <Tags.USER: 'user'>, <Tags.USER_ID: 'user_id'>, <Tags.ID: 'id'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/model-data/aliccp/categories/categories/unique.user_id.parquet', 'embedding_sizes': {'cardinality': 753.0, 'dimension': 65.0}, 'domain': {'min': 0, 'max': 752, 'name': 'user_id'}}, 'dtype': dtype('int32'), 'is_list': False, 'is_ragged': False}, {'name': 'item_id', 'tags': {<Tags.ITEM: 'item'>, <Tags.CATEGORICAL: 'categorical'>, <Tags.ID: 'id'>, <Tags.ITEM_ID: 'item_id'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/model-data/aliccp/categories/categories/unique.item_id.parquet', 'embedding_sizes': {'cardinality': 754.0, 'dimension': 65.0}, 'domain': {'min': 0, 'max': 753, 'name': 'item_id'}}, 'dtype': dtype('int32'), 'is_list': False, 'is_ragged': False}, {'name': 'item_category', 'tags': {<Tags.ITEM: 'item'>, <Tags.CATEGORICAL: 'categorical'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/model-data/aliccp/categories/categories/unique.item_category.parquet', 'embedding_sizes': {'cardinality': 754.0, 'dimension': 65.0}, 'domain': {'min': 0, 'max': 753, 'name': 'item_category'}}, 'dtype': dtype('int32'), 'is_list': False, 'is_ragged': False}, {'name': 'item_shop', 'tags': {<Tags.ITEM: 'item'>, <Tags.CATEGORICAL: 'categorical'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/model-data/aliccp/categories/categories/unique.item_shop.parquet', 'embedding_sizes': {'cardinality': 754.0, 'dimension': 65.0}, 'domain': {'min': 0, 'max': 753, 'name': 'item_shop'}}, 'dtype': dtype('int32'), 'is_list': False, 'is_ragged': False}, {'name': 'item_brand', 'tags': {<Tags.ITEM: 'item'>, <Tags.CATEGORICAL: 'categorical'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/model-data/aliccp/categories/categories/unique.item_brand.parquet', 'embedding_sizes': {'cardinality': 754.0, 'dimension': 65.0}, 'domain': {'min': 0, 'max': 753, 'name': 'item_brand'}}, 'dtype': dtype('int32'), 'is_list': False, 'is_ragged': False}, {'name': 'user_shops', 'tags': {<Tags.CATEGORICAL: 'categorical'>, <Tags.USER: 'user'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/model-data/aliccp/categories/categories/unique.user_shops.parquet', 'embedding_sizes': {'cardinality': 753.0, 'dimension': 65.0}, 'domain': {'min': 0, 'max': 752, 'name': 'user_shops'}}, 'dtype': dtype('int32'), 'is_list': False, 'is_ragged': False}, {'name': 'user_profile', 'tags': {<Tags.CATEGORICAL: 'categorical'>, <Tags.USER: 'user'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/model-data/aliccp/categories/categories/unique.user_profile.parquet', 'embedding_sizes': {'cardinality': 51.0, 'dimension': 16.0}, 'domain': {'min': 0, 'max': 50, 'name': 'user_profile'}}, 'dtype': dtype('int32'), 'is_list': False, 'is_ragged': False}, {'name': 'user_group', 'tags': {<Tags.CATEGORICAL: 'categorical'>, <Tags.USER: 'user'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/model-data/aliccp/categories/categories/unique.user_group.parquet', 'embedding_sizes': {'cardinality': 10.0, 'dimension': 16.0}, 'domain': {'min': 0, 'max': 9, 'name': 'user_group'}}, 'dtype': dtype('int32'), 'is_list': False, 'is_ragged': False}, {'name': 'user_gender', 'tags': {<Tags.CATEGORICAL: 'categorical'>, <Tags.USER: 'user'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/model-data/aliccp/categories/categories/unique.user_gender.parquet', 'embedding_sizes': {'cardinality': 3.0, 'dimension': 16.0}, 'domain': {'min': 0, 'max': 2, 'name': 'user_gender'}}, 'dtype': dtype('int32'), 'is_list': False, 'is_ragged': False}, {'name': 'user_age', 'tags': {<Tags.CATEGORICAL: 'categorical'>, <Tags.USER: 'user'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/model-data/aliccp/categories/categories/unique.user_age.parquet', 'embedding_sizes': {'cardinality': 7.0, 'dimension': 16.0}, 'domain': {'min': 0, 'max': 6, 'name': 'user_age'}}, 'dtype': dtype('int32'), 'is_list': False, 'is_ragged': False}, {'name': 'user_consumption_2', 'tags': {<Tags.CATEGORICAL: 'categorical'>, <Tags.USER: 'user'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/model-data/aliccp/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('int32'), 'is_list': False, 'is_ragged': False}, {'name': 'user_is_occupied', 'tags': {<Tags.CATEGORICAL: 'categorical'>, <Tags.USER: 'user'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/model-data/aliccp/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('int32'), 'is_list': False, 'is_ragged': False}, {'name': 'user_geography', 'tags': {<Tags.CATEGORICAL: 'categorical'>, <Tags.USER: 'user'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/model-data/aliccp/categories/categories/unique.user_geography.parquet', 'embedding_sizes': {'cardinality': 5.0, 'dimension': 16.0}, 'domain': {'min': 0, 'max': 4, 'name': 'user_geography'}}, 'dtype': dtype('int32'), 'is_list': False, 'is_ragged': False}, {'name': 'user_intentions', 'tags': {<Tags.CATEGORICAL: 'categorical'>, <Tags.USER: 'user'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/model-data/aliccp/categories/categories/unique.user_intentions.parquet', 'embedding_sizes': {'cardinality': 753.0, 'dimension': 65.0}, 'domain': {'min': 0, 'max': 752, 'name': 'user_intentions'}}, 'dtype': dtype('int32'), 'is_list': False, 'is_ragged': False}, {'name': 'user_brands', 'tags': {<Tags.CATEGORICAL: 'categorical'>, <Tags.USER: 'user'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/model-data/aliccp/categories/categories/unique.user_brands.parquet', 'embedding_sizes': {'cardinality': 753.0, 'dimension': 65.0}, 'domain': {'min': 0, 'max': 752, 'name': 'user_brands'}}, 'dtype': dtype('int32'), 'is_list': False, 'is_ragged': False}, {'name': 'user_categories', 'tags': {<Tags.CATEGORICAL: 'categorical'>, <Tags.USER: 'user'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/model-data/aliccp/categories/categories/unique.user_categories.parquet', 'embedding_sizes': {'cardinality': 753.0, 'dimension': 65.0}, 'domain': {'min': 0, 'max': 752, 'name': 'user_categories'}}, 'dtype': dtype('int32'), 'is_list': False, 'is_ragged': False}]"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Inspect the schema!\n",
    "schema"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "994ddb84-fd90-4547-ae58-c8f678d3dcd7",
   "metadata": {},
   "source": [
    "### 3.2 - Designing Retrieval Model Architecture\n",
    "\n",
    "The **Two-Tower** model consists of a **User tower** (where all user features are fed) and an **Item tower** (where all item features are fed).\n",
    "\n",
    "The User tower generates an embedding for the User. Then it computes the positive interaction \"score\" (likelihood of interaction event) using the dot-product between the User embedding and the Item embedding, in addition to sampled \"negative\" Items within a batch.\n",
    "\n",
    "##### About 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. \n",
    "\n",
    "In Merlin Models -- NVIDIA provides some scalable negative sampling algorithms for this Item Retrieval task. 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": "code",
   "execution_count": 11,
   "id": "30446249-70e8-46c2-98eb-40a5e358d86e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to create the two tower retrieval model\n",
    "\n",
    "def create_two_tower(tower_dim: int, encoder_dim: int, optimizer: str, k: int, tags) -> mm.TwoTowerModelV2:\n",
    "    # User/Query Tower\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(\n",
    "        user_inputs,\n",
    "        mm.MLPBlock([encoder_dim, tower_dim], no_activation_last_layer=True)\n",
    "    )\n",
    "\n",
    "    # Item/Candidate Tower\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(\n",
    "        item_inputs,\n",
    "        mm.MLPBlock([encoder_dim, tower_dim], no_activation_last_layer=True)\n",
    "    )\n",
    "    \n",
    "    # Build Model Class\n",
    "    model = mm.TwoTowerModelV2(query, candidate)\n",
    "    model.compile(\n",
    "        optimizer=optimizer,\n",
    "        run_eagerly=False,\n",
    "        loss=\"categorical_crossentropy\",\n",
    "        metrics=[mm.RecallAt(k), mm.NDCGAt(k)]\n",
    "    )\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25407870-5e59-4562-8c39-2869a13ac8bd",
   "metadata": {},
   "source": [
    "**Notes:**\n",
    "- `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).\n",
    "- In the `TwoTowerModelV2` function we did not set `negative_samplers` arg. By default, it uses contrastive learning and `in-batch` negative sampling strategy.\n",
    "- Two metrics are used to judge the quality of the recommendations: **Normalized Discounted Cumulative Gain (NDCG@K)** and **Recall@K**.\n",
    "    - NDCG@K 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",
    "    - Recall (Also known as HitRate@K) when there is only one relevant item in the recommendation list. Recall just verifies whether the relevant item is among the top-k items.\n",
    "- 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. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "1280bb4f",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/2\n",
      "86/86 [==============================] - 9s 38ms/step - loss: 8.2952 - recall_at_10: 0.0232 - ndcg_at_10: 0.0208 - regularization_loss: 0.0000e+00 - loss_batch: 8.2788 - val_loss: 8.2906 - val_recall_at_10: 0.0297 - val_ndcg_at_10: 0.0297 - val_regularization_loss: 0.0000e+00 - val_loss_batch: 7.7285\n",
      "Epoch 2/2\n",
      "86/86 [==============================] - 2s 28ms/step - loss: 8.2947 - recall_at_10: 0.0303 - ndcg_at_10: 0.0300 - regularization_loss: 0.0000e+00 - loss_batch: 8.2783 - val_loss: 8.2911 - val_recall_at_10: 0.0221 - val_ndcg_at_10: 0.0203 - val_regularization_loss: 0.0000e+00 - val_loss_batch: 7.7291\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<keras.callbacks.History at 0x7fdb549578e0>"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Initialize model\n",
    "retrieval_model = create_two_tower(\n",
    "    tower_dim=64,\n",
    "    encoder_dim=128,\n",
    "    optimizer=\"adam\",\n",
    "    k=10,\n",
    "    tags=Tags\n",
    ")\n",
    "\n",
    "# Fit model\n",
    "retrieval_model.fit(train_retrieval, validation_data=valid_retrieval, batch_size=4096, epochs=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "423c10c2-ae46-45bf-9623-6214c7570ac8",
   "metadata": {},
   "source": [
    "### 3.3 - Evaluate the model accuracy"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "155af447-97c4-4875-97ad-84e678fd7b40",
   "metadata": {},
   "source": [
    "The validation metric values during training 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.**\n",
    "\n",
    "To determine the exact accuracy, we need to compute the similarity score between a given query and all possible candidates. Below, by using the `topk_model` we can evaluate the trained retrieval model using the entire item catalog (brute force)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "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>3</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>38</th>\n",
       "      <td>3</td>\n",
       "      <td>3</td>\n",
       "      <td>3</td>\n",
       "      <td>3</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>56</th>\n",
       "      <td>4</td>\n",
       "      <td>4</td>\n",
       "      <td>4</td>\n",
       "      <td>4</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>50</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",
       "3         1              1          1           1\n",
       "6         2              2          2           2\n",
       "38        3              3          3           3\n",
       "56        4              4          4           4\n",
       "50        5              5          5           5"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create candidate/item features for evaluation\n",
    "candidate_features = unique_rows_by_features(train_retrieval, Tags.ITEM, Tags.ITEM_ID)\n",
    "\n",
    "# Here's a sneek peek of the item features data\n",
    "candidate_features.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "288bbf45-2afa-4b1a-ae96-a3a4e525de80",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert model to a top_k_encoder\n",
    "topk_model = retrieval_model.to_top_k_encoder(candidate_features, k=20, batch_size=128)\n",
    "topk_model.compile(run_eagerly=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "2b42eaa0-996e-43e3-a951-fc013e2d0d3d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "147/147 [==============================] - 3s 9ms/step - loss: 0.4388 - recall_at_10: 0.0929 - mrr_at_10: 0.0351 - ndcg_at_10: 0.0485 - map_at_10: 0.0351 - precision_at_10: 0.0093 - regularization_loss: 0.0000e+00 - loss_batch: 0.4407\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'loss': 0.43882274627685547,\n",
       " 'recall_at_10': 0.09288118034601212,\n",
       " 'mrr_at_10': 0.03513556718826294,\n",
       " 'ndcg_at_10': 0.04852905124425888,\n",
       " 'map_at_10': 0.03513556718826294,\n",
       " 'precision_at_10': 0.00928812101483345,\n",
       " 'regularization_loss': 0.0,\n",
       " 'loss_batch': 0.6025397777557373}"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create data loader for validation data\n",
    "eval_loader = mm.Loader(valid_retrieval, batch_size=1024).map(mm.ToTarget(schema, \"item_id\"))\n",
    "\n",
    "# Evaluation\n",
    "metrics = topk_model.evaluate(eval_loader, return_dict=True)\n",
    "metrics"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f49037a0-a1c2-4a6f-a89e-e67e0e696445",
   "metadata": {},
   "source": [
    "### 3.4 - Generate top-K candidates"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "454a1acf-5baf-4f38-8ce4-2c05c6e685cb",
   "metadata": {},
   "source": [
    "Let's generate top-K (k=20 in our example) recommendation candidates for a given batch of 8 samples.\n",
    "- The `to_top_k_encoder()` method uses the item/candidate features dataset to compute and store all item/candidate embeddings in an index.\n",
    "- The forward method of `topk_model` takes as the query/user features as input, and computes the dot product scores between the given query/user embeddings and all the candidates of the top-k index.\n",
    "- Then, it returns the top-k (k=20) item ids with the highest scores."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "96d15826-5e18-4061-8834-c36406455e32",
   "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>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>3372</th>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>39</td>\n",
       "      <td>5</td>\n",
       "      <td>1</td>\n",
       "      <td>3</td>\n",
       "      <td>2</td>\n",
       "      <td>1</td>\n",
       "      <td>2</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>18</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>22</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>11</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",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "      user_id  user_shops  user_profile  user_group  user_gender  user_age  \\\n",
       "3372        0           0            39           5            1         3   \n",
       "18          1           1             1           1            1         1   \n",
       "22          2           2             1           1            1         1   \n",
       "2           3           3             1           1            1         1   \n",
       "11          4           4             1           1            1         1   \n",
       "\n",
       "      user_consumption_2  user_is_occupied  user_geography  user_intentions  \\\n",
       "3372                   2                 1               2                0   \n",
       "18                     1                 1               1                1   \n",
       "22                     1                 1               1                2   \n",
       "2                      1                 1               1                3   \n",
       "11                     1                 1               1                4   \n",
       "\n",
       "      user_brands  user_categories  \n",
       "3372            0                0  \n",
       "18              1                1  \n",
       "22              2                2  \n",
       "2               3                3  \n",
       "11              4                4  "
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create query/user features for evaluation\n",
    "user_features = unique_rows_by_features(valid_retrieval, Tags.USER, Tags.USER_ID)\n",
    "user_features.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "e9e3c032-0084-46c8-b61c-eb7b8595d71b",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[[0]\n",
      " [1]\n",
      " [2]\n",
      " [3]\n",
      " [4]\n",
      " [5]\n",
      " [6]\n",
      " [7]], shape=(8, 1), dtype=int32)\n"
     ]
    }
   ],
   "source": [
    "# Check out one batch of 8 Users\n",
    "loader = mm.Loader(user_features, batch_size=8, shuffle=False)\n",
    "batch = next(iter(loader))\n",
    "print(batch[0]['user_id'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b7b8048-1a90-4311-8d19-be6e8715fc7f",
   "metadata": {},
   "source": [
    "The recommended top 20 item ids and scores are returned below from the candidate retrieval model for each of the 8 selected users (from the validation set)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "0e720af4-70cd-40c8-857b-75e127e7d867",
   "metadata": {},
   "outputs": [],
   "source": [
    "scores, reccommended_item_ids = topk_model(batch[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "ed2b07aa-ada9-4f37-a5d2-a4dea5940403",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(8, 20), dtype=int32, numpy=\n",
       "array([[  2,  53, 219, 339,  17, 203,  88,   8,  10, 591, 102,   9, 193,\n",
       "        263, 296, 313, 366, 170, 208, 104],\n",
       "       [  2, 339,  53,  15,   7,   8, 104, 219, 313,  44, 391,  16, 581,\n",
       "        367, 122,  36, 131,  12,  49,  88],\n",
       "       [184, 267, 189,  14,  37, 242,   5,  13,   7,  43, 182, 209,  24,\n",
       "        323, 469, 344,  61, 315, 202, 346],\n",
       "       [  5,  39,  12, 353,  71, 384,  98,  47, 182, 496, 276, 226,  37,\n",
       "         67, 552, 196,   3, 423, 359,  86],\n",
       "       [  2,   4, 102,  53,   1, 239,  75,   3, 207, 291, 349,  31, 416,\n",
       "        303,  26, 103, 283, 131, 751, 328],\n",
       "       [ 11,   5, 421,  16,  37, 496, 226,  39, 448, 274,   7, 201, 251,\n",
       "        360, 279,  14, 625,  12, 150, 210],\n",
       "       [182,  71,  20, 384, 184,  61,   3,   5,  14, 315,  25,  39, 271,\n",
       "        234, 746,  33, 115, 253, 189, 285],\n",
       "       [ 53, 391,  49,  42, 339,  13,  44,  36,  77, 124,  15, 131,   2,\n",
       "         91,  64, 300, 335, 388, 177, 367]], dtype=int32)>"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Recommended Items for the batch of 8 Users\n",
    "reccommended_item_ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "59a74673-73ba-4706-be99-c3037f2bd6ba",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(8, 20), dtype=float32, numpy=\n",
       "array([[0.06132706, 0.05495376, 0.05188668, 0.04804639, 0.04525334,\n",
       "        0.04467893, 0.04359638, 0.0429577 , 0.04236389, 0.04190419,\n",
       "        0.04150067, 0.03876395, 0.03788529, 0.03769156, 0.03755951,\n",
       "        0.03725768, 0.03700415, 0.03667157, 0.03543973, 0.03480223],\n",
       "       [0.06057974, 0.05873544, 0.05457463, 0.04985515, 0.0474914 ,\n",
       "        0.04664886, 0.04660673, 0.04630806, 0.04628112, 0.04596485,\n",
       "        0.04518457, 0.04455723, 0.04452522, 0.04431298, 0.04381118,\n",
       "        0.04294571, 0.04294293, 0.0429119 , 0.04242484, 0.04143431],\n",
       "       [0.08359075, 0.08266199, 0.07058696, 0.06940745, 0.06668985,\n",
       "        0.0660269 , 0.06584951, 0.06411734, 0.06376088, 0.06375296,\n",
       "        0.06344488, 0.06105455, 0.06002223, 0.05888034, 0.05773782,\n",
       "        0.05707346, 0.05687609, 0.05232539, 0.05226946, 0.05199999],\n",
       "       [0.11616069, 0.08919618, 0.08878286, 0.08471525, 0.08420472,\n",
       "        0.08336139, 0.08041547, 0.07971578, 0.07602289, 0.07405829,\n",
       "        0.07403331, 0.0731288 , 0.06993872, 0.06915572, 0.06598826,\n",
       "        0.06594489, 0.06438265, 0.06390512, 0.06255455, 0.06220961],\n",
       "       [0.05636432, 0.05528677, 0.03927067, 0.0382652 , 0.03549547,\n",
       "        0.03413099, 0.03387703, 0.03256701, 0.03203993, 0.03157425,\n",
       "        0.03142566, 0.03089456, 0.03064269, 0.0303258 , 0.03032339,\n",
       "        0.02989603, 0.02922622, 0.02809661, 0.0278715 , 0.02711965],\n",
       "       [0.07619128, 0.06807565, 0.06013621, 0.05950812, 0.05655721,\n",
       "        0.05644257, 0.05487436, 0.05457787, 0.05408981, 0.0524896 ,\n",
       "        0.05196728, 0.05184248, 0.05125651, 0.05065448, 0.05015856,\n",
       "        0.05003097, 0.04912207, 0.04907685, 0.04902352, 0.04806277],\n",
       "       [0.07970415, 0.07465717, 0.07225692, 0.06710847, 0.06676053,\n",
       "        0.06339377, 0.06201241, 0.06161601, 0.0611356 , 0.06069867,\n",
       "        0.06024875, 0.05982666, 0.05912331, 0.05726478, 0.05624647,\n",
       "        0.05578331, 0.05473503, 0.05261823, 0.05231078, 0.05225336],\n",
       "       [0.0881146 , 0.08730791, 0.07916678, 0.07495458, 0.07264654,\n",
       "        0.07154912, 0.07076031, 0.06975465, 0.06837998, 0.06535089,\n",
       "        0.06464162, 0.06234738, 0.0617562 , 0.06088701, 0.06032664,\n",
       "        0.06032572, 0.05976337, 0.0570838 , 0.0565482 , 0.05555537]],\n",
       "      dtype=float32)>"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Recommendation \"scores\" (higher is better) for each recommended item for the batch of 8 Users\n",
    "scores"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4fd2484-38df-4b28-9b5b-961e3e46b29f",
   "metadata": {},
   "source": [
    "At this point - we now have a trained **Candidate Retrieval Model**. Input a **User** and find the topK most likely **Items** to be interacted with. These will serve as inputs to the next model in the pipeline..."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ac898a1-ebbf-4771-8f4d-fc8b6f47f587",
   "metadata": {},
   "source": [
    "## 4.0 - Ranking Model\n",
    "\n",
    "Ranking models are a type of machine learning models that are commonly used in recommendation systems to rank items based on their relevance or likelihood of being interacted with by a user. These models can be used to generate personalized recommendations for individual users by taking into account their preferences and past interactions such as clicks, purchases, or ratings. \n",
    "\n",
    "There are several types of ranking models that can be used in recommendation systems, such as:\n",
    "\n",
    "1. Collaborative Filtering models, which use the past interactions of users to generate recommendations\n",
    "2. Content-Based models, which generate recommendations based on the characteristics of the items\n",
    "3. Hybrid models, which combine the above two approaches (we will use this one)\n",
    "\n",
    "We will use a Deep Learning Recommendation Model [(DLRM)](https://arxiv.org/abs/1906.00091) architecture (published by Facebook/Meta in 2019) to score and rank User/Item pairs. The model was introduced as a personalization deep learning model that uses embeddings to process sparse features that represent categorical data and a multilayer perceptron (MLP) to process dense features, then interacts these features explicitly using the statistical techniques proposed in [here](https://ieeexplore.ieee.org/stamp/stamp.jsp?tp=&arnumber=5694074). To learn more about DLRM architetcture please visit [this notebook](https://github.com/NVIDIA-Merlin/models/blob/main/examples/04-Exporting-ranking-models.ipynb) in the Merlin Models GH repo.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d6559bcc-de8c-48e8-b709-2f6ed2bb25d0",
   "metadata": {},
   "source": [
    "### 4.1 - Feature Engineering for DLRM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "6b1667ea-73dc-4dee-bafc-bf93af122996",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define train and valid dataset objects\n",
    "train_rank = Dataset(os.path.join(OUTPUT_DATA_DIR, \"train\", \"*.parquet\"), part_size=\"500MB\")\n",
    "valid_rank = Dataset(os.path.join(OUTPUT_DATA_DIR, \"valid\", \"*.parquet\"), part_size=\"500MB\")\n",
    "\n",
    "# Define schema object\n",
    "schema = train_rank.schema.without(['user_id_raw', 'item_id_raw'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "782fc973-47e4-40ef-89af-57d3bcd28f67",
   "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.CATEGORICAL, Tags.USER, Tags.USER_ID, Ta...</td>\n",
       "      <td>int32</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>/model-data/aliccp/categories/categories/uniqu...</td>\n",
       "      <td>753.0</td>\n",
       "      <td>65.0</td>\n",
       "      <td>0</td>\n",
       "      <td>752</td>\n",
       "      <td>user_id</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>item_id</td>\n",
       "      <td>(Tags.ITEM, Tags.CATEGORICAL, Tags.ID, Tags.IT...</td>\n",
       "      <td>int32</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>/model-data/aliccp/categories/categories/uniqu...</td>\n",
       "      <td>754.0</td>\n",
       "      <td>65.0</td>\n",
       "      <td>0</td>\n",
       "      <td>753</td>\n",
       "      <td>item_id</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>item_category</td>\n",
       "      <td>(Tags.ITEM, Tags.CATEGORICAL)</td>\n",
       "      <td>int32</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>/model-data/aliccp/categories/categories/uniqu...</td>\n",
       "      <td>754.0</td>\n",
       "      <td>65.0</td>\n",
       "      <td>0</td>\n",
       "      <td>753</td>\n",
       "      <td>item_category</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>item_shop</td>\n",
       "      <td>(Tags.ITEM, Tags.CATEGORICAL)</td>\n",
       "      <td>int32</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>/model-data/aliccp/categories/categories/uniqu...</td>\n",
       "      <td>754.0</td>\n",
       "      <td>65.0</td>\n",
       "      <td>0</td>\n",
       "      <td>753</td>\n",
       "      <td>item_shop</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>item_brand</td>\n",
       "      <td>(Tags.ITEM, Tags.CATEGORICAL)</td>\n",
       "      <td>int32</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>/model-data/aliccp/categories/categories/uniqu...</td>\n",
       "      <td>754.0</td>\n",
       "      <td>65.0</td>\n",
       "      <td>0</td>\n",
       "      <td>753</td>\n",
       "      <td>item_brand</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>user_shops</td>\n",
       "      <td>(Tags.CATEGORICAL, Tags.USER)</td>\n",
       "      <td>int32</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>/model-data/aliccp/categories/categories/uniqu...</td>\n",
       "      <td>753.0</td>\n",
       "      <td>65.0</td>\n",
       "      <td>0</td>\n",
       "      <td>752</td>\n",
       "      <td>user_shops</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>user_profile</td>\n",
       "      <td>(Tags.CATEGORICAL, Tags.USER)</td>\n",
       "      <td>int32</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>/model-data/aliccp/categories/categories/uniqu...</td>\n",
       "      <td>51.0</td>\n",
       "      <td>16.0</td>\n",
       "      <td>0</td>\n",
       "      <td>50</td>\n",
       "      <td>user_profile</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>user_group</td>\n",
       "      <td>(Tags.CATEGORICAL, Tags.USER)</td>\n",
       "      <td>int32</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>/model-data/aliccp/categories/categories/uniqu...</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.CATEGORICAL, Tags.USER)</td>\n",
       "      <td>int32</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>/model-data/aliccp/categories/categories/uniqu...</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.CATEGORICAL, Tags.USER)</td>\n",
       "      <td>int32</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>/model-data/aliccp/categories/categories/uniqu...</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.CATEGORICAL, Tags.USER)</td>\n",
       "      <td>int32</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>/model-data/aliccp/categories/categories/uniqu...</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.CATEGORICAL, Tags.USER)</td>\n",
       "      <td>int32</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>/model-data/aliccp/categories/categories/uniqu...</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.CATEGORICAL, Tags.USER)</td>\n",
       "      <td>int32</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>/model-data/aliccp/categories/categories/uniqu...</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.CATEGORICAL, Tags.USER)</td>\n",
       "      <td>int32</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>/model-data/aliccp/categories/categories/uniqu...</td>\n",
       "      <td>753.0</td>\n",
       "      <td>65.0</td>\n",
       "      <td>0</td>\n",
       "      <td>752</td>\n",
       "      <td>user_intentions</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14</th>\n",
       "      <td>user_brands</td>\n",
       "      <td>(Tags.CATEGORICAL, Tags.USER)</td>\n",
       "      <td>int32</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>/model-data/aliccp/categories/categories/uniqu...</td>\n",
       "      <td>753.0</td>\n",
       "      <td>65.0</td>\n",
       "      <td>0</td>\n",
       "      <td>752</td>\n",
       "      <td>user_brands</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15</th>\n",
       "      <td>user_categories</td>\n",
       "      <td>(Tags.CATEGORICAL, Tags.USER)</td>\n",
       "      <td>int32</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>/model-data/aliccp/categories/categories/uniqu...</td>\n",
       "      <td>753.0</td>\n",
       "      <td>65.0</td>\n",
       "      <td>0</td>\n",
       "      <td>752</td>\n",
       "      <td>user_categories</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16</th>\n",
       "      <td>click</td>\n",
       "      <td>(Tags.TARGET, Tags.BINARY_CLASSIFICATION)</td>\n",
       "      <td>int64</td>\n",
       "      <td>False</td>\n",
       "      <td>False</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>NaN</td>\n",
       "      <td>0</td>\n",
       "      <td>2</td>\n",
       "      <td>click</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "[{'name': 'user_id', 'tags': {<Tags.CATEGORICAL: 'categorical'>, <Tags.USER: 'user'>, <Tags.USER_ID: 'user_id'>, <Tags.ID: 'id'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/model-data/aliccp/categories/categories/unique.user_id.parquet', 'embedding_sizes': {'cardinality': 753.0, 'dimension': 65.0}, 'domain': {'min': 0, 'max': 752, 'name': 'user_id'}}, 'dtype': dtype('int32'), 'is_list': False, 'is_ragged': False}, {'name': 'item_id', 'tags': {<Tags.ITEM: 'item'>, <Tags.CATEGORICAL: 'categorical'>, <Tags.ID: 'id'>, <Tags.ITEM_ID: 'item_id'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/model-data/aliccp/categories/categories/unique.item_id.parquet', 'embedding_sizes': {'cardinality': 754.0, 'dimension': 65.0}, 'domain': {'min': 0, 'max': 753, 'name': 'item_id'}}, 'dtype': dtype('int32'), 'is_list': False, 'is_ragged': False}, {'name': 'item_category', 'tags': {<Tags.ITEM: 'item'>, <Tags.CATEGORICAL: 'categorical'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/model-data/aliccp/categories/categories/unique.item_category.parquet', 'embedding_sizes': {'cardinality': 754.0, 'dimension': 65.0}, 'domain': {'min': 0, 'max': 753, 'name': 'item_category'}}, 'dtype': dtype('int32'), 'is_list': False, 'is_ragged': False}, {'name': 'item_shop', 'tags': {<Tags.ITEM: 'item'>, <Tags.CATEGORICAL: 'categorical'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/model-data/aliccp/categories/categories/unique.item_shop.parquet', 'embedding_sizes': {'cardinality': 754.0, 'dimension': 65.0}, 'domain': {'min': 0, 'max': 753, 'name': 'item_shop'}}, 'dtype': dtype('int32'), 'is_list': False, 'is_ragged': False}, {'name': 'item_brand', 'tags': {<Tags.ITEM: 'item'>, <Tags.CATEGORICAL: 'categorical'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/model-data/aliccp/categories/categories/unique.item_brand.parquet', 'embedding_sizes': {'cardinality': 754.0, 'dimension': 65.0}, 'domain': {'min': 0, 'max': 753, 'name': 'item_brand'}}, 'dtype': dtype('int32'), 'is_list': False, 'is_ragged': False}, {'name': 'user_shops', 'tags': {<Tags.CATEGORICAL: 'categorical'>, <Tags.USER: 'user'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/model-data/aliccp/categories/categories/unique.user_shops.parquet', 'embedding_sizes': {'cardinality': 753.0, 'dimension': 65.0}, 'domain': {'min': 0, 'max': 752, 'name': 'user_shops'}}, 'dtype': dtype('int32'), 'is_list': False, 'is_ragged': False}, {'name': 'user_profile', 'tags': {<Tags.CATEGORICAL: 'categorical'>, <Tags.USER: 'user'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/model-data/aliccp/categories/categories/unique.user_profile.parquet', 'embedding_sizes': {'cardinality': 51.0, 'dimension': 16.0}, 'domain': {'min': 0, 'max': 50, 'name': 'user_profile'}}, 'dtype': dtype('int32'), 'is_list': False, 'is_ragged': False}, {'name': 'user_group', 'tags': {<Tags.CATEGORICAL: 'categorical'>, <Tags.USER: 'user'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/model-data/aliccp/categories/categories/unique.user_group.parquet', 'embedding_sizes': {'cardinality': 10.0, 'dimension': 16.0}, 'domain': {'min': 0, 'max': 9, 'name': 'user_group'}}, 'dtype': dtype('int32'), 'is_list': False, 'is_ragged': False}, {'name': 'user_gender', 'tags': {<Tags.CATEGORICAL: 'categorical'>, <Tags.USER: 'user'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/model-data/aliccp/categories/categories/unique.user_gender.parquet', 'embedding_sizes': {'cardinality': 3.0, 'dimension': 16.0}, 'domain': {'min': 0, 'max': 2, 'name': 'user_gender'}}, 'dtype': dtype('int32'), 'is_list': False, 'is_ragged': False}, {'name': 'user_age', 'tags': {<Tags.CATEGORICAL: 'categorical'>, <Tags.USER: 'user'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/model-data/aliccp/categories/categories/unique.user_age.parquet', 'embedding_sizes': {'cardinality': 7.0, 'dimension': 16.0}, 'domain': {'min': 0, 'max': 6, 'name': 'user_age'}}, 'dtype': dtype('int32'), 'is_list': False, 'is_ragged': False}, {'name': 'user_consumption_2', 'tags': {<Tags.CATEGORICAL: 'categorical'>, <Tags.USER: 'user'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/model-data/aliccp/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('int32'), 'is_list': False, 'is_ragged': False}, {'name': 'user_is_occupied', 'tags': {<Tags.CATEGORICAL: 'categorical'>, <Tags.USER: 'user'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/model-data/aliccp/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('int32'), 'is_list': False, 'is_ragged': False}, {'name': 'user_geography', 'tags': {<Tags.CATEGORICAL: 'categorical'>, <Tags.USER: 'user'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/model-data/aliccp/categories/categories/unique.user_geography.parquet', 'embedding_sizes': {'cardinality': 5.0, 'dimension': 16.0}, 'domain': {'min': 0, 'max': 4, 'name': 'user_geography'}}, 'dtype': dtype('int32'), 'is_list': False, 'is_ragged': False}, {'name': 'user_intentions', 'tags': {<Tags.CATEGORICAL: 'categorical'>, <Tags.USER: 'user'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/model-data/aliccp/categories/categories/unique.user_intentions.parquet', 'embedding_sizes': {'cardinality': 753.0, 'dimension': 65.0}, 'domain': {'min': 0, 'max': 752, 'name': 'user_intentions'}}, 'dtype': dtype('int32'), 'is_list': False, 'is_ragged': False}, {'name': 'user_brands', 'tags': {<Tags.CATEGORICAL: 'categorical'>, <Tags.USER: 'user'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/model-data/aliccp/categories/categories/unique.user_brands.parquet', 'embedding_sizes': {'cardinality': 753.0, 'dimension': 65.0}, 'domain': {'min': 0, 'max': 752, 'name': 'user_brands'}}, 'dtype': dtype('int32'), 'is_list': False, 'is_ragged': False}, {'name': 'user_categories', 'tags': {<Tags.CATEGORICAL: 'categorical'>, <Tags.USER: 'user'>}, 'properties': {'num_buckets': None, 'freq_threshold': 0.0, 'max_size': 0.0, 'start_index': 0.0, 'cat_path': '/model-data/aliccp/categories/categories/unique.user_categories.parquet', 'embedding_sizes': {'cardinality': 753.0, 'dimension': 65.0}, 'domain': {'min': 0, 'max': 752, 'name': 'user_categories'}}, 'dtype': dtype('int32'), 'is_list': False, 'is_ragged': False}, {'name': 'click', 'tags': {<Tags.TARGET: 'target'>, <Tags.BINARY_CLASSIFICATION: 'binary_classification'>}, 'properties': {'domain': {'min': 0, 'max': 2, 'name': 'click'}}, 'dtype': dtype('int64'), 'is_list': False, 'is_ragged': False}]"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Inspect schema - DLRM takes all of these features as input!\n",
    "schema"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "6fbf24bf-faf9-44c8-95f7-36554172346c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'click'"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Target column here is \"click\" (binary classification)\n",
    "target_column = schema.select_by_tag(Tags.TARGET).column_names[0]\n",
    "target_column"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a4b9bd9-3f7e-4bbf-baa7-f39453b2dbbb",
   "metadata": {},
   "source": [
    "### 4.2 - Train DLRM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "fc3fe4bc-54e4-4b98-b5f2-39e14c328854",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_dlrm(optimizer: str, schema, target_column: str, embedding_dim: int):\n",
    "    model = mm.DLRMModel(\n",
    "        schema,\n",
    "        embedding_dim=embedding_dim,\n",
    "        bottom_block=mm.MLPBlock([128, 64]),\n",
    "        top_block=mm.MLPBlock([128, 64, 32]),\n",
    "        prediction_tasks=mm.BinaryClassificationTask(target_column),\n",
    "    )\n",
    "    model.compile(optimizer=optimizer, run_eagerly=False, metrics=[tf.keras.metrics.AUC()])\n",
    "    return model\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "37756986-80a7-41ff-bb61-d275139dced7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/3\n",
      "43/43 [==============================] - 4s 29ms/step - loss: 0.6932 - auc: 0.5004 - regularization_loss: 0.0000e+00 - loss_batch: 0.6932 - val_loss: 0.6931 - val_auc: 0.5000 - val_regularization_loss: 0.0000e+00 - val_loss_batch: 0.6931\n",
      "Epoch 2/3\n",
      "43/43 [==============================] - 1s 15ms/step - loss: 0.6931 - auc: 0.5016 - regularization_loss: 0.0000e+00 - loss_batch: 0.6931 - val_loss: 0.6932 - val_auc: 0.4999 - val_regularization_loss: 0.0000e+00 - val_loss_batch: 0.6932\n",
      "Epoch 3/3\n",
      "43/43 [==============================] - 1s 15ms/step - loss: 0.6929 - auc: 0.5127 - regularization_loss: 0.0000e+00 - loss_batch: 0.6929 - val_loss: 0.6935 - val_auc: 0.5000 - val_regularization_loss: 0.0000e+00 - val_loss_batch: 0.6936\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<keras.callbacks.History at 0x7fdb2d8d7880>"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dlrm = create_dlrm(\n",
    "    optimizer=\"adam\",\n",
    "    schema=schema,\n",
    "    target_column=target_column,\n",
    "    embedding_dim=64\n",
    ")\n",
    "\n",
    "dlrm.fit(train_rank, validation_data=valid_rank, batch_size=16 * 1024, epochs=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "46cdbd5a-ea3d-4523-99cb-a22bc92c4dde",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[[24]\n",
      " [ 5]\n",
      " [ 5]\n",
      " [ 2]\n",
      " [ 1]\n",
      " [35]\n",
      " [ 3]\n",
      " [14]], shape=(8, 1), dtype=int32) tf.Tensor(\n",
      "[[ 4]\n",
      " [ 3]\n",
      " [10]\n",
      " [ 1]\n",
      " [ 1]\n",
      " [31]\n",
      " [14]\n",
      " [ 2]], shape=(8, 1), dtype=int32)\n"
     ]
    }
   ],
   "source": [
    "# Check out one batch of users\n",
    "loader = mm.Loader(valid_rank, batch_size=8, shuffle=False)\n",
    "batch = next(iter(loader))\n",
    "print(batch[0]['user_id'], batch[0]['item_id'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "f768e0ec-8770-4471-bcad-fde461300988",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(8, 1), dtype=float32, numpy=\n",
       "array([[0.48569584],\n",
       "       [0.51287526],\n",
       "       [0.49173698],\n",
       "       [0.51630616],\n",
       "       [0.5096519 ],\n",
       "       [0.50963724],\n",
       "       [0.48768458],\n",
       "       [0.49276185]], dtype=float32)>"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Test the DLRM model I/O with the user batch\n",
    "dlrm(batch[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d811ff92-49f5-481b-a4b1-513edbf80c51",
   "metadata": {},
   "source": [
    "### 4.3 - Test Retrieval and Rank together\n",
    "\n",
    "So now that we have a **Candidate Retrieval Model** and a **Ranking Model** (DLRM) we can test the end to end flow locally. In order to do this live, like in the next multi-stage section, you need a **Feature Store** (Redis/Feast), **ANN Index** (Redis), and means to host/serve the entire DAG of operations (Triton).\n",
    "\n",
    "*Below is a very brute force approach for a single batch of users.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "ac6c095b-0e8d-4f06-b4aa-5bceeb2511cd",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Test full pipeline\n",
    "import numpy as np\n",
    "\n",
    "# Reload Dataset\n",
    "train_retrieval = Dataset(os.path.join(OUTPUT_RETRIEVAL_DATA_DIR, \"train\", \"*.parquet\"), part_size=\"500MB\")\n",
    "schema = train_retrieval.schema.select_by_tag([Tags.ITEM_ID, Tags.USER_ID, Tags.ITEM, Tags.USER]).without(['click'])\n",
    "train_retrieval.schema = schema\n",
    "\n",
    "# User and Item Offline Feature \"Stores\"\n",
    "item_fs = unique_rows_by_features(train_retrieval, Tags.ITEM, Tags.ITEM_ID).to_ddf().compute()\n",
    "user_fs = unique_rows_by_features(train_retrieval, Tags.USER, Tags.USER_ID).to_ddf().compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "99ddc0f2-3e9d-41cc-9793-4f93a11f5e17",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(8, 1), dtype=int32, numpy=\n",
       "array([[1],\n",
       "       [2],\n",
       "       [3],\n",
       "       [4],\n",
       "       [5],\n",
       "       [6],\n",
       "       [7],\n",
       "       [8]], dtype=int32)>"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# User batch loader\n",
    "user_loader = mm.Loader(unique_rows_by_features(train_retrieval, Tags.USER, Tags.USER_ID), batch_size=8, shuffle=False)\n",
    "\n",
    "# Sample User batch\n",
    "user_batch = next(iter(user_loader))\n",
    "users = user_batch[0]['user_id']\n",
    "users"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "ca5c2c58-0b1c-4bf7-b997-738a5a12f95c",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(8, 20), dtype=int32, numpy=\n",
       "array([[  2, 339,  53,  15,   7,   8, 104, 219, 313,  44, 391,  16, 581,\n",
       "        367, 122,  36, 131,  12,  49,  88],\n",
       "       [184, 267, 189,  14,  37, 242,   5,  13,   7,  43, 182, 209,  24,\n",
       "        323, 469, 344,  61, 315, 202, 346],\n",
       "       [  5,  39,  12, 353,  71, 384,  98,  47, 182, 496, 276, 226,  37,\n",
       "         67, 552, 196,   3, 423, 359,  86],\n",
       "       [  2,   4, 102,  53,   1, 239,  75,   3, 207, 291, 349,  31, 416,\n",
       "        303,  26, 103, 283, 131, 751, 328],\n",
       "       [ 11,   5, 421,  16,  37, 496, 226,  39, 448, 274,   7, 201, 251,\n",
       "        360, 279,  14, 625,  12, 150, 210],\n",
       "       [182,  71,  20, 384, 184,  61,   3,   5,  14, 315,  25,  39, 271,\n",
       "        234, 746,  33, 115, 253, 189, 285],\n",
       "       [ 53, 391,  49,  42, 339,  13,  44,  36,  77, 124,  15, 131,   2,\n",
       "         91,  64, 300, 335, 388, 177, 367],\n",
       "       [  2,   4, 334, 360,  27, 244, 688,  53, 243,  79, 287,   1, 416,\n",
       "         78,   3, 283,   8, 102, 751,  15]], dtype=int32)>"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Retrieve candidate Items for each User\n",
    "_, candidate_item_ids = topk_model(user_batch[0])\n",
    "candidate_item_ids"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ab7d2b0-dcfd-46b7-8cad-4c55353a61fa",
   "metadata": {},
   "source": [
    "### 4.4 - Define Functions to run Recommendation Pipeline\n",
    "\n",
    "Next we define our functions to run the entire pipeline (as tested above)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "033901bf-a1f3-4dcd-bdc0-ba06add7ba9b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Softmax sample function for ordering\n",
    "def softmax_sample(recs: np.array, scores: np.array) -> np.array:\n",
    "    arr = np.exp(scores)/sum(np.exp(scores))\n",
    "    top_item_idx = (-arr).argsort()\n",
    "    return recs[top_item_idx]\n",
    "\n",
    "def generate_recommendations_batch(topk_model, dlrm, input_data, batch_size: int, tags):\n",
    "    user_loader = mm.Loader(unique_rows_by_features(input_data, tags.USER, tags.USER_ID), batch_size=batch_size, shuffle=False)\n",
    "    \n",
    "    # Load a batch\n",
    "    for batch in user_loader:\n",
    "            \n",
    "        # Generate candidates for this batch of users\n",
    "        users = batch[0]['user_id']\n",
    "        _, candidate_item_ids = topk_model(batch[0])\n",
    "        \n",
    "        # For each user + candidate items, score with the DLRM\n",
    "        for user, candidates in zip(users.numpy(), candidate_item_ids.numpy()):\n",
    "            try:\n",
    "                num_recs = len(candidates)\n",
    "                user_id = user[0]\n",
    "\n",
    "                # get user features\n",
    "                user_features = user_fs[user_fs.user_id == user_id]\n",
    "                raw_user_id = user_features.user_id_raw.to_numpy()[0]\n",
    "                user_features = user_features.append([user_features]*(num_recs-1), ignore_index=True)\n",
    "\n",
    "                # get item features\n",
    "                item_features = item_fs[item_fs.item_id.isin(candidates)].reset_index(drop=True)\n",
    "                raw_item_ids = item_features.item_id_raw.to_numpy()\n",
    "\n",
    "                # combine into feature vectors\n",
    "                item_features[user_features.columns] = user_features\n",
    "                item_features = Dataset(item_features)\n",
    "                item_features.schema = schema.without(['click'])\n",
    "\n",
    "                # Score with ranking model\n",
    "                inputs = mm.Loader(item_features, batch_size=num_recs)\n",
    "                inputs = next(iter(inputs))\n",
    "                scores = dlrm(inputs[0]).numpy().reshape(-1)\n",
    "\n",
    "                # Rank\n",
    "                recs = softmax_sample(raw_item_ids, scores)\n",
    "\n",
    "                yield raw_user_id, recs\n",
    "            except Exception as e:\n",
    "                logging.info(user_id, str(e))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "f29e6b18-84c5-4d3d-9c9a-fbce08358f47",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(7,\n",
       " array([ 53, 132, 122,   6,  87,  18, 624,  11,  36,  15,  17, 409, 329,\n",
       "        363,  49,  44, 330, 104, 222,   4], dtype=int32))"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Test Recommendation Generator\n",
    "next(generate_recommendations_batch(topk_model, dlrm, train_retrieval, 32, Tags))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea442dd4-e570-4500-80ae-7996f9acaed5",
   "metadata": {},
   "source": [
    "## 5.0 - Write Recommendations to Redis\n",
    "\n",
    "In a batch-offline recommendation system, recommendations are generated in batches and are not generated in real-time. This means that the system will run on a schedule (e.g. once a day) to generate recommendations for all users at once. Redis is used to store these recommendations in a way that allows for quick retrieval when a user requests their recommendations.\n",
    "\n",
    "Redis is a high-performance in-memory data store that is well-suited for this use case because it can handle high read and write loads and can retrieve data very quickly. This allows the recommendation system to return recommendations to users in real-time, while still taking advantage of the offline batch processing to generate the recommendations.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "c8fd929c-ec80-49d8-ba44-7a7a81ce9a26",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import needed Redis libraries\n",
    "import asyncio\n",
    "import typing as t\n",
    "import redis.asyncio as redis\n",
    "from redis.commands.json.path import Path\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "1dd3991a-3810-4683-bad6-8438ba4a5fd4",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "async def store_recommendations(recommendations: t.Iterable, redis_conn: redis.Redis):\n",
    "    \"\"\"\n",
    "    Store the recommendations generated for each User in Redis.\n",
    "\n",
    "    Parameters:\n",
    "        recommendations (t.Iterable): A generator over a dictionary where the keys are user_ids and the values are lists of recommendations\n",
    "        redis_conn (redis.Redis): A Redis connection object used to store the recommendations in Redis\n",
    "    \"\"\"\n",
    "    async def store_as_json(user_id: str, recs: list):\n",
    "        \"\"\"\n",
    "        Store an individual user's latest recommendations in Redis.\n",
    "\n",
    "        Parameters:\n",
    "            user_id (str): The user id of the user whose recommendations are being stored\n",
    "            recs (list): A list of item_ids representing the recommendations for the user\n",
    "        \"\"\"\n",
    "        entry = {\n",
    "            \"user_id\": int(user_id),\n",
    "            \"recommendations\": [int(rec) for rec in recs]\n",
    "        }\n",
    "        # Set the JSON object in Redis\n",
    "        await redis_conn.json().set(f\"USER:{user_id}\", Path.root_path(), entry)\n",
    "        \n",
    "    # Write the recommendations to Redis as a JSON object\n",
    "    for user_id, recs in recommendations:\n",
    "        await store_as_json(user_id, recs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "8bd57904-4419-4b77-9ece-eb323c0a3349",
   "metadata": {},
   "outputs": [],
   "source": [
    "redis_conn = redis.Redis(\n",
    "    host=\"redis-inference-store\",\n",
    "    port=6379,\n",
    "    decode_responses=True\n",
    ")\n",
    "\n",
    "# Create Recommendation Denerator\n",
    "recommendations = generate_recommendations_batch(topk_model, dlrm, train_retrieval, 32, Tags)\n",
    "\n",
    "# Run the process - may take a few minutes\n",
    "await store_recommendations(recommendations, redis_conn=redis_conn)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe25e4e6-11af-42b5-922d-c5252ec335fe",
   "metadata": {},
   "source": [
    "## 6.0 - Fetch Recommendations from Redis\n",
    "\n",
    "Now that you've written a list of recommended item IDs for each User - we can use the Redis CLI to take a peek at one example. Your application can now fetch these when needed (at low latency)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "b2b44758-b274-4f39-b1f0-b21193f2e94b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\"{\\\"user_id\\\":1,\\\"recommendations\\\":[26,222,377,18,443,12,612,62,289,25,615,14,257,13,459,624,419,19,8,244]}\"\n"
     ]
    }
   ],
   "source": [
    "!redis-cli -h redis-inference-store -p 6379 JSON.GET USER:1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "ee5f1991-5160-45aa-97a9-4b13fe877581",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simple benchmark with async Python client\n",
    "times = []\n",
    "for i in range(300):\n",
    "    t = time.time()\n",
    "    await redis_conn.json().get(f\"USER:{i}\")\n",
    "    times.append(time.time()-t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "6ce7bce2-8356-4b3e-836d-23768a60bd2c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.00020427783330281574"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Average read time from Redis --> less than a ms likely\n",
    "np.average(times)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b19062ef",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## 7.0 - Export Models"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "abf8cb35",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "So far we have trained and evaluated our models for the recsys pipeline. We need to export the models and artifacts in order to deploy them \"online\" in the next notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a01df11e",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### 7.1 - Export Retrieval Model"
   ]
  },
  {
   "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": 39,
   "id": "07c68eb4",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "query_tower = retrieval_model.query_encoder\n",
    "query_tower.save(os.path.join(DATA_DIR, \"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_DIR, 'query_tower'))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4c031f8",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### 7.2 - Export 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. These will get stored in our Feature Store later on."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "1f861ac4",
   "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",
       "      <th>user_id_raw</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",
       "      <td>7</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",
       "      <td>8</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",
       "      <td>6</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",
       "      <td>9</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",
       "      <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  user_id_raw  \n",
       "0            1                1            7  \n",
       "1            2                2            8  \n",
       "2            3                3            6  \n",
       "3            4                4            9  \n",
       "4            5                5            5  "
      ]
     },
     "execution_count": 40,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "user_features = (\n",
    "    unique_rows_by_features(train_retrieval, Tags.USER, Tags.USER_ID).compute().reset_index(drop=True)\n",
    ")\n",
    "user_features.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "33758417",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# save to disk\n",
    "user_features.to_parquet(os.path.join(DATA_DIR, \"user_features.parquet\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "184f6cf6-cc5e-459e-9d87-5fdcd01d8c85",
   "metadata": {},
   "source": [
    "### 7.3 - Export User Embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "5d31a12c-ca41-4b02-9288-7d6b31d54e39",
   "metadata": {},
   "outputs": [],
   "source": [
    "queries = retrieval_model.query_embeddings(Dataset(user_features, schema=schema), batch_size=1024, index=Tags.USER_ID)\n",
    "query_embs_df = queries.compute(scheduler=\"synchronous\").reset_index()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "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.005831</td>\n",
       "      <td>-0.017723</td>\n",
       "      <td>-0.005784</td>\n",
       "      <td>0.049866</td>\n",
       "      <td>0.026280</td>\n",
       "      <td>0.049169</td>\n",
       "      <td>-0.054866</td>\n",
       "      <td>0.022130</td>\n",
       "      <td>0.036576</td>\n",
       "      <td>...</td>\n",
       "      <td>-0.021381</td>\n",
       "      <td>-0.000743</td>\n",
       "      <td>-0.060986</td>\n",
       "      <td>-0.037169</td>\n",
       "      <td>0.040978</td>\n",
       "      <td>0.017704</td>\n",
       "      <td>0.003427</td>\n",
       "      <td>-0.028408</td>\n",
       "      <td>0.047877</td>\n",
       "      <td>-0.025089</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2</td>\n",
       "      <td>-0.050697</td>\n",
       "      <td>0.068015</td>\n",
       "      <td>-0.071699</td>\n",
       "      <td>0.043323</td>\n",
       "      <td>-0.001567</td>\n",
       "      <td>0.032226</td>\n",
       "      <td>-0.120845</td>\n",
       "      <td>-0.016043</td>\n",
       "      <td>0.016297</td>\n",
       "      <td>...</td>\n",
       "      <td>0.055824</td>\n",
       "      <td>0.003816</td>\n",
       "      <td>-0.061081</td>\n",
       "      <td>-0.002499</td>\n",
       "      <td>0.051202</td>\n",
       "      <td>-0.060756</td>\n",
       "      <td>-0.022746</td>\n",
       "      <td>0.015439</td>\n",
       "      <td>0.078743</td>\n",
       "      <td>-0.037100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3</td>\n",
       "      <td>0.039121</td>\n",
       "      <td>-0.006548</td>\n",
       "      <td>-0.148172</td>\n",
       "      <td>0.098371</td>\n",
       "      <td>-0.025619</td>\n",
       "      <td>0.046444</td>\n",
       "      <td>-0.166427</td>\n",
       "      <td>-0.044223</td>\n",
       "      <td>0.086496</td>\n",
       "      <td>...</td>\n",
       "      <td>0.000787</td>\n",
       "      <td>0.065980</td>\n",
       "      <td>-0.131302</td>\n",
       "      <td>-0.064237</td>\n",
       "      <td>0.037129</td>\n",
       "      <td>-0.035728</td>\n",
       "      <td>-0.082380</td>\n",
       "      <td>0.062503</td>\n",
       "      <td>0.179321</td>\n",
       "      <td>-0.091110</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>4</td>\n",
       "      <td>-0.031507</td>\n",
       "      <td>-0.061412</td>\n",
       "      <td>0.008147</td>\n",
       "      <td>0.063988</td>\n",
       "      <td>0.013863</td>\n",
       "      <td>0.104590</td>\n",
       "      <td>-0.053230</td>\n",
       "      <td>0.054759</td>\n",
       "      <td>0.083436</td>\n",
       "      <td>...</td>\n",
       "      <td>-0.031580</td>\n",
       "      <td>-0.020756</td>\n",
       "      <td>-0.025745</td>\n",
       "      <td>0.008281</td>\n",
       "      <td>0.030552</td>\n",
       "      <td>0.033858</td>\n",
       "      <td>0.061298</td>\n",
       "      <td>0.085968</td>\n",
       "      <td>0.093174</td>\n",
       "      <td>-0.027022</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>5</td>\n",
       "      <td>-0.029882</td>\n",
       "      <td>0.066317</td>\n",
       "      <td>-0.077858</td>\n",
       "      <td>0.026251</td>\n",
       "      <td>-0.046349</td>\n",
       "      <td>0.006982</td>\n",
       "      <td>-0.123309</td>\n",
       "      <td>-0.120219</td>\n",
       "      <td>0.059101</td>\n",
       "      <td>...</td>\n",
       "      <td>-0.074184</td>\n",
       "      <td>0.012038</td>\n",
       "      <td>0.067444</td>\n",
       "      <td>-0.047427</td>\n",
       "      <td>0.108614</td>\n",
       "      <td>0.067772</td>\n",
       "      <td>0.048659</td>\n",
       "      <td>-0.034214</td>\n",
       "      <td>0.144912</td>\n",
       "      <td>-0.046109</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.005831 -0.017723 -0.005784  0.049866  0.026280  0.049169   \n",
       "1        2 -0.050697  0.068015 -0.071699  0.043323 -0.001567  0.032226   \n",
       "2        3  0.039121 -0.006548 -0.148172  0.098371 -0.025619  0.046444   \n",
       "3        4 -0.031507 -0.061412  0.008147  0.063988  0.013863  0.104590   \n",
       "4        5 -0.029882  0.066317 -0.077858  0.026251 -0.046349  0.006982   \n",
       "\n",
       "          6         7         8  ...        54        55        56        57  \\\n",
       "0 -0.054866  0.022130  0.036576  ... -0.021381 -0.000743 -0.060986 -0.037169   \n",
       "1 -0.120845 -0.016043  0.016297  ...  0.055824  0.003816 -0.061081 -0.002499   \n",
       "2 -0.166427 -0.044223  0.086496  ...  0.000787  0.065980 -0.131302 -0.064237   \n",
       "3 -0.053230  0.054759  0.083436  ... -0.031580 -0.020756 -0.025745  0.008281   \n",
       "4 -0.123309 -0.120219  0.059101  ... -0.074184  0.012038  0.067444 -0.047427   \n",
       "\n",
       "         58        59        60        61        62        63  \n",
       "0  0.040978  0.017704  0.003427 -0.028408  0.047877 -0.025089  \n",
       "1  0.051202 -0.060756 -0.022746  0.015439  0.078743 -0.037100  \n",
       "2  0.037129 -0.035728 -0.082380  0.062503  0.179321 -0.091110  \n",
       "3  0.030552  0.033858  0.061298  0.085968  0.093174 -0.027022  \n",
       "4  0.108614  0.067772  0.048659 -0.034214  0.144912 -0.046109  \n",
       "\n",
       "[5 rows x 65 columns]"
      ]
     },
     "execution_count": 43,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query_embs_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "595b1578",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### 7.4 - Export Item Features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "3ae320ec",
   "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",
       "      <th>item_id_raw</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>7</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",
       "      <td>6</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",
       "      <td>8</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",
       "      <td>9</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",
       "      <td>5</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   item_id  item_category  item_shop  item_brand  item_id_raw\n",
       "0        1              1          1           1            7\n",
       "1        2              2          2           2            6\n",
       "2        3              3          3           3            8\n",
       "3        4              4          4           4            9\n",
       "4        5              5          5           5            5"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "item_features = (\n",
    "    unique_rows_by_features(train_retrieval, Tags.ITEM, Tags.ITEM_ID).compute().reset_index(drop=True)\n",
    ")\n",
    "item_features.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "40102380",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# save to disk\n",
    "item_features.to_parquet(os.path.join(DATA_DIR, \"item_features.parquet\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f57f5a3e",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### 7.5 - Export Item Embeddings\n",
    "These will get stored in our ANN Index (in Redis)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "30f50d2a-dbf3-4459-af46-1a8194730896",
   "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>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.034885</td>\n",
       "      <td>-0.000131</td>\n",
       "      <td>0.018455</td>\n",
       "      <td>0.037430</td>\n",
       "      <td>0.026332</td>\n",
       "      <td>0.012729</td>\n",
       "      <td>0.006760</td>\n",
       "      <td>0.069112</td>\n",
       "      <td>0.044133</td>\n",
       "      <td>0.031141</td>\n",
       "      <td>...</td>\n",
       "      <td>-0.027153</td>\n",
       "      <td>-0.029950</td>\n",
       "      <td>-0.020070</td>\n",
       "      <td>-0.067773</td>\n",
       "      <td>0.002420</td>\n",
       "      <td>-0.001353</td>\n",
       "      <td>-0.055582</td>\n",
       "      <td>0.042481</td>\n",
       "      <td>0.013875</td>\n",
       "      <td>0.021228</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>0.021357</td>\n",
       "      <td>-0.026375</td>\n",
       "      <td>0.069090</td>\n",
       "      <td>-0.011445</td>\n",
       "      <td>0.025277</td>\n",
       "      <td>-0.010337</td>\n",
       "      <td>0.008437</td>\n",
       "      <td>0.042574</td>\n",
       "      <td>0.060663</td>\n",
       "      <td>-0.004809</td>\n",
       "      <td>...</td>\n",
       "      <td>-0.037270</td>\n",
       "      <td>-0.039209</td>\n",
       "      <td>0.013558</td>\n",
       "      <td>-0.006484</td>\n",
       "      <td>-0.029601</td>\n",
       "      <td>0.073999</td>\n",
       "      <td>0.009857</td>\n",
       "      <td>-0.022534</td>\n",
       "      <td>-0.009440</td>\n",
       "      <td>-0.025069</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>-0.018197</td>\n",
       "      <td>0.017502</td>\n",
       "      <td>0.002263</td>\n",
       "      <td>0.008534</td>\n",
       "      <td>0.015912</td>\n",
       "      <td>0.006360</td>\n",
       "      <td>-0.001660</td>\n",
       "      <td>0.007613</td>\n",
       "      <td>0.054932</td>\n",
       "      <td>0.042344</td>\n",
       "      <td>...</td>\n",
       "      <td>-0.045789</td>\n",
       "      <td>0.033707</td>\n",
       "      <td>-0.025606</td>\n",
       "      <td>-0.020231</td>\n",
       "      <td>0.068983</td>\n",
       "      <td>0.030158</td>\n",
       "      <td>-0.054312</td>\n",
       "      <td>-0.006741</td>\n",
       "      <td>0.026637</td>\n",
       "      <td>-0.040934</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>-0.018756</td>\n",
       "      <td>-0.057435</td>\n",
       "      <td>0.027142</td>\n",
       "      <td>0.069214</td>\n",
       "      <td>-0.014137</td>\n",
       "      <td>0.063484</td>\n",
       "      <td>0.049648</td>\n",
       "      <td>-0.000459</td>\n",
       "      <td>0.041440</td>\n",
       "      <td>0.020341</td>\n",
       "      <td>...</td>\n",
       "      <td>-0.050948</td>\n",
       "      <td>-0.007804</td>\n",
       "      <td>0.001069</td>\n",
       "      <td>-0.059237</td>\n",
       "      <td>-0.018273</td>\n",
       "      <td>-0.005572</td>\n",
       "      <td>-0.017192</td>\n",
       "      <td>0.033178</td>\n",
       "      <td>0.050670</td>\n",
       "      <td>0.040354</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>0.044985</td>\n",
       "      <td>0.015847</td>\n",
       "      <td>-0.041081</td>\n",
       "      <td>-0.006620</td>\n",
       "      <td>-0.003196</td>\n",
       "      <td>-0.045210</td>\n",
       "      <td>-0.031615</td>\n",
       "      <td>-0.093638</td>\n",
       "      <td>0.007464</td>\n",
       "      <td>0.033940</td>\n",
       "      <td>...</td>\n",
       "      <td>-0.014779</td>\n",
       "      <td>0.057923</td>\n",
       "      <td>-0.015743</td>\n",
       "      <td>-0.048929</td>\n",
       "      <td>0.000438</td>\n",
       "      <td>-0.043618</td>\n",
       "      <td>-0.137103</td>\n",
       "      <td>-0.019580</td>\n",
       "      <td>0.025585</td>\n",
       "      <td>0.028937</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5 rows × 64 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                0         1         2         3         4         5         6  \\\n",
       "item_id                                                                         \n",
       "1       -0.034885 -0.000131  0.018455  0.037430  0.026332  0.012729  0.006760   \n",
       "2        0.021357 -0.026375  0.069090 -0.011445  0.025277 -0.010337  0.008437   \n",
       "3       -0.018197  0.017502  0.002263  0.008534  0.015912  0.006360 -0.001660   \n",
       "4       -0.018756 -0.057435  0.027142  0.069214 -0.014137  0.063484  0.049648   \n",
       "5        0.044985  0.015847 -0.041081 -0.006620 -0.003196 -0.045210 -0.031615   \n",
       "\n",
       "                7         8         9  ...        54        55        56  \\\n",
       "item_id                                ...                                 \n",
       "1        0.069112  0.044133  0.031141  ... -0.027153 -0.029950 -0.020070   \n",
       "2        0.042574  0.060663 -0.004809  ... -0.037270 -0.039209  0.013558   \n",
       "3        0.007613  0.054932  0.042344  ... -0.045789  0.033707 -0.025606   \n",
       "4       -0.000459  0.041440  0.020341  ... -0.050948 -0.007804  0.001069   \n",
       "5       -0.093638  0.007464  0.033940  ... -0.014779  0.057923 -0.015743   \n",
       "\n",
       "               57        58        59        60        61        62        63  \n",
       "item_id                                                                        \n",
       "1       -0.067773  0.002420 -0.001353 -0.055582  0.042481  0.013875  0.021228  \n",
       "2       -0.006484 -0.029601  0.073999  0.009857 -0.022534 -0.009440 -0.025069  \n",
       "3       -0.020231  0.068983  0.030158 -0.054312 -0.006741  0.026637 -0.040934  \n",
       "4       -0.059237 -0.018273 -0.005572 -0.017192  0.033178  0.050670  0.040354  \n",
       "5       -0.048929  0.000438 -0.043618 -0.137103 -0.019580  0.025585  0.028937  \n",
       "\n",
       "[5 rows x 64 columns]"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "item_embs = retrieval_model.candidate_embeddings(Dataset(item_features, schema=schema), batch_size=1024, index=Tags.ITEM_ID)\n",
    "item_embs_df = item_embs.compute(scheduler=\"synchronous\")\n",
    "item_embs_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "5d152e82",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# save to disk\n",
    "item_embs_df.to_parquet(os.path.join(DATA_DIR, \"item_embeddings.parquet\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32144ec6-4c10-49bc-ba58-fb21f32e7602",
   "metadata": {},
   "source": [
    "### 7.6 - Export DLRM"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "edae0a21-c4d1-47e9-b386-7af0fed63268",
   "metadata": {},
   "outputs": [],
   "source": [
    "dlrm.save(os.path.join(DATA_DIR, \"dlrm\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "442c3f88-cda8-4202-b023-4f926fd52d2d",
   "metadata": {},
   "source": [
    "## Conclusion"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "1eed43c3",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "No we have learned how to train and evaluate your Two-Tower **Candidate Retrieval Model** and **DLRM**. We also wrote recommendations \"offline\" to a key-value store (Redis) for use in our applications.\n",
    "\n",
    "![img](https://raw.githubusercontent.com/RedisVentures/Redis-Recsys/master/assets/OfflineBatchRecsys.png)\n",
    "\n",
    "*This completes the **\"Offline Batch Recommender System\"** intro with Redis + NVIDIA*. In the next tutorial, we will focus on deploying these assets to the [Triton Inference Server](https://github.com/triton-inference-server/server) for live \"online\" recommendations."
   ]
  }
 ],
 "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.9 (default, Apr 13 2022, 08:48:06) \n[Clang 13.1.6 (clang-1316.0.21.2.5)]"
  },
  "merlin": {
   "containers": [
    "nvcr.io/nvidia/merlin/merlin-tensorflow:latest"
   ]
  },
  "vscode": {
   "interpreter": {
    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}