{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tutorial for the EvalRS competition\n",
    "# Retrieval models with Merlin\n",
    "\n",
    "In this notebook, we present a tutorial on how to use the open-source [NVIDIA Merlin](https://github.com/NVIDIA-Merlin/) framework to build and train a retrieval model for EvalRS competition. We are going to use the [NVTabular](https://github.com/NVIDIA-Merlin/NVTabular) library for preprocessing and [Merlin Models](https://github.com/NVIDIA-Merlin/models) library for building and training Tensorflow-based retrieval models.\n",
    "\n",
    "### Retrieval models and Two-stage RecSys\n",
    "Retrieval models are recsys scalable models that are able to retrieve a large number of candidate items for recommendation. They are typically used in two-stage recsys pipelines, where the retrieval stage scores hundreds of thousands or millions of items and then the ranking stage model rescores the candidate items, using for that more features or with a more powerful architecture.  \n",
    "For ML-based candidate retrieval model, as it needs to quickly score millions of items for a given user, the retrieval models typically produce recommendation scores by just computing the dot product between user and item representations. Popular choices of such models are Matrix Factorization, which learns low-rank user and item embeddings, and the Two-Tower architecture, which is a neural network with two MLP towers where both user and item features are fed to generate user and item embeddings in the output.  \n",
    "\n",
    "In this notebook you will learn how to implement a pipeline for the competition that preprocess the available categorical and continuous features in a suitable format for neural networks and trains retrieval models (Matrix Factorization and Two-Tower architecture)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Installs EvalRS dependencies\n",
    "#!pip install -r ../../requirements.txt\n",
    "# Installs Merlin dependencies\n",
    "#!pip install -r requirements.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "_Basic imports, read the credentials from the env file_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "import shutil\n",
    "# Imports EvalRS dependencies\n",
    "sys.path.insert(0, '../../')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv('../../upload.env')\n",
    "\n",
    "EMAIL = os.getenv('EMAIL')  # the e-mail you used to sign up\n",
    "assert EMAIL != '' and EMAIL is not None\n",
    "BUCKET_NAME = os.getenv('BUCKET_NAME') # you received it in your e-mail\n",
    "PARTICIPANT_ID = os.getenv('PARTICIPANT_ID') # you received it in your e-mail\n",
    "AWS_ACCESS_KEY = os.getenv('AWS_ACCESS_KEY') # you received it in your e-mail\n",
    "AWS_SECRET_KEY = os.getenv('AWS_SECRET_KEY') # you received it in your e-mail"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "_Specify some other global variables to improve local iteration and debugging, for example setting a LIMIT to work with a smaller, faster test set_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "LIMIT = 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "_NOTE: as long as there is a limit specified, the runner won't upload results: make sure to have LIMIT=0 when you want to submit to the leaderboard!_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from evaluation.EvalRSRunner import EvalRSRunner\n",
    "from evaluation.EvalRSRunner import ChallengeDataset\n",
    "from reclist.abstractions import RecModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "dataset = ChallengeDataset(force_download=False)  # note, if YES, the dataset will be donwloaded again"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "runner = EvalRSRunner(\n",
    "    dataset=dataset,\n",
    "    aws_access_key_id=AWS_ACCESS_KEY,\n",
    "    aws_secret_access_key=AWS_SECRET_KEY,\n",
    "    participant_id=PARTICIPANT_ID,\n",
    "    bucket_name=BUCKET_NAME,\n",
    "    email=EMAIL\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creates a retrieval recsys pipeline with Merlin"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nvtabular as nvt\n",
    "import nvtabular.ops as ops\n",
    "from merlin.dag import ColumnSelector\n",
    "from merlin.schema import Schema, Tags"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import merlin.models.tf as mm\n",
    "from merlin.io import Dataset\n",
    "from merlin.models.tf.dataset import BatchedDataset\n",
    "from merlin.schema.tags import Tags\n",
    "from merlin.models.utils import schema_utils\n",
    "from merlin.models.tf.core.transformations import PopularityLogitsCorrection\n",
    "import tensorflow as tf\n",
    "from tensorflow.keras import regularizers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_item_frequencies(train_ds):   \n",
    "    \"\"\"Utility function that returns a TF tensor with the items (tracks) frequency\"\"\"\n",
    "    schema = train_ds.schema\n",
    "    # Gets the item ids cardinality\n",
    "    item_id_feature = schema.select_by_tag(Tags.ITEM_ID)\n",
    "    item_id_feature_name = item_id_feature.column_names[0]\n",
    "    cardinalities = schema_utils.categorical_cardinalities(schema)\n",
    "    item_id_cardinality = cardinalities[item_id_feature_name]\n",
    "\n",
    "    item_id_feature_name = schema.select_by_tag(Tags.ITEM_ID).column_names[0]\n",
    "\n",
    "    item_frequencies_df = (\n",
    "        train_ds.to_ddf()\n",
    "        .groupby(item_id_feature_name)\n",
    "        .size()\n",
    "        .to_frame(\"freq\")\n",
    "        .compute()\n",
    "    )\n",
    "    assert len(item_frequencies_df) <= item_id_cardinality\n",
    "    assert item_frequencies_df.index.max() < item_id_cardinality\n",
    "\n",
    "    # Completing the missing item ids and filling freq with 0\n",
    "    item_frequencies_df = item_frequencies_df.reindex(\n",
    "        np.arange(0, item_id_cardinality)\n",
    "    ).fillna(0)\n",
    "    assert len(item_frequencies_df) == item_id_cardinality\n",
    "\n",
    "    item_frequencies_df = item_frequencies_df.sort_index()\n",
    "    item_frequencies_df[\"dummy\"] = 1\n",
    "    item_frequencies_df[\"expected_id\"] = item_frequencies_df[\"dummy\"].cumsum() - 1\n",
    "    assert (\n",
    "        item_frequencies_df.index == item_frequencies_df[\"expected_id\"]\n",
    "    ).all(), f\"The item id feature ({item_id_feature_name}) should be contiguous from 0 to {item_id_cardinality-1}\"\n",
    "\n",
    "    item_frequencies = tf.convert_to_tensor(\n",
    "        item_frequencies_df[\"freq\"].values\n",
    "    )\n",
    "\n",
    "    return item_frequencies"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### NVTabular\n",
    "For preprocessing, we use the [NVTabular](https://github.com/NVIDIA-Merlin/NVTabular) library, which provides very handful features for common operations like (label) encoding categorical features for embeddings and normalization for continuous features.\n",
    "NVTabular works both with CPUs and GPUs. For this example we are installing NVTabular using `pip` for CPU usage, but if you want to speedup preprocessing with GPUs you can install NVTabular using `conda` as explained in its [repo](https://github.com/NVIDIA-Merlin/NVTabular).  \n",
    "\n",
    "In the `MyRetrievalModel` class, the `get_nvtabular_preproc_workflow()` method defines the preprocessing workflow, which is in the `preprocess_dataset()` method."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Merlin Models\n",
    "For model definition and training we use [Merlin Models](https://github.com/NVIDIA-Merlin/models). It provides a Tensorflow (Keras) API where you can easily build state-of-the-art retrieval and ranking models. One of the core ideas of Models is to leverage the schema generated during preprocessing by **NVTabular** to create automatically the necessary embedding tables for categorical features and define the target of the model.  You can find more information about retrieval with Merlin Models in this [example notebook](https://github.com/NVIDIA-Merlin/models/blob/main/examples/05-Retrieval-Model.ipynb).\n",
    "\n",
    "Here we will be implementing two retrieval models: **Matrix Factorization** and **Two-Tower architecture**.\n",
    "You will notice that we create a base `MyRetrievalModel` class, which we specialize later as `MyMFModel` and `MyTwoTowerModel` to override the `get_model()` method to return a `MatrixFactorizationModel` or `TwoTowerModel` instance."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Defining your RecModel\n",
    "\n",
    "Here we define our `RecModel` for EvalRS, which must implement the `train()` and the `predict()` methods. Within `train()` we receive a dataframe with the fold data and we need to define our full pipeline for preprocessing and training data. The `predict()` method receives a dataframe with the user ids, and it expects you to return the top-k predicted items for each user."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyRetrievalModel(RecModel):\n",
    "    \n",
    "    def __init__(self, items_df: pd.DataFrame, users_df: pd.DataFrame, top_k: int = 100, \n",
    "                 predict_batch_size=1024, **kwargs):\n",
    "        super(RecModel, self).__init__()\n",
    "        self.items_df = items_df\n",
    "        self.users_df = users_df\n",
    "        self.top_k = top_k\n",
    "        self.predict_batch_size = predict_batch_size\n",
    "        self.hparams = kwargs\n",
    "    \n",
    "    def get_nvtabular_preproc_workflow(self):\n",
    "        \"\"\"Defines an NVTabular preprocessing workflow\"\"\"\n",
    "        \n",
    "        user_id_col = [\"user_id\"] \n",
    "        user_cat_cols = [\"country\", \"gender\"]\n",
    "        user_age_col = [\"age\"]\n",
    "\n",
    "        user_continuous_cols = [\n",
    "            'novelty_artist_avg_month', 'novelty_artist_avg_6months', 'novelty_artist_avg_year',\n",
    "            'mainstreaminess_avg_month', 'mainstreaminess_avg_6months', 'mainstreaminess_avg_year', 'mainstreaminess_global',\n",
    "            'relative_le_per_weekday1', 'relative_le_per_weekday2', 'relative_le_per_weekday3',\n",
    "             'relative_le_per_weekday4', 'relative_le_per_weekday5','relative_le_per_weekday6', 'relative_le_per_weekday7',\n",
    "             'relative_le_per_hour0', 'relative_le_per_hour1','relative_le_per_hour2', 'relative_le_per_hour3',\n",
    "             'relative_le_per_hour4', 'relative_le_per_hour5', 'relative_le_per_hour6', 'relative_le_per_hour7', \n",
    "             'relative_le_per_hour8', 'relative_le_per_hour9', 'relative_le_per_hour10', 'relative_le_per_hour11',\n",
    "             'relative_le_per_hour12', 'relative_le_per_hour13', 'relative_le_per_hour14', 'relative_le_per_hour15',\n",
    "             'relative_le_per_hour16', 'relative_le_per_hour17','relative_le_per_hour18', 'relative_le_per_hour19',\n",
    "             'relative_le_per_hour20', 'relative_le_per_hour21', 'relative_le_per_hour22', 'relative_le_per_hour23', \n",
    "        ]\n",
    "\n",
    "        user_counts_cols = ['playcount', 'cnt_listeningevents', 'cnt_distinct_tracks', \n",
    "                                'cnt_distinct_artists', 'cnt_listeningevents_per_week']\n",
    "\n",
    "        item_id = [\"track_id\"]\n",
    "        item_features_cols = [\"artist_id\", \"album_id\"] \n",
    "\n",
    "\n",
    "        user_id = user_id_col >> ops.Categorify() >> ops.TagAsUserID()\n",
    "        user_feat_cat = user_cat_cols >> ops.Categorify() >> ops.TagAsUserFeatures()\n",
    "        age_boundaries = list(np.arange(0,100,5))\n",
    "        user_age = user_age_col >> ops.FillMissing(0) >> ops.Bucketize(age_boundaries) >> ops.Categorify() >> ops.TagAsUserFeatures()\n",
    "        user_feat_cont = user_continuous_cols >> ops.FillMedian() >> ops.Normalize() >> ops.TagAsUserFeatures()\n",
    "        user_feat_count = user_counts_cols >> ops.Clip(min_value=1) >> ops.FillMedian() >> ops.LogOp() >> ops.Normalize() >> ops.TagAsUserFeatures()\n",
    "        user_features = user_id + user_feat_cat + user_age + user_feat_cont + user_feat_count\n",
    "\n",
    "        item_id = item_id >> ops.Categorify() >> ops.TagAsItemID()\n",
    "        item_cat_feat = item_features_cols >> ops.Categorify() >> ops.TagAsItemFeatures()\n",
    "        item_features = item_id + item_cat_feat\n",
    "\n",
    "        outputs = user_features + item_features\n",
    "        workflow = nvt.Workflow(outputs)\n",
    "        return workflow\n",
    "    \n",
    "    def preprocess_dataset(self, events_df):\n",
    "        \"\"\"Preprocess the dataset using an NVTabular workflow\"\"\"        \n",
    "        nvt_dataset = nvt.Dataset(events_df)\n",
    "        \n",
    "        CATEG_MAPPING_FOLDER = 'categories/'\n",
    "        shutil.rmtree(CATEG_MAPPING_FOLDER, ignore_errors=True)\n",
    "        \n",
    "        nvt_workflow = self.get_nvtabular_preproc_workflow()\n",
    "        nvt_workflow.fit(nvt_dataset)\n",
    "        schema = nvt_workflow.output_schema\n",
    "        \n",
    "        # Loads mapping of categ features after the workflow is fit\n",
    "        self.user_ids_mapping_df = pd.read_parquet(CATEG_MAPPING_FOLDER+'unique.user_id.parquet')[['user_id']]\n",
    "        self.track_ids_mapping_df = pd.read_parquet(CATEG_MAPPING_FOLDER+'unique.track_id.parquet')[['track_id']]\n",
    "        \n",
    "        transformed_dataset = nvt_workflow.transform(nvt_dataset)        \n",
    "        transformed_df = transformed_dataset.persist().repartition(npartitions=10)  \n",
    "        \n",
    "        return transformed_df, schema\n",
    "    \n",
    "        \n",
    "    def get_item_retrieval_task(self):        \n",
    "        \"\"\"Defines the item retrieval task to be used by the retrieval models\"\"\"\n",
    "        items_frequencies = get_item_frequencies(self.train_dataset)\n",
    "        \n",
    "        post_logits = None\n",
    "        reg_factor = self.hparams['logq_correction_factor']\n",
    "        if reg_factor > 0.0:\n",
    "            post_logits = PopularityLogitsCorrection(\n",
    "                items_frequencies, reg_factor=reg_factor, schema=self.schema\n",
    "            )\n",
    "        \n",
    "        item_retrieval_task = mm.ItemRetrievalTask(self.schema, \n",
    "                                        logits_temperature=self.hparams['logits_temperature'],\n",
    "                                         post_logits=post_logits)\n",
    "        \n",
    "        return item_retrieval_task\n",
    "\n",
    "    \n",
    "    def get_model(self):\n",
    "        \"\"\"Defines the model architecture. Needs to be overridden by the child class\"\"\"\n",
    "        raise NotImplementedError()\n",
    "    \n",
    "    def compile_model(self, model):\n",
    "        \"\"\"Compiles the Keras model setting the metrics, loss, learning rate and optimizer\"\"\"\n",
    "        metrics = [mm.TopKMetricsAggregator(mm.RecallAt(20), mm.MRRAt(20))]\n",
    "        \n",
    "        lerning_rate = tf.keras.optimizers.schedules.ExponentialDecay(\n",
    "            self.hparams['lr'],\n",
    "            decay_steps=self.hparams['lr_decay_steps'],\n",
    "            decay_rate=self.hparams['lr_decay_rate'],\n",
    "            staircase=True,\n",
    "        )\n",
    "        optimizer = tf.keras.optimizers.Adam(learning_rate=lerning_rate)\n",
    "        loss = tf.keras.losses.CategoricalCrossentropy(\n",
    "            from_logits=True, label_smoothing=self.hparams['label_smoothing'],\n",
    "        )\n",
    "        \n",
    "        model.compile(optimizer, loss=loss, metrics=metrics, run_eagerly=False)\n",
    "        \n",
    "    def get_items_topk_recommender_model(\n",
    "        self,\n",
    "        train_dataset: Dataset, \n",
    "        schema, \n",
    "        model, \n",
    "        ):        \n",
    "        \"\"\"Converts a retrieval model into a Top-k recommender model, which\n",
    "        takes only user features as input, generates the user representations \n",
    "        (e.g. by taking the user embedding for MF or using the user tower to generate it)\n",
    "        and uses scores all cached item representations to return the most similar items \n",
    "        (P.s. This procedure would be done by an ANN engine in production)\"\"\"\n",
    "        item_features = schema.select_by_tag(Tags.ITEM).column_names\n",
    "        item_dataset = train_dataset.to_ddf()[item_features].drop_duplicates(subset=['track_id'], keep='last').compute()\n",
    "        item_dataset = Dataset(item_dataset)\n",
    "\n",
    "        return model.to_top_k_recommender(item_dataset, k=self.top_k)\n",
    "        \n",
    "    def train_model(self, model, train_dataset):\n",
    "        model.fit(\n",
    "            train_dataset,\n",
    "            epochs=self.hparams['epochs'],\n",
    "            batch_size=self.hparams['train_batch_size'],\n",
    "            #steps_per_epoch=5,\n",
    "            shuffle=True,\n",
    "            drop_last=True,\n",
    "            train_metrics_steps=100,\n",
    "        )\n",
    "        \n",
    "\n",
    "    def train(self, train_df: pd.DataFrame):\n",
    "        \"\"\"\n",
    "        Implement here your training logic. Since our example method is a simple random model,\n",
    "        we actually don't use any training data to build the model, but you should ;-)\n",
    "        At the end of training, make sure the class contains a trained model you can use in the predict method.\n",
    "        \"\"\"\n",
    "            \n",
    "        print(\"Merging events and user features\")\n",
    "        events_df = train_df.merge(self.users_df, on='user_id', how='inner')\n",
    "\n",
    "        print(\"Start preprocessing\")\n",
    "        transformed_df, schema = self.preprocess_dataset(events_df)\n",
    "        train_dataset = Dataset(transformed_df, schema=schema)\n",
    "            \n",
    "        self.train_dataset = train_dataset\n",
    "        self.schema = self.train_dataset.schema\n",
    "        \n",
    "        print(\"Building the model\")\n",
    "        model = self.get_model()\n",
    "        self.compile_model(model)                \n",
    "        \n",
    "        print(\"Start training\")\n",
    "        self.train_model(model, train_dataset)        \n",
    "        self.trained_model = model\n",
    "        \n",
    "        print(\"Preparing retrieval model for prediction\")\n",
    "        self.topk_rec_model = self.get_items_topk_recommender_model(train_dataset, schema, model)\n",
    "        \n",
    "        print(\"Caching users transformed features\")\n",
    "        self.users_schema = schema.select_by_tag(Tags.USER)\n",
    "        user_features = self.users_schema.column_names\n",
    "        self.users_transformed_df = train_dataset.to_ddf()[user_features].drop_duplicates(subset=['user_id'], keep='last').compute()\n",
    "        # Adding the raw (original) user id to the dataframe\n",
    "        self.users_transformed_df = self.users_transformed_df.merge(self.user_ids_mapping_df.rename({'user_id': 'raw_user_id'}, axis=1),\n",
    "                                                                    left_on='user_id', right_index=True)\n",
    "\n",
    "        print(\"Training completed!\")         \n",
    "\n",
    "    def predict(self, user_ids: pd.DataFrame) -> pd.DataFrame:\n",
    "        \"\"\"        \n",
    "        This function takes as input all the users that we want to predict the top-k items for, and \n",
    "        returns all the predicted songs.\n",
    "\n",
    "        While in this example is just a random generator, the same logic in your implementation \n",
    "        would allow for batch predictions of all the target data points.        \n",
    "        \"\"\"\n",
    "            \n",
    "        self.predict_user_ids = user_ids\n",
    "        \n",
    "        print(\"Start prediction\")\n",
    "        print(\"# users:\",len(user_ids))\n",
    "        test_users_df = user_ids.rename({'user_id': 'raw_user_id'}, axis=1).merge(self.users_transformed_df, \n",
    "                                                       on='raw_user_id', how='left')\n",
    "        self.test_users_df = test_users_df\n",
    "        test_users_found_df = test_users_df[~test_users_df[test_users_df.columns[1]].isna()]\n",
    "        test_users_not_found_df = test_users_df[test_users_df[test_users_df.columns[1]].isna()]        \n",
    "\n",
    "        test_users_dataset = Dataset(test_users_found_df, self.users_schema)\n",
    "        test_batched_dataset = BatchedDataset(\n",
    "            test_users_dataset, batch_size=self.predict_batch_size, shuffle=False, schema=self.users_schema\n",
    "        )        \n",
    "        \n",
    "        print(f\"Predicting Top-{self.top_k} items for test users\")\n",
    "        predictions = self.topk_rec_model.predict(test_batched_dataset)[1].astype(np.int32) \n",
    "        \n",
    "        print(f\"Converting user ids and predicted item ids to the original ids\")\n",
    "        predictions_converted = self.convert_prediction_item_ids(predictions)\n",
    "        \n",
    "        user_ids_found_converted = test_users_found_df['raw_user_id'].values.astype(np.int32)        \n",
    "        user_ids_not_found_converted = test_users_not_found_df['raw_user_id'].values.astype(np.int32)\n",
    "        \n",
    "        # Merging raw user id with top-k predictions\n",
    "        user_predictions_found_converted = np.concatenate((np.expand_dims(user_ids_found_converted, -1), predictions_converted), axis=1)\n",
    "        user_predictions_not_found_converted = np.concatenate((np.expand_dims(user_ids_not_found_converted, -1), \n",
    "                                                               np.zeros((user_ids_not_found_converted.shape[0], predictions_converted.shape[1]))), axis=1)\n",
    "        # Combining predictions of users found and not found\n",
    "        user_predictions_converted = np.vstack([user_predictions_found_converted, user_predictions_not_found_converted])        \n",
    "        user_predictions_df = pd.DataFrame(user_predictions_converted, columns=['user_id', *[str(i) for i in range(predictions_converted.shape[1])]]).set_index('user_id')\n",
    "        # Ensures predictions output dataframe is sorted the same as input user_ids order\n",
    "        user_predictions_df = user_predictions_df.loc[user_ids['user_id'].values]\n",
    "        self.user_predictions_df = user_predictions_df\n",
    "        print(\"Finish prediction\")\n",
    "        return user_predictions_df\n",
    "    \n",
    "    def convert_user_ids(self, user_ids):\n",
    "        \"\"\"Converts the encoded user ids into the original ids\"\"\"\n",
    "        raw_user_ids = self.user_ids_mapping_df['user_id'].loc[user_ids].values\n",
    "        return raw_user_ids\n",
    "    \n",
    "    def convert_prediction_item_ids(self,predictions):\n",
    "        \"\"\"Converts the encoded predicted item ids into the original item ids\"\"\"\n",
    "        raw_topk_predicted_item_ids = self.track_ids_mapping_df['track_id'].loc[predictions.reshape(-1)].values\n",
    "        raw_topk_predicted_item_ids = np.reshape(raw_topk_predicted_item_ids, (-1, predictions.shape[1]))\n",
    "        return raw_topk_predicted_item_ids"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Matrix Factorization\n",
    "\n",
    "Matrix Factorization is a [classic algorithm](https://arxiv.org/pdf/1205.2618.pdf) that have been used for Collaborative Filtering, where the user and item embeddings are factorized and the user preference over items is estimated by computing the dot product between the user embedding and item embeddings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyMFModel(MyRetrievalModel):\n",
    "    \n",
    "    def get_model(self):      \n",
    "        item_retrieval_task = self.get_item_retrieval_task()\n",
    "\n",
    "        model = mm.MatrixFactorizationModel(\n",
    "            self.schema,\n",
    "            dim=self.hparams['mf_dim'],\n",
    "            prediction_tasks=item_retrieval_task,\n",
    "            embeddings_l2_reg=self.hparams['embeddings_l2_reg']\n",
    "        )\n",
    "        \n",
    "        return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mf_model = MyMFModel(\n",
    "    items_df=dataset.df_tracks,\n",
    "    users_df=dataset.df_users,\n",
    "    \n",
    "    # Training hparams\n",
    "    epochs=5,\n",
    "    train_batch_size=8192,\n",
    "    lr=1e-3,\n",
    "    lr_decay_steps=100,\n",
    "    lr_decay_rate=0.96,\n",
    "    label_smoothing=0.0,\n",
    "    \n",
    "    # Model hparams\n",
    "    logq_correction_factor=1.0,\n",
    "    embeddings_l2_reg=5e-6,\n",
    "    logits_temperature=1.8,\n",
    "    mf_dim=128,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Runs the evaluation code: remember, if LIMIT is not 0, your submission won't be uploaded but the loop may still be useful for you to debug / iterate locally_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "runner.evaluate(model=mf_model, limit=LIMIT)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Two-tower architecture"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A Two-Tower Model consists of item (candidate) and user (query) encoder towers. The model learns representations (embeddings) for user and candidate items separately and user-item affinity is computed by dot product between the output of the towers. Differently from Matrix Factorization, multiple user/query and item features can be used for more accurate recommendations.\n",
    "\n",
    "<img src=\"./images/TwoTower.png\"  width=\"30%\">\n",
    "\n",
    "Image Adapted from: [Off-policy Learning in Two-stage Recommender Systems](https://dl.acm.org/doi/abs/10.1145/3366423.3380130)\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyTwoTowerModel(MyRetrievalModel):\n",
    "    \n",
    "    def get_model(self):   \n",
    "        item_retrieval_task = self.get_item_retrieval_task()\n",
    "        \n",
    "        model = mm.TwoTowerModel(\n",
    "            self.schema,\n",
    "            query_tower=mm.MLPBlock(\n",
    "                self.hparams['tt_mlp_layers'],\n",
    "                activation=self.hparams['tt_mlp_activation'],\n",
    "                no_activation_last_layer=True,    \n",
    "                dropout=self.hparams['tt_mlp_dropout'],                \n",
    "                kernel_regularizer=regularizers.l2(self.hparams['tt_mlp_l2_reg']),\n",
    "                bias_regularizer=regularizers.l2(self.hparams['tt_mlp_l2_reg']),\n",
    "            ),\n",
    "            embedding_options=mm.EmbeddingOptions(\n",
    "                infer_embedding_sizes=True,\n",
    "                infer_embedding_sizes_multiplier=self.hparams['tt_infer_embedding_sizes_multiplier'],\n",
    "                embeddings_l2_reg=self.hparams['embeddings_l2_reg'],\n",
    "            ),\n",
    "            prediction_tasks=item_retrieval_task\n",
    "        )\n",
    "\n",
    "        return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tt_model = MyTwoTowerModel(\n",
    "    items_df=dataset.df_tracks,\n",
    "    users_df=dataset.df_users,\n",
    "    \n",
    "    # Training hparams\n",
    "    epochs=5,\n",
    "    train_batch_size=8192,\n",
    "    lr=1e-3,\n",
    "    lr_decay_steps=100,\n",
    "    lr_decay_rate=0.96,\n",
    "    label_smoothing=0.0,\n",
    "    \n",
    "    # Model hparams\n",
    "    logq_correction_factor=1.0,\n",
    "    embeddings_l2_reg=1e-5,\n",
    "    logits_temperature=1.8,\n",
    "    tt_mlp_layers=[128,64],\n",
    "    tt_mlp_activation=\"relu\",\n",
    "    tt_mlp_dropout=0.3,\n",
    "    tt_mlp_l2_reg=5e-5,\n",
    "    tt_infer_embedding_sizes_multiplier=2.0\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Runs the evaluation code: remember, if LIMIT is not 0, your submission won't be uploaded but the loop may still be useful for you to debug / iterate locally_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "runner.evaluate(model=tt_model, limit=LIMIT)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook you have learned how build retrieval models (MF, Two-Tower) with the [Merlin](https://github.com/NVIDIA-Merlin/) open-source framework for the EvalRS competition. \n",
    "Feel free to improve these models using Tensorflow/Keras."
   ]
  }
 ],
 "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
