2_retrieval_model_training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9f53a397",
   "metadata": {},
   "source": [
    "## <span style=\"color:#ff5f27\">🧬 Train Retrieval Model </span>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9856c6ab-078c-4dc8-a989-353f8926206d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install hopsworks==3.7.0 --quiet\n",
    "# !pip install tensorflow==2.15 --quiet\n",
    "# !pip install tensorflow-recommenders==0.7.3 --quiet"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9faa182c",
   "metadata": {},
   "source": [
    "## <span style=\"color:#ff5f27\">📝 Imports </span>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91d80acd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "from tensorflow.keras.layers import StringLookup, Normalization\n",
    "import tensorflow_recommenders as tfrs\n",
    "import pandas as pd\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a09283b4",
   "metadata": {},
   "source": [
    "## <span style=\"color:#ff5f27\">🔮 Connect to Hopsworks Feature Store </span>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03b4ad6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import hopsworks\n",
    "\n",
    "project = hopsworks.login()\n",
    "\n",
    "fs = project.get_feature_store()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3261b352",
   "metadata": {},
   "source": [
    "## <span style=\"color:#ff5f27\">🔪 Feature Selection </span>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c149d5c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "users_fg = fs.get_feature_group(\n",
    "    name=\"users\",\n",
    "    version=1,\n",
    ")\n",
    "\n",
    "videos_fg = fs.get_feature_group(\n",
    "    name=\"videos\",\n",
    "    version=1,\n",
    ")\n",
    "\n",
    "interactions_fg = fs.get_feature_group(\n",
    "    name=\"interactions\",\n",
    "    version=1,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd2c14e9",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94a22d99",
   "metadata": {},
   "source": [
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d72398e",
   "metadata": {},
   "outputs": [],
   "source": [
    "QUERY_FEATURES = [\"user_id\", \"gender\", \"age\", \"country\"]\n",
    "CANDIDATE_FEATURES = [\"video_id\", \"category\", \"views\", \"likes\", \"video_length\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f091291d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Select features for training data\n",
    "selected_features = interactions_fg.select([\"interaction_id\"])\\\n",
    "    .join(users_fg.select(QUERY_FEATURES), on=\"user_id\")\\\n",
    "    .join(videos_fg.select(CANDIDATE_FEATURES), on=\"video_id\")\n",
    "\n",
    "# Uncomment this if you would like to view your selected features\n",
    "# selected_features.show(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87707f44",
   "metadata": {},
   "source": [
    "## <span style=\"color:#ff5f27\">⚙️ Feature View Creation </span>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af97a264",
   "metadata": {},
   "outputs": [],
   "source": [
    "feature_view = fs.get_or_create_feature_view(\n",
    "    name='retrieval',\n",
    "    version=1,\n",
    "    query=selected_features,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bad1dc1f",
   "metadata": {},
   "source": [
    "## <span style=\"color:#ff5f27\">🏋️ Training Dataset </span>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "437f3034",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df, val_df, test_df, _, _, _ = feature_view.train_validation_test_split(\n",
    "    validation_size=0.1, \n",
    "    test_size=0.01,\n",
    "    description='Retrieval dataset splits',\n",
    ")\n",
    "train_df.head(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66935588",
   "metadata": {},
   "source": [
    "You will train your retrieval model with a subset of features.\n",
    "\n",
    "For the query embedding you will use:\n",
    "- `user_id`: ID of a user.\n",
    "- `gender`: Gender of a user.\n",
    "- `age`: age of a user.\n",
    "- `country`: country if a user.\n",
    "\n",
    "For the candidate embedding you will use:\n",
    "- `video_id`: ID of a video.\n",
    "- `category`: Video Category.\n",
    "- `views`: Number of video views.\n",
    "- `likes`: Number of video likes.\n",
    "- `video_length`: Length of video.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f95d63e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def df_to_ds(df):\n",
    "    return tf.data.Dataset.from_tensor_slices({col: df[col] for col in df})\n",
    "\n",
    "BATCH_SIZE = 2048\n",
    "train_ds = df_to_ds(train_df).batch(BATCH_SIZE).cache().shuffle(BATCH_SIZE*10)\n",
    "val_ds = df_to_ds(val_df).batch(BATCH_SIZE).cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0bb3dd6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Query Features \n",
    "user_id_list = train_df[\"user_id\"].unique().tolist()\n",
    "countries_list = train_df[\"country\"].unique().tolist()\n",
    "gender_list = train_df[\"gender\"].unique().tolist()\n",
    "\n",
    "# Item Features\n",
    "video_id_list = train_df[\"video_id\"].unique().tolist()\n",
    "category_list = train_df[\"category\"].unique().tolist()\n",
    "\n",
    "print(f\"⛳️ Number of users: {len(user_id_list)}\")\n",
    "print(f\"⛳️ Number of items: {len(video_id_list)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "434c4058",
   "metadata": {},
   "source": [
    "## <span style=\"color:#ff5f27\">🏰 Two Tower Model </span>\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8db71a34",
   "metadata": {},
   "outputs": [],
   "source": [
    "EMB_DIM = 16"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63b6462f",
   "metadata": {},
   "outputs": [],
   "source": [
    "class QueryTower(tf.keras.Model):\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "        self.emb_dim = EMB_DIM\n",
    "        self.user_embedding = tf.keras.Sequential([\n",
    "            StringLookup(\n",
    "                vocabulary=user_id_list,\n",
    "                mask_token=None\n",
    "            ),\n",
    "            tf.keras.layers.Embedding(\n",
    "                # You add an additional embedding to account for unknown tokens.\n",
    "                len(user_id_list) + 1,\n",
    "                self.emb_dim\n",
    "            )\n",
    "        ])\n",
    "\n",
    "        self.normalized_age = Normalization(axis=None)\n",
    "        \n",
    "        # Converts strings into integer indices (scikit-learn LabelEncoder analog)\n",
    "        self.gender_tokenizer = StringLookup(\n",
    "            vocabulary=gender_list,\n",
    "            mask_token=None,\n",
    "        )\n",
    "        \n",
    "        self.country_tokenizer = StringLookup(\n",
    "            vocabulary=countries_list, \n",
    "            mask_token=None,\n",
    "        )\n",
    "\n",
    "        self.fnn = tf.keras.Sequential([\n",
    "            tf.keras.layers.Dense(self.emb_dim, activation=\"relu\"),\n",
    "            tf.keras.layers.Dense(self.emb_dim)\n",
    "        ])\n",
    "\n",
    "    def call(self, inputs):\n",
    "        gender_embedding = tf.one_hot(\n",
    "            self.gender_tokenizer(inputs[\"gender\"]),\n",
    "            len(gender_list),\n",
    "        )\n",
    "        \n",
    "        country_embedding = tf.one_hot(\n",
    "            self.country_tokenizer(inputs[\"country\"]),\n",
    "            len(countries_list),\n",
    "        )\n",
    "        \n",
    "        concatenated_inputs = tf.concat([\n",
    "            self.user_embedding(inputs[\"user_id\"]),\n",
    "            tf.reshape(self.normalized_age(inputs[\"age\"]), (-1,1)),\n",
    "            gender_embedding,\n",
    "            country_embedding,\n",
    "        ], axis=1)\n",
    "\n",
    "        outputs = self.fnn(concatenated_inputs)\n",
    "\n",
    "        return outputs\n",
    "\n",
    "\n",
    "query_model = QueryTower()\n",
    "\n",
    "query_model.normalized_age.adapt(train_ds.map(lambda x : x[\"age\"]))\n",
    "\n",
    "# Initialize model with inputs.\n",
    "query_df = train_df[QUERY_FEATURES]\n",
    "query_ds = df_to_ds(query_df).batch(1)\n",
    "query_model(next(iter(query_ds)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04965563",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ItemTower(tf.keras.Model):\n",
    "\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "        self.emb_dim = EMB_DIM\n",
    "        self.video_embedding = tf.keras.Sequential([\n",
    "            StringLookup(\n",
    "                vocabulary=video_id_list,\n",
    "                mask_token=None\n",
    "            ),\n",
    "            tf.keras.layers.Embedding(\n",
    "                # You add an additional embedding to account for unknown tokens.\n",
    "                len(video_id_list) + 1,\n",
    "                self.emb_dim,\n",
    "            )\n",
    "        ])\n",
    "        \n",
    "        # Converts strings into integer indices (scikit-learn LabelEncoder analog)\n",
    "        self.category_tokenizer = StringLookup(\n",
    "            vocabulary=category_list, \n",
    "            mask_token=None,\n",
    "        )\n",
    "        \n",
    "        self.normalized_views = Normalization(axis=None)\n",
    "        self.normalized_likes = Normalization(axis=None)\n",
    "        self.normalized_video_length = Normalization(axis=None)\n",
    "\n",
    "        self.fnn = tf.keras.Sequential([\n",
    "            tf.keras.layers.Dense(self.emb_dim, activation=\"relu\"),\n",
    "            tf.keras.layers.Dense(self.emb_dim)\n",
    "        ])\n",
    "\n",
    "    def call(self, inputs):\n",
    "        category_embedding = tf.one_hot(\n",
    "            self.category_tokenizer(inputs[\"category\"]),\n",
    "            len(category_list),\n",
    "        )\n",
    "\n",
    "        concatenated_inputs = tf.concat([\n",
    "            self.video_embedding(inputs[\"video_id\"]),\n",
    "            category_embedding,\n",
    "            tf.reshape(self.normalized_views(inputs[\"views\"]), (-1,1)),\n",
    "            tf.reshape(self.normalized_views(inputs[\"likes\"]), (-1,1)),\n",
    "            tf.reshape(self.normalized_views(inputs[\"video_length\"]), (-1,1)),\n",
    "        ], axis=1)\n",
    "\n",
    "        outputs = self.fnn(concatenated_inputs)\n",
    "\n",
    "        return outputs\n",
    "\n",
    "    \n",
    "item_model = ItemTower()\n",
    "\n",
    "item_df = train_df[CANDIDATE_FEATURES]\n",
    "item_df.drop_duplicates(subset=\"video_id\", inplace=True)\n",
    "item_ds = df_to_ds(item_df)\n",
    "\n",
    "item_model(next(iter(item_ds.batch(1))))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "440193f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TwoTowerModel(tf.keras.Model):\n",
    "    def __init__(self, query_model, item_model):\n",
    "        super().__init__()\n",
    "        self.query_model = query_model\n",
    "        self.item_model = item_model\n",
    "        self.task = tfrs.tasks.Retrieval(\n",
    "            metrics=tfrs.metrics.FactorizedTopK(\n",
    "                candidates=item_ds.batch(BATCH_SIZE).map(self.item_model)\n",
    "            )\n",
    "        )\n",
    "\n",
    "    def train_step(self, batch) -> tf.Tensor:\n",
    "        # Set up a gradient tape to record gradients.\n",
    "        with tf.GradientTape() as tape:\n",
    "\n",
    "            # Loss computation.\n",
    "            user_embeddings = self.query_model(batch)\n",
    "            item_embeddings = self.item_model(batch)\n",
    "            loss = self.task(\n",
    "                user_embeddings, \n",
    "                item_embeddings,\n",
    "                compute_metrics=False,\n",
    "            )\n",
    "\n",
    "            # Handle regularization losses as well.\n",
    "            regularization_loss = sum(self.losses)\n",
    "\n",
    "            total_loss = loss + regularization_loss\n",
    "\n",
    "        gradients = tape.gradient(total_loss, self.trainable_variables)\n",
    "        self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))\n",
    "\n",
    "        metrics = {\n",
    "            \"loss\": loss,\n",
    "            \"regularization_loss\": regularization_loss,\n",
    "            \"total_loss\": total_loss\n",
    "        }\n",
    "\n",
    "        return metrics\n",
    "\n",
    "    def test_step(self, batch) -> tf.Tensor:\n",
    "        # Loss computation.\n",
    "        user_embeddings = self.query_model(batch)\n",
    "        item_embeddings = self.item_model(batch)\n",
    "\n",
    "        loss = self.task(\n",
    "            user_embeddings, \n",
    "            item_embeddings,\n",
    "            compute_metrics=False,\n",
    "        )\n",
    "\n",
    "        # Handle regularization losses as well.\n",
    "        regularization_loss = sum(self.losses)\n",
    "\n",
    "        total_loss = loss + regularization_loss\n",
    "\n",
    "        metrics = {metric.name: metric.result() for metric in self.metrics}\n",
    "        metrics[\"loss\"] = loss\n",
    "        metrics[\"regularization_loss\"] = regularization_loss\n",
    "        metrics[\"total_loss\"] = total_loss\n",
    "\n",
    "        return metrics"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5bf66678",
   "metadata": {},
   "source": [
    "### <span style=\"color:#ff5f27\">🏃🏻‍♂️ Model Training </span>\n",
    "\n",
    "You'll train our model using the AdamW optimizer, which applies weight regularization during training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e96e26c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a TwoTowerModel with the specified query_model and item_model\n",
    "model = TwoTowerModel(query_model, item_model)\n",
    "\n",
    "# Define an optimizer using AdamW with a learning rate of 0.01\n",
    "optimizer = tf.keras.optimizers.AdamW(\n",
    "    weight_decay=0.001, \n",
    "    learning_rate=0.01,\n",
    ")\n",
    "\n",
    "# Compile the model using the specified optimizer\n",
    "model.compile(optimizer=optimizer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46ea848e",
   "metadata": {},
   "outputs": [],
   "source": [
    "model.fit(\n",
    "    train_ds, \n",
    "    validation_data=val_ds, \n",
    "    epochs=5,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6eb20b1e",
   "metadata": {},
   "source": [
    "## <span style=\"color:#ff5f27\">🗄️ Upload Model to Model Registry </span>\n",
    "\n",
    "One of the features in Hopsworks is the model registry. This is where you can store different versions of models and compare their performance. Models from the registry can then be served as API endpoints.\n",
    "\n",
    "Let's connect to the model registry using the [HSML library](https://docs.hopsworks.ai/machine-learning-api/latest) from Hopsworks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aac5bd7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "mr = project.get_model_registry()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e3fe643",
   "metadata": {},
   "outputs": [],
   "source": [
    "class QueryModelModule(tf.Module):\n",
    "    def __init__(self, query_model):\n",
    "        self.query_model = query_model\n",
    "\n",
    "    @tf.function()\n",
    "    def compute_emb(self, instances):\n",
    "        query_emb = self.query_model(instances)\n",
    "        return {\n",
    "            \"user_id\": instances[\"user_id\"],\n",
    "            \"gender\": instances[\"gender\"],\n",
    "            \"age\": instances[\"age\"],\n",
    "            \"country\": instances[\"country\"],\n",
    "            \"query_emb\": query_emb,\n",
    "        }\n",
    "\n",
    "# wrap query_model:   query_model -> query_model_module\n",
    "query_model = QueryModelModule(model.query_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49a54c08",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the input specifications for the instances\n",
    "instances_spec = {\n",
    "    'user_id': tf.TensorSpec(shape=(None,), dtype=tf.string, name='user_id'),   # Specification for user IDs\n",
    "    'gender': tf.TensorSpec(shape=(None,), dtype=tf.string, name='gender'),     # Specification for gender\n",
    "    'country': tf.TensorSpec(shape=(None,), dtype=tf.string, name='country'),    # Specification for country\n",
    "    'age': tf.TensorSpec(shape=(None,), dtype=tf.int64, name='age'),              # Specification for age\n",
    "}\n",
    "\n",
    "# Get the concrete function for the query_model's compute_emb function using the specified input signatures\n",
    "signatures = query_model.compute_emb.get_concrete_function(instances_spec)\n",
    "\n",
    "# Save the query_model along with the concrete function signatures\n",
    "tf.saved_model.save(\n",
    "    query_model,           # The model to save\n",
    "    \"query_model\",         # Path to save the model\n",
    "    signatures=signatures, # Concrete function signatures to include\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ed39777",
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.saved_model.save(\n",
    "    model.item_model,    # The model to save\n",
    "    \"candidate_model\",   # Path to save the model\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "baa9ce26",
   "metadata": {},
   "outputs": [],
   "source": [
    "from hsml.schema import Schema\n",
    "from hsml.model_schema import ModelSchema\n",
    "\n",
    "# Infer input schema from data.\n",
    "query_model_input_schema = Schema(query_df)\n",
    "\n",
    "# Manually specify output schema.\n",
    "query_model_output_schema = Schema([{\n",
    "    \"name\": \"query_embedding\",\n",
    "    \"type\": \"float32\",\n",
    "    \"shape\": [EMB_DIM],\n",
    "}])\n",
    "\n",
    "query_model_schema = ModelSchema(\n",
    "    input_schema=query_model_input_schema,\n",
    "    output_schema=query_model_output_schema,\n",
    ")\n",
    "\n",
    "query_model_schema.to_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "013612bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sample a query example from the query DataFrame\n",
    "query_example = query_df.sample().to_dict(\"records\")\n",
    "\n",
    "# Create a tensorflow model for the query_model in the Model Registry \n",
    "mr_query_model = mr.tensorflow.create_model(\n",
    "    name=\"query_model\",                                           # Name of the model\n",
    "    description=\"Model that generates query embeddings from user features\",  # Description of the model\n",
    "    input_example=query_example,                                  # Example input for the model\n",
    "    model_schema=query_model_schema,                              # Schema of the model\n",
    ")\n",
    "\n",
    "# Save the query_model to the Model Registry\n",
    "mr_query_model.save(\"query_model\")                                # Path to save the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4535c5ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the input schema for the candidate_model based on item_df\n",
    "candidate_model_input_schema = Schema(item_df)\n",
    "\n",
    "# Define the output schema for the candidate_model, specifying the shape and type of the output\n",
    "candidate_model_output_schema = Schema([{\n",
    "    \"name\": \"candidate_embedding\",   # Name of the output feature\n",
    "    \"type\": \"float32\",               # Data type of the output feature\n",
    "    \"shape\": [EMB_DIM],              # Shape of the output feature\n",
    "}])\n",
    "\n",
    "# Combine the input and output schemas to create the overall model schema for the candidate_model\n",
    "candidate_model_schema = ModelSchema(\n",
    "    input_schema=candidate_model_input_schema,    # Input schema for the model\n",
    "    output_schema=candidate_model_output_schema,  # Output schema for the model\n",
    ")\n",
    "\n",
    "# Sample a candidate example from the item DataFrame\n",
    "candidate_example = item_df.sample().to_dict(\"records\")\n",
    "\n",
    "# Create a tensorflow model for the candidate_model in the Model Registry\n",
    "mr_candidate_model = mr.tensorflow.create_model(\n",
    "    name=\"candidate_model\",                                        # Name of the model\n",
    "    description=\"Model that generates candidate embeddings from video features\",  # Description of the model\n",
    "    input_example=candidate_example,                              # Example input for the model\n",
    "    model_schema=candidate_model_schema,                          # Schema of the model\n",
    ")\n",
    "\n",
    "# Save the candidate_model to the Model Registry\n",
    "mr_candidate_model.save(\"candidate_model\")                        # Path to save the model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "657e0ba1",
   "metadata": {},
   "source": [
    "---"
   ]
  }
 ],
 "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}