tfrs_retrieval.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import tensorflow as tf\n",
    "import tensorflow_datasets as tfds\n",
    "import tensorflow_recommenders as tfrs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2 stages:\n",
    "- Retrieval model - selects an initial set of candidates from all possible candidates\n",
    "- Ranking model - ranks the selected candidates and selects the best ones"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [],
   "source": [
    "def show(dataset, n=2):\n",
    "    return pd.DataFrame(dataset.take(n).as_numpy_iterator())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Retrieval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [],
   "source": [
    "ratings = tfds.load(\"movielens/100k-ratings\", split=\"train\")\n",
    "movies = tfds.load(\"movielens/100k-movies\", split=\"train\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "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>bucketized_user_age</th>\n",
       "      <th>movie_genres</th>\n",
       "      <th>movie_id</th>\n",
       "      <th>movie_title</th>\n",
       "      <th>raw_user_age</th>\n",
       "      <th>timestamp</th>\n",
       "      <th>user_gender</th>\n",
       "      <th>user_id</th>\n",
       "      <th>user_occupation_label</th>\n",
       "      <th>user_occupation_text</th>\n",
       "      <th>user_rating</th>\n",
       "      <th>user_zip_code</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>45.0</td>\n",
       "      <td>[7]</td>\n",
       "      <td>b'357'</td>\n",
       "      <td>b\"One Flew Over the Cuckoo's Nest (1975)\"</td>\n",
       "      <td>46.0</td>\n",
       "      <td>879024327</td>\n",
       "      <td>True</td>\n",
       "      <td>b'138'</td>\n",
       "      <td>4</td>\n",
       "      <td>b'doctor'</td>\n",
       "      <td>4.0</td>\n",
       "      <td>b'53211'</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>25.0</td>\n",
       "      <td>[4, 14]</td>\n",
       "      <td>b'709'</td>\n",
       "      <td>b'Strictly Ballroom (1992)'</td>\n",
       "      <td>32.0</td>\n",
       "      <td>875654590</td>\n",
       "      <td>True</td>\n",
       "      <td>b'92'</td>\n",
       "      <td>5</td>\n",
       "      <td>b'entertainment'</td>\n",
       "      <td>2.0</td>\n",
       "      <td>b'80525'</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   bucketized_user_age movie_genres movie_id  \\\n",
       "0                 45.0          [7]   b'357'   \n",
       "1                 25.0      [4, 14]   b'709'   \n",
       "\n",
       "                                 movie_title  raw_user_age  timestamp  \\\n",
       "0  b\"One Flew Over the Cuckoo's Nest (1975)\"          46.0  879024327   \n",
       "1                b'Strictly Ballroom (1992)'          32.0  875654590   \n",
       "\n",
       "   user_gender user_id  user_occupation_label user_occupation_text  \\\n",
       "0         True  b'138'                      4            b'doctor'   \n",
       "1         True   b'92'                      5     b'entertainment'   \n",
       "\n",
       "   user_rating user_zip_code  \n",
       "0          4.0      b'53211'  \n",
       "1          2.0      b'80525'  "
      ]
     },
     "execution_count": 62,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# user_rating -> explicit feedback\n",
    "# first stage of the recommender (retrieval model) will use implicit feedback\n",
    "# (converted from the explicit user ratings)\n",
    "# every movie a user rated (watched) -> positive example\n",
    "# every movie a user have not rated -> negative example\n",
    "show(ratings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "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>movie_genres</th>\n",
       "      <th>movie_id</th>\n",
       "      <th>movie_title</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>[4]</td>\n",
       "      <td>b'1681'</td>\n",
       "      <td>b'You So Crazy (1994)'</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>[4, 7]</td>\n",
       "      <td>b'1457'</td>\n",
       "      <td>b'Love Is All There Is (1996)'</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "  movie_genres movie_id                     movie_title\n",
       "0          [4]  b'1681'          b'You So Crazy (1994)'\n",
       "1       [4, 7]  b'1457'  b'Love Is All There Is (1996)'"
      ]
     },
     "execution_count": 63,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# move_genres -> list of genre ids the movie belngs to\n",
    "show(movies)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [],
   "source": [
    "# no explicit rating\n",
    "seen_by_users = ratings.map(\n",
    "    lambda x: {'movie_title': x['movie_title'], 'user_id': x['user_id']})\n",
    "\n",
    "movie_titles = movies.map(lambda x: x['movie_title'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_samples = 100_000\n",
    "train_split = 0.8\n",
    "n_train_samples = round(train_split * n_samples)\n",
    "n_test_samples = n_samples = n_train_samples\n",
    "\n",
    "tf.random.set_seed(0)\n",
    "seen_by_users = seen_by_users.shuffle(\n",
    "    100_000, seed=0, reshuffle_each_iteration=False\n",
    ")\n",
    "\n",
    "train = seen_by_users.take(n_train_samples)\n",
    "test = seen_by_users.skip(n_train_samples).take(n_test_samples)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [],
   "source": [
    "unique_movie_titles = np.unique(list(movie_titles.as_numpy_iterator()))\n",
    "unique_user_ids = np.unique(\n",
    "    list(seen_by_users.map(lambda x: x['user_id']).as_numpy_iterator())\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[1, 2, 2],\n",
       "       [1, 2, 3]])"
      ]
     },
     "execution_count": 67,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# string lookup example\n",
    "vocab = ['a', 'b', 'c']\n",
    "# data = tf.constant([[\"a\", \"c\", \"d\"], [\"d\", \"z\", \"b\"]])\n",
    "test_data = np.array([['a', 'b', 'b'], ['a', 'b', 'c']])\n",
    "layer = tf.keras.layers.experimental.preprocessing.StringLookup(vocabulary=vocab)\n",
    "layer(test_data).numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 3, 2), dtype=float32, numpy=\n",
       "array([[[-0.02080249, -0.02934336],\n",
       "        [ 0.00353907,  0.00612575],\n",
       "        [-0.00833255,  0.03078279]],\n",
       "\n",
       "       [[-0.02080249, -0.02934336],\n",
       "        [-0.02080249, -0.02934336],\n",
       "        [-0.00833255,  0.03078279]]], dtype=float32)>"
      ]
     },
     "execution_count": 68,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# embedding layer example\n",
    "# embedding dim=2: 2x3 input -> 2x3x2 output\n",
    "layer_2 = tf.keras.layers.Embedding(len(vocab), 2, input_length=3)\n",
    "i = np.array([[0, 1, 2], [0, 0, 2]])\n",
    "layer_2(i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "# query model\n",
    "# https://stats.stackexchange.com/questions/270546/how-does-keras-embedding-layer-work\n",
    "def embedding_model_factory(n_embedding_dimensions=32, vocabulary):\n",
    "    layers = [\n",
    "        # Maps strings from a vocabulary to a range of integer indices\n",
    "        tf.keras.layers.experimental.preprocessing.StringLookup(\n",
    "            vocabulary=vocabulary, mask_token=None\n",
    "        ),\n",
    "        # + 1 to add one aditional embedding entry for unknown tokens\n",
    "        tf.keras.layers.Embedding(\n",
    "            len(vocabulary) + 1, n_embedding_dimensions\n",
    "        )\n",
    "    ]\n",
    "    \n",
    "    return tf.keras.Sequential(layers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(1, 32), dtype=float32, numpy=\n",
       "array([[-2.4332905e-02,  1.5845884e-02,  2.8408874e-02,  4.8857283e-02,\n",
       "        -2.8385568e-02, -4.8242688e-02, -8.9645386e-05, -4.0224455e-02,\n",
       "         4.6875920e-02,  3.5669688e-02,  3.5790887e-02, -2.7830116e-03,\n",
       "         2.4314746e-03,  1.8764388e-02,  8.4960945e-03,  2.2447530e-02,\n",
       "        -4.4737469e-02, -2.6634825e-02,  4.8829857e-02, -3.5618961e-02,\n",
       "         3.8881149e-02,  2.6854027e-02, -2.4254238e-02, -4.8504997e-02,\n",
       "         2.9529762e-02, -6.4371899e-04,  2.0598780e-02, -2.7926341e-03,\n",
       "         2.4232868e-02, -1.7427541e-02, -6.1231032e-03, -2.7959382e-02]],\n",
       "      dtype=float32)>"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "movie_model = embedding_model_factory(\n",
    "    n_embedding_dimensions=32,\n",
    "    vocabulary=unique_movie_titles\n",
    ")\n",
    "# embedding for the movie \"Til There Was You (1997)\"\n",
    "movie_model(np.array([b\"'Til There Was You (1997)\"]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(1, 32), dtype=float32, numpy=\n",
       "array([[-0.036576  , -0.03089347, -0.04926305, -0.00108545, -0.02889257,\n",
       "        -0.02835456, -0.02898588,  0.03876593,  0.01880166,  0.01347914,\n",
       "         0.03804502, -0.04827794, -0.03280818,  0.00111272, -0.04606495,\n",
       "        -0.03216724,  0.04409171,  0.01784435,  0.04915215, -0.0402068 ,\n",
       "        -0.04015896,  0.02858837, -0.01993331,  0.02911016, -0.01273555,\n",
       "        -0.02697011, -0.0156634 ,  0.00053151,  0.0050525 ,  0.04768539,\n",
       "         0.02215948,  0.00570469]], dtype=float32)>"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "user_model = embedding_model_factory(\n",
    "    n_embedding_dimensions=32,\n",
    "    vocabulary=unique_user_ids\n",
    ")\n",
    "# embedding for the user \"10\"\n",
    "user_model(np.array([b'10']))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(), dtype=float32, numpy=0.0>"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# affinity score := dot product of user and movie embedding vectors\n",
    "# for an accurate model, the score for positive pairs (pairs of movies/users in the dataset)\n",
    "# should be higher than for any other user/movie pairs\n",
    "\n",
    "metrics = tfrs.metrics.FactorizedTopK(\n",
    "  # candidates is an nx32 embedding matrix containing all the movies\n",
    "  # used as implicit negative pairs?\n",
    "  candidates=movie_titles.batch(128).map(movie_model)\n",
    ")\n",
    "\n",
    "# setup the task for the retrival model with appropriate loss\n",
    "# default is categorical cross entropy\n",
    "task = tfrs.tasks.Retrieval(metrics=metrics)\n",
    "\n",
    "u = user_model(np.array(['b102']))\n",
    "v = movie_model(np.array([b'1-900 (1994)']))\n",
    "# task is a layer that takes query and candidate vectors and returns loss\n",
    "task(u, v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FullModel(tfrs.Model):\n",
    "    \n",
    "    def __init__(self, user_model, movie_model, task):\n",
    "        super().__init__()\n",
    "        self.user_model = user_model\n",
    "        self.movie_model = movie_model\n",
    "        self.task = task\n",
    "        \n",
    "    def compute_loss(self, features, training=False):\n",
    "        # -> features: Dict[Text, tf.Tensor]\n",
    "        # -> tf.Tensor\n",
    "\n",
    "        # get embeddings for given user\n",
    "        user_embeddings = self.user_model(features['user_id'])\n",
    "        # get embedding for given movie (that the user has seen) -> positive signal\n",
    "        positive_movie_embeddings = self.movie_model(features['movie_title'])\n",
    "        # return loss a metrics\n",
    "        return self.task(user_embeddings, positive_movie_embeddings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = FullModel(user_model, movie_model, task)\n",
    "model.compile(optimizer=tf.keras.optimizers.Adagrad(learning_rate=0.1))\n",
    "\n",
    "cached_train = train.shuffle(100_000).batch(8192).cache()\n",
    "cached_test = test.batch(4096).cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/3\n",
      "10/10 [==============================] - 11s 932ms/step - factorized_top_k/top_1_categorical_accuracy: 0.0014 - factorized_top_k/top_5_categorical_accuracy: 0.0101 - factorized_top_k/top_10_categorical_accuracy: 0.0221 - factorized_top_k/top_50_categorical_accuracy: 0.1027 - factorized_top_k/top_100_categorical_accuracy: 0.1801 - loss: 69825.5795 - regularization_loss: 0.0000e+00 - total_loss: 69825.5795\n",
      "Epoch 2/3\n",
      "10/10 [==============================] - 9s 916ms/step - factorized_top_k/top_1_categorical_accuracy: 0.0030 - factorized_top_k/top_5_categorical_accuracy: 0.0195 - factorized_top_k/top_10_categorical_accuracy: 0.0400 - factorized_top_k/top_50_categorical_accuracy: 0.1688 - factorized_top_k/top_100_categorical_accuracy: 0.2894 - loss: 67515.8928 - regularization_loss: 0.0000e+00 - total_loss: 67515.8928\n",
      "Epoch 3/3\n",
      "10/10 [==============================] - 9s 936ms/step - factorized_top_k/top_1_categorical_accuracy: 0.0037 - factorized_top_k/top_5_categorical_accuracy: 0.0238 - factorized_top_k/top_10_categorical_accuracy: 0.0475 - factorized_top_k/top_50_categorical_accuracy: 0.1909 - factorized_top_k/top_100_categorical_accuracy: 0.3193 - loss: 66289.9339 - regularization_loss: 0.0000e+00 - total_loss: 66289.9339\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<keras.callbacks.History at 0x7f7e6442fe10>"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.fit(cached_train, epochs=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5/5 [==============================] - 3s 360ms/step - factorized_top_k/top_1_categorical_accuracy: 0.0013 - factorized_top_k/top_5_categorical_accuracy: 0.0093 - factorized_top_k/top_10_categorical_accuracy: 0.0206 - factorized_top_k/top_50_categorical_accuracy: 0.1223 - factorized_top_k/top_100_categorical_accuracy: 0.2326 - loss: 31095.9199 - regularization_loss: 0.0000e+00 - total_loss: 31095.9199\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'factorized_top_k/top_1_categorical_accuracy': 0.0013000000035390258,\n",
       " 'factorized_top_k/top_5_categorical_accuracy': 0.009349999949336052,\n",
       " 'factorized_top_k/top_10_categorical_accuracy': 0.020649999380111694,\n",
       " 'factorized_top_k/top_50_categorical_accuracy': 0.12229999899864197,\n",
       " 'factorized_top_k/top_100_categorical_accuracy': 0.2326499968767166,\n",
       " 'loss': 28265.875,\n",
       " 'regularization_loss': 0,\n",
       " 'total_loss': 28265.875}"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.evaluate(cached_test, return_dict=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tensorflow_recommenders.layers.factorized_top_k.BruteForce at 0x7f7e5c6d4d50>"
      ]
     },
     "execution_count": 74,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# making predictions\n",
    "\n",
    "# create a model that takes in raw query features, and\n",
    "index = tfrs.layers.factorized_top_k.BruteForce(model.user_model)\n",
    "# recommends movies out of the entire movies dataset\n",
    "dataset = tf.data.Dataset.zip(\n",
    "      (movie_titles.batch(100), movie_titles.batch(100).map(model.movie_model))\n",
    ")\n",
    "index.index_from_dataset(dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[[b'Angels in the Outfield (1994)' b'Rudy (1993)'\n",
      "  b'Homeward Bound: The Incredible Journey (1993)'\n",
      "  b'Affair to Remember, An (1957)' b'Circle of Friends (1995)'\n",
      "  b'Client, The (1994)' b'Firm, The (1993)'\n",
      "  b'Miracle on 34th Street (1994)'\n",
      "  b'Winnie the Pooh and the Blustery Day (1968)' b'Michael (1996)']], shape=(1, 10), dtype=string)\n"
     ]
    }
   ],
   "source": [
    "# get recommendations\n",
    "_, result = index(tf.constant([\"42\"]))\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tf.Tensor(\n",
      "[[b'Juror, The (1996)' b'Down Periscope (1996)' b'Phantom, The (1996)'\n",
      "  b'Eddie (1996)' b'Nutty Professor, The (1996)' b'Happy Gilmore (1996)'\n",
      "  b'Bulletproof (1996)' b'Father of the Bride Part II (1995)'\n",
      "  b'Craft, The (1996)' b'Space Jam (1996)']], shape=(1, 10), dtype=string)\n"
     ]
    }
   ],
   "source": [
    "# using ScaNN index\n",
    "scann_index = tfrs.layers.factorized_top_k.ScaNN(model.user_model)\n",
    "scann_index.index_from_dataset(dataset)\n",
    "\n",
    "_, scann_result = index(tf.constant([\"101\"]))\n",
    "print(scann_result)\n",
    "\n",
    "# to persist the model, just save the index object"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "airbus",
   "language": "python",
   "name": "airbus"
  },
  "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.7.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}