{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_dEaVsqSgNyQ"
      },
      "source": [
        "##### Copyright 2021 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4FyfuZX-gTKS"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sT8AyHRMNh41"
      },
      "source": [
        "# Recommend movies for users with TensorFlow Ranking\n",
        "\n",
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://www.tensorflow.org/ranking/tutorials/quickstart\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" /\u003eView on TensorFlow.org\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/ranking/blob/master/docs/tutorials/quickstart.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun in Google Colab\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://github.com/tensorflow/ranking/blob/master/docs/tutorials/quickstart.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca href=\"https://storage.googleapis.com/tensorflow_docs/ranking/docs/tutorials/quickstart.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/download_logo_32px.png\" /\u003eDownload notebook\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\u003c/table\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8f-reQ11gbLB"
      },
      "source": [
        "In this tutorial, we build a simple two tower ranking model using the [MovieLens 100K dataset](https://grouplens.org/datasets/movielens/100k/) with TF-Ranking. We can use this model to rank and recommend movies for a given user according to their predicted user ratings."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qA00wBE2Ntdm"
      },
      "source": [
        "## Setup\n",
        "\n",
        "Install and import the TF-Ranking library:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6yzAaM85Z12D"
      },
      "outputs": [],
      "source": [
        "!pip install -q tensorflow-ranking\n",
        "!pip install -q --upgrade tensorflow-datasets"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "n3oYt3R6Nr9l"
      },
      "outputs": [],
      "source": [
        "from typing import Dict, Tuple\n",
        "\n",
        "import tensorflow as tf\n",
        "\n",
        "import tensorflow_datasets as tfds\n",
        "import tensorflow_ranking as tfr"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zCxQ1CZcO2wh"
      },
      "source": [
        "## Read the data"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A0sY6-Rtt_Co"
      },
      "source": [
        "Prepare to train a model by creating a ratings dataset and movies dataset. Use `user_id` as the query input feature, `movie_title` as the document input feature, and `user_rating` as the label to train the ranking model."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "M-mxBYjdO5m7"
      },
      "outputs": [],
      "source": [
        "%%capture --no-display\n",
        "# Ratings data.\n",
        "ratings = tfds.load('movielens/100k-ratings', split=\"train\")\n",
        "# Features of all the available movies.\n",
        "movies = tfds.load('movielens/100k-movies', split=\"train\")\n",
        "\n",
        "# Select the basic features.\n",
        "ratings = ratings.map(lambda x: {\n",
        "    \"movie_title\": x[\"movie_title\"],\n",
        "    \"user_id\": x[\"user_id\"],\n",
        "    \"user_rating\": x[\"user_rating\"]\n",
        "})"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5W0HSfmSNCWm"
      },
      "source": [
        "Build vocabularies to convert all user ids and all movie titles into integer indices for embedding layers:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9I1VTEjHzpfX"
      },
      "outputs": [],
      "source": [
        "movies = movies.map(lambda x: x[\"movie_title\"])\n",
        "users = ratings.map(lambda x: x[\"user_id\"])\n",
        "\n",
        "user_ids_vocabulary = tf.keras.layers.experimental.preprocessing.StringLookup(\n",
        "    mask_token=None)\n",
        "user_ids_vocabulary.adapt(users.batch(1000))\n",
        "\n",
        "movie_titles_vocabulary = tf.keras.layers.experimental.preprocessing.StringLookup(\n",
        "    mask_token=None)\n",
        "movie_titles_vocabulary.adapt(movies.batch(1000))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zMsmoqWTOTKo"
      },
      "source": [
        "Group by `user_id` to form lists for ranking models:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lXY7kX7nOSwH"
      },
      "outputs": [],
      "source": [
        "key_func = lambda x: user_ids_vocabulary(x[\"user_id\"])\n",
        "reduce_func = lambda key, dataset: dataset.batch(100)\n",
        "ds_train = ratings.group_by_window(\n",
        "    key_func=key_func, reduce_func=reduce_func, window_size=100)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "57r87tdQlkcT"
      },
      "outputs": [],
      "source": [
        "for x in ds_train.take(1):\n",
        "  for key, value in x.items():\n",
        "    print(f\"Shape of {key}: {value.shape}\")\n",
        "    print(f\"Example values of {key}: {value[:5].numpy()}\")\n",
        "    print()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YcZJf2qxOeWU"
      },
      "source": [
        "Generate batched features and labels:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ctq2RTOqOfAo"
      },
      "outputs": [],
      "source": [
        "def _features_and_labels(\n",
        "    x: Dict[str, tf.Tensor]) -\u003e Tuple[Dict[str, tf.Tensor], tf.Tensor]:\n",
        "  labels = x.pop(\"user_rating\")\n",
        "  return x, labels\n",
        "\n",
        "\n",
        "ds_train = ds_train.map(_features_and_labels)\n",
        "\n",
        "ds_train = ds_train.apply(\n",
        "    tf.data.experimental.dense_to_ragged_batch(batch_size=32))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RJUU3mv-_VdQ"
      },
      "source": [
        "The `user_id` and `movie_title` tensors generated in `ds_train` are of shape `[32, None]`, where the second dimension is 100 in most cases except for the batches when less than 100 items grouped in lists. A model working on ragged tensors is thus used."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GTquqk1GkIfd"
      },
      "outputs": [],
      "source": [
        "for x, label in ds_train.take(1):\n",
        "  for key, value in x.items():\n",
        "    print(f\"Shape of {key}: {value.shape}\")\n",
        "    print(f\"Example values of {key}: {value[:3, :3].numpy()}\")\n",
        "    print()\n",
        "  print(f\"Shape of label: {label.shape}\")\n",
        "  print(f\"Example values of label: {label[:3, :3].numpy()}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Lrch6rVBOB9Q"
      },
      "source": [
        "## Define a model\n",
        "\n",
        "Define a ranking model by inheriting from `tf.keras.Model` and implementing the `call` method:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "e5dNbDZwOIHR"
      },
      "outputs": [],
      "source": [
        "class MovieLensRankingModel(tf.keras.Model):\n",
        "\n",
        "  def __init__(self, user_vocab, movie_vocab):\n",
        "    super().__init__()\n",
        "\n",
        "    # Set up user and movie vocabulary and embedding.\n",
        "    self.user_vocab = user_vocab\n",
        "    self.movie_vocab = movie_vocab\n",
        "    self.user_embed = tf.keras.layers.Embedding(user_vocab.vocabulary_size(),\n",
        "                                                64)\n",
        "    self.movie_embed = tf.keras.layers.Embedding(movie_vocab.vocabulary_size(),\n",
        "                                                 64)\n",
        "\n",
        "  def call(self, features: Dict[str, tf.Tensor]) -\u003e tf.Tensor:\n",
        "    # Define how the ranking scores are computed: \n",
        "    # Take the dot-product of the user embeddings with the movie embeddings.\n",
        "\n",
        "    user_embeddings = self.user_embed(self.user_vocab(features[\"user_id\"]))\n",
        "    movie_embeddings = self.movie_embed(\n",
        "        self.movie_vocab(features[\"movie_title\"]))\n",
        "\n",
        "    return tf.reduce_sum(user_embeddings * movie_embeddings, axis=2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BMV0HpzmJGWk"
      },
      "source": [
        "Create the model, and then compile it with ranking `tfr.keras.losses` and `tfr.keras.metrics`, which are the core of the TF-Ranking package. \n",
        "\n",
        "This example uses a ranking-specific **softmax loss**, which is a listwise loss introduced to promote all relevant items in the ranking list with better chances on top of the irrelevant ones. In contrast to the softmax loss in the multi-class classification problem, where only one class is positive and the rest are negative, the TF-Ranking library supports multiple relevant documents in a query list and non-binary relevance labels.\n",
        "\n",
        "For ranking metrics, this example uses in specific **Normalized Discounted Cumulative Gain (NDCG)** and **Mean Reciprocal Rank (MRR)**, which calculate the user utility of a ranked query list with position discounts. For more details about ranking metrics, review evaluation measures [offline metrics](https://en.wikipedia.org/wiki/Evaluation_measures_(information_retrieval)#Offline_metrics)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "H2tQDhqkOKf1"
      },
      "outputs": [],
      "source": [
        "# Create the ranking model, trained with a ranking loss and evaluated with\n",
        "# ranking metrics.\n",
        "model = MovieLensRankingModel(user_ids_vocabulary, movie_titles_vocabulary)\n",
        "optimizer = tf.keras.optimizers.Adagrad(0.5)\n",
        "loss = tfr.keras.losses.get(\n",
        "    loss=tfr.keras.losses.RankingLossKey.SOFTMAX_LOSS, ragged=True)\n",
        "eval_metrics = [\n",
        "    tfr.keras.metrics.get(key=\"ndcg\", name=\"metric/ndcg\", ragged=True),\n",
        "    tfr.keras.metrics.get(key=\"mrr\", name=\"metric/mrr\", ragged=True)\n",
        "]\n",
        "model.compile(optimizer=optimizer, loss=loss, metrics=eval_metrics)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NeBnBFMfVLzP"
      },
      "source": [
        "## Train and evaluate the model\n",
        "\n",
        "Train the model with `model.fit`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bzGm7WqSVNyP"
      },
      "outputs": [],
      "source": [
        "model.fit(ds_train, epochs=3)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "V5uuSRXZoOKW"
      },
      "source": [
        "Generate predictions and evaluate."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6Hryvj3cPnvK"
      },
      "outputs": [],
      "source": [
        "# Get movie title candidate list.\n",
        "for movie_titles in movies.batch(2000):\n",
        "  break\n",
        "\n",
        "# Generate the input for user 42.\n",
        "inputs = {\n",
        "    \"user_id\":\n",
        "        tf.expand_dims(tf.repeat(\"42\", repeats=movie_titles.shape[0]), axis=0),\n",
        "    \"movie_title\":\n",
        "        tf.expand_dims(movie_titles, axis=0)\n",
        "}\n",
        "\n",
        "# Get movie recommendations for user 42.\n",
        "scores = model(inputs)\n",
        "titles = tfr.utils.sort_by_scores(scores,\n",
        "                                  [tf.expand_dims(movie_titles, axis=0)])[0]\n",
        "print(f\"Top 5 recommendations for user 42: {titles[0, :5]}\")"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "quickstart.ipynb",
      "private_outputs": true,
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}