Build a Movie Recommendation System
Train a Two-Tower model and deploy a recommendation API on AWS.
Introduction
30 Second Summary
Every time you open a streaming app, an invisible system decides what to show you next. Behind every "recommended for you" row is a pipeline of models retrieving candidates, ranking them, and learning from your clicks.
In this project, you will build that entire system from scratch. You will create a data pipeline, train a retrieval model, add a reinforcement learning layer that improves from user feedback, and deploy the whole thing as a live API.
What You'll Build
A production-grade movie recommendation API that uses a Two-Tower model for fast candidate retrieval, a REINFORCE policy gradient for re-ranking, and a Databricks medallion architecture pipeline for data preparation.
By the end of this project, you'll have:
- A Databricks ETL pipeline that transforms raw MovieLens data through Bronze, Silver, and Gold layers using Delta Lake tables.
- A trained Two-Tower retrieval model with full experiment tracking in MLflow (hyperparameters, metrics, and model artifacts).
- A live FastAPI service deployed on AWS ECS Fargate with a REINFORCE policy network that learns from user feedback to improve recommendations over time.
- Secret Mission: batch feedback replay that processes multiple feedback signals at once for more stable policy updates.
Are there any prerequisites?
You should be comfortable with basic Python and using the command line. A basic understanding of neural networks (what layers and embeddings are) will help, but you do not need prior experience with recommendation systems or reinforcement learning.
Build the Medallion ETL Pipeline in Databricks
Raw data is messy. Before any machine learning model can learn patterns, its training data needs to pass through structured layers of cleaning and transformation.
In this step, you will build a medallion architecture ETL pipeline in Databricks. The pipeline takes raw MovieLens data through three layers: Bronze (raw ingestion), Silver (cleaned and validated), and Gold (ML-ready features).
In this step, get ready to:
- Ingest raw MovieLens CSV data into a Bronze Delta table with metadata.
- Transform Bronze data into a cleaned Silver layer with proper types and deduplication.
- Build Gold-layer tables with user features and ML-ready interactions, then export for local training.
Create a notebook and upload the MovieLens dataset
The medallion architecture organizes data into three layers. Bronze holds raw, unmodified data. Silver contains cleaned and validated data. Gold stores business-ready datasets optimized for specific use cases like ML training.
- Sign in to your Databricks Free Edition workspace.
- Download the MovieLens ml-latest-small dataset to your computer by visiting this link:
https://files.grouplens.org/datasets/movielens/ml-latest-small.zip
- Unzip the downloaded file. You will find a folder called ml-latest-small containing several CSV files. You need ratings.csv and movies.csv.
- In the Databricks sidebar, click + New, then select Notebook.
- Rename the notebook to Medallion ETL Pipeline by clicking the title at the top and typing the new name.
What is a Databricks notebook?
A Databricks notebook is an interactive coding environment that runs Python code in cells. Each cell can be executed independently, which makes it easy to build a pipeline step by step.
In Databricks Free Edition, notebooks run on serverless compute. You do not need to configure any cluster.
- Upload the MovieLens data to a Unity Catalog volume. In the sidebar, click + New, then Add or upload data.
- Click Upload files to a volume.
- Click browse and select both ratings.csv and movies.csv from the unzipped ml-latest-small folder.
- Under Destination volume, select the default volume (or create a new volume called raw_data if none exists). Note the catalog and schema names shown (typically main and default).
- Wait for the upload to complete, then return to your notebook.
Why upload via the UI?
Databricks Free Edition restricts outbound internet access on serverless compute. You cannot download files from external URLs using code. Uploading via the workspace UI is the supported method for getting data into your volume.
Now write the Bronze layer code. The Bronze layer ingests raw data exactly as-is, adding only an ingestion_timestamp metadata column to track when the data was loaded.
- Paste the following code into the first cell of your notebook and run it by pressing Shift+Enter:
from pyspark.sql import functions as F
# Define paths to the uploaded CSV files in your volume
volume_path = "/Volumes/main/default/raw_data"
# Read raw ratings CSV preserving all original columns
bronze_ratings = (
spark.read.format("csv")
.option("header", "true")
.load(f"{volume_path}/ratings.csv")
.withColumn("ingestion_timestamp", F.current_timestamp())
)
# Write as a Bronze Delta table
bronze_ratings.write.format("delta").mode("overwrite").saveAsTable("main.default.bronze_ratings")
print(f"Bronze ratings: {bronze_ratings.count()} rows")
bronze_ratings.show(5)
What does this code do?
- spark.read.format("csv") reads the raw CSV file from the Unity Catalog volume using the path /Volumes/main/default/raw_data/ratings.csv.
- withColumn("ingestion_timestamp", ...) adds a metadata column recording exactly when the data was ingested. This is a Bronze-layer best practice for traceability.
- write.format("delta").saveAsTable(...) persists the data as a Delta Lake table, which gives you versioning, ACID transactions, and time-travel capabilities.
- Run the cell. You should see output showing 100,836 rows and a preview of the first 5 records with columns: userId, movieId, rating, timestamp, and ingestion_timestamp.
Seeing a file not found error?
The volume path must match where you uploaded the files. If you used a different volume name or your catalog is not main, update the volume_path variable to match. Check your actual path in the Catalog Explorer sidebar.
Help me fix my volume path in Databricks
Transform Bronze to Silver
The Silver layer is where data gets cleaned and validated. You will cast columns to proper data types, remove invalid entries, and deduplicate records.
- Add a new cell below and paste the following Silver layer transformation code:
from pyspark.sql.window import Window
# Read from the Bronze table
bronze_df = spark.table("main.default.bronze_ratings")
# Cast data types: userId and movieId to integer, rating to float, timestamp to datetime
silver_df = (
bronze_df
.withColumn("userId", F.col("userId").cast("integer"))
.withColumn("movieId", F.col("movieId").cast("integer"))
.withColumn("rating", F.col("rating").cast("float"))
.withColumn("timestamp", F.from_unixtime(F.col("timestamp")).cast("timestamp"))
)
# Remove nulls and invalid ratings (must be between 0.5 and 5.0)
silver_df = silver_df.filter(
(F.col("userId").isNotNull()) &
(F.col("movieId").isNotNull()) &
(F.col("rating").isNotNull()) &
(F.col("rating") >= 0.5) &
(F.col("rating") <= 5.0)
)
# Deduplicate on (userId, movieId) keeping the latest rating
window = Window.partitionBy("userId", "movieId").orderBy(F.col("timestamp").desc())
silver_df = (
silver_df
.withColumn("row_num", F.row_number().over(window))
.filter(F.col("row_num") == 1)
.drop("row_num")
)
# Write the Silver Delta table
silver_df.write.format("delta").mode("overwrite").saveAsTable("main.default.silver_ratings")
print(f"Silver ratings: {silver_df.count()} rows")
silver_df.printSchema()
What does this code do?
- Type casting ensures userId and movieId are integers (not strings from CSV), rating is a float, and timestamp becomes a proper datetime column.
- The filter removes any rows with null values or ratings outside the valid 0.5 to 5.0 range.
- The window function with row_number() handles cases where a user rated the same movie multiple times. It keeps only the most recent rating by ordering timestamps descending and taking row 1.
- Run the cell. You should see the row count (slightly less than 100,836 due to deduplication) and a schema showing the correct data types: integer for IDs, float for rating, and timestamp for the time column.
Seeing a table not found error?
Make sure you ran the Bronze cell first. The Silver layer reads from main.default.bronze_ratings, which must exist before this cell can execute.
Help me debug my Silver layer transformation
Build the Gold layer and export for training
The Gold layer produces ML-ready datasets. You will create two tables: one with per-user features (for future use) and one with the clean interactions that your Two-Tower model will train on.
- Add a new cell and paste the following code to create the gold_user_features table:
# Read Silver data and movies for genre info
silver_df = spark.table("main.default.silver_ratings")
movies_df = (
spark.read.format("csv")
.option("header", "true")
.load(f"{volume_path}/movies.csv")
)
# Join ratings with movies to access genre information
ratings_with_genres = silver_df.join(movies_df, on="movieId", how="left")
# Build per-user features: rating count, average rating, genre preferences
gold_user_features = (
ratings_with_genres
.groupBy("userId")
.agg(
F.count("rating").alias("rating_count"),
F.avg("rating").alias("avg_rating"),
F.collect_set("genres").alias("genre_preferences")
)
)
# Save as Gold Delta table
gold_user_features.write.format("delta").mode("overwrite").saveAsTable("main.default.gold_user_features")
print(f"Gold user features: {gold_user_features.count()} rows")
gold_user_features.show(5, truncate=False)
What does this code do?
- The join combines each rating with its movie's genre information (the genres column from movies.csv).
- groupBy("userId") aggregates all ratings by user, computing count, average, and the set of genres they have watched.
- These features could be used for cold-start handling or user segmentation in a production system.
- Run the cell. You should see 610 rows (one per user) with columns for rating count, average rating, and genre preferences.
Now create the gold_interactions table and export it as a Parquet file for local model training.
- Add a new cell and paste the following code:
from pyspark.sql.functions import min as spark_min, max as spark_max
# Calculate min and max rating for normalization
stats = silver_df.agg(
spark_min("rating").alias("min_rating"),
spark_max("rating").alias("max_rating")
).collect()[0]
min_rating = stats["min_rating"]
max_rating = stats["max_rating"]
# Create interactions table with normalized ratings
gold_interactions = (
silver_df
.select("userId", "movieId", "rating", "timestamp")
.withColumn(
"normalized_rating",
(F.col("rating") - min_rating) / (max_rating - min_rating)
)
)
# Save as Gold Delta table
gold_interactions.write.format("delta").mode("overwrite").saveAsTable("main.default.gold_interactions")
# Export as Parquet for local model training
gold_interactions.write.format("parquet").mode("overwrite").save(f"{volume_path}/gold_interactions.parquet")
print(f"Gold interactions: {gold_interactions.count()} rows")
print("Exported gold_interactions.parquet to volume for download.")
gold_interactions.show(5)
What does this code do?
- The normalized rating scales values from the original 0.5-5.0 range into 0.0-1.0. This helps the model train more efficiently.
- The table keeps userId and movieId as the core interaction pair your Two-Tower model needs.
- The final .write.format("parquet") exports the data to a file you can download and use for local training outside of Databricks.
- Run the cell. You should see the interaction count and a preview showing userId, movieId, rating, timestamp, and normalized_rating columns.
Export failed with a permission error?
Ensure you have WRITE VOLUME permission on your volume. In Databricks Free Edition, the default volume in main.default should grant this automatically. If you created a custom volume, check its permissions in the Catalog Explorer.
Help me fix my Parquet export permissions
✔️ Awesome, I've got everything!
Great. Double check that all four cells ran without errors and you see the exported Parquet confirmation message in Cell 4.
ⓧ I'd like to double check the full code
Here is the complete notebook code across all four cells:
# Cell 1: Bronze Layer
from pyspark.sql import functions as F
# Define paths to the uploaded CSV files in your volume
volume_path = "/Volumes/main/default/raw_data"
# Read raw ratings CSV preserving all original columns
bronze_ratings = (
spark.read.format("csv")
.option("header", "true")
.load(f"{volume_path}/ratings.csv")
.withColumn("ingestion_timestamp", F.current_timestamp())
)
# Write as a Bronze Delta table
bronze_ratings.write.format("delta").mode("overwrite").saveAsTable("main.default.bronze_ratings")
print(f"Bronze ratings: {bronze_ratings.count()} rows")
bronze_ratings.show(5)
# Cell 2: Silver Layer
from pyspark.sql.window import Window
# Read from the Bronze table
bronze_df = spark.table("main.default.bronze_ratings")
# Cast data types: userId and movieId to integer, rating to float, timestamp to datetime
silver_df = (
bronze_df
.withColumn("userId", F.col("userId").cast("integer"))
.withColumn("movieId", F.col("movieId").cast("integer"))
.withColumn("rating", F.col("rating").cast("float"))
.withColumn("timestamp", F.from_unixtime(F.col("timestamp")).cast("timestamp"))
)
# Remove nulls and invalid ratings (must be between 0.5 and 5.0)
silver_df = silver_df.filter(
(F.col("userId").isNotNull()) &
(F.col("movieId").isNotNull()) &
(F.col("rating").isNotNull()) &
(F.col("rating") >= 0.5) &
(F.col("rating") <= 5.0)
)
# Deduplicate on (userId, movieId) keeping the latest rating
window = Window.partitionBy("userId", "movieId").orderBy(F.col("timestamp").desc())
silver_df = (
silver_df
.withColumn("row_num", F.row_number().over(window))
.filter(F.col("row_num") == 1)
.drop("row_num")
)
# Write the Silver Delta table
silver_df.write.format("delta").mode("overwrite").saveAsTable("main.default.silver_ratings")
print(f"Silver ratings: {silver_df.count()} rows")
silver_df.printSchema()
# Cell 3: Gold User Features
# Read Silver data and movies for genre info
silver_df = spark.table("main.default.silver_ratings")
movies_df = (
spark.read.format("csv")
.option("header", "true")
.load(f"{volume_path}/movies.csv")
)
# Join ratings with movies to access genre information
ratings_with_genres = silver_df.join(movies_df, on="movieId", how="left")
# Build per-user features: rating count, average rating, genre preferences
gold_user_features = (
ratings_with_genres
.groupBy("userId")
.agg(
F.count("rating").alias("rating_count"),
F.avg("rating").alias("avg_rating"),
F.collect_set("genres").alias("genre_preferences")
)
)
# Save as Gold Delta table
gold_user_features.write.format("delta").mode("overwrite").saveAsTable("main.default.gold_user_features")
print(f"Gold user features: {gold_user_features.count()} rows")
gold_user_features.show(5, truncate=False)
# Cell 4: Gold Interactions + Export
from pyspark.sql.functions import min as spark_min, max as spark_max
# Calculate min and max rating for normalization
stats = silver_df.agg(
spark_min("rating").alias("min_rating"),
spark_max("rating").alias("max_rating")
).collect()[0]
min_rating = stats["min_rating"]
max_rating = stats["max_rating"]
# Create interactions table with normalized ratings
gold_interactions = (
silver_df
.select("userId", "movieId", "rating", "timestamp")
.withColumn(
"normalized_rating",
(F.col("rating") - min_rating) / (max_rating - min_rating)
)
)
# Save as Gold Delta table
gold_interactions.write.format("delta").mode("overwrite").saveAsTable("main.default.gold_interactions")
# Export as Parquet for local model training
gold_interactions.write.format("parquet").mode("overwrite").save(f"{volume_path}/gold_interactions.parquet")
print(f"Gold interactions: {gold_interactions.count()} rows")
print("Exported gold_interactions.parquet to volume for download.")
gold_interactions.show(5)
Your medallion architecture is complete. You have raw data preserved in Bronze, clean validated data in Silver, and ML-ready features in Gold. Next up, you will download the exported Parquet file and set up your local environment for model training.
Set Up Local Project and Install Dependencies
Your Databricks ETL pipeline is complete. The Gold layer data is ready for model training, sitting in a clean Parquet file with normalized user-item interactions.
Now you need to bring that data to your local machine and set up the Python environment where you will train the Two-Tower model and build the FastAPI service. You will also configure MLflow for experiment tracking so every training run is logged and reproducible.
In this step, get ready to:
- Create a local project directory with a Python virtual environment and install all dependencies.
- Download the Gold-layer data and MovieLens metadata for local use.
- Initialize MLflow experiment tracking and confirm the UI is accessible.
Create the project directory and install dependencies
Every ML project benefits from an isolated Python environment. A virtual environment ensures your project dependencies do not conflict with other Python tools on your machine.
- Open your terminal. Press Cmd+Space (macOS) or the Windows key (Windows) to open your search bar.
- Type Terminal (macOS) or PowerShell (Windows) and press Enter to open it.
First, confirm you have Python 3.10 or higher installed. Check your version by running:
python3 --version
✔️ I see version 3.10 or higher
You are ready to continue. Your Python version supports all the dependencies this project uses.
ⓧ I see an older version
You need Python 3.10 or higher. On macOS, your system Python may be outdated even if you installed a newer version separately.
- Download and install Python 3.10+ from python.org/downloads.
- After installing, close and reopen your terminal, then run python3 --version again to confirm.
ⓧ Command not found
Python is not installed or not in your PATH.
- Download and install Python 3.10+ from python.org/downloads.
- On Windows, check the Add Python to PATH option during installation.
- After installing, close and reopen your terminal, then run python3 --version again.
- Move to your Desktop by running this command:
cd ~/Desktop
- Create your project folder and move into it by running these commands:
mkdir rec-system
cd rec-system
- Create a Python virtual environment by running:
python3 -m venv .venv
- Activate the virtual environment:
macOS
source .venv/bin/activate
Windows
.venv\Scripts\activate
You should see (.venv) appear at the beginning of your terminal prompt. This confirms the virtual environment is active.
- Create a requirements.txt file for the project's production dependencies by running:
cat > requirements.txt << 'EOF'
torch==2.12.0
fastapi[standard]==0.136.3
pandas==2.2.3
numpy==2.2.6
EOF
Why a separate requirements.txt?
This file lists only the packages your production Docker container will need. MLflow is a development tool for experiment tracking, so it gets installed separately and does not ship with the deployed API.
- Install the production dependencies and MLflow by running:
pip install -r requirements.txt mlflow==2.22.5
What does this install?
- PyTorch (torch==2.12.0) provides the neural network framework for the Two-Tower model.
- FastAPI (fastapi[standard]==0.136.3) is the web framework for the recommendation API.
- MLflow (mlflow==2.22.5) handles experiment tracking for hyperparameters, metrics, and model artifacts.
- pandas and numpy handle data loading and numerical operations.
This installation may take a few minutes. PyTorch is a large package.
Seeing installation errors?
- Make sure (.venv) appears in your prompt. If not, re-run the activation command for your platform.
- If you see a permissions error, confirm you are inside the virtual environment and not using sudo.
- If a specific package fails to resolve, confirm your Python version is 3.10+ by running python3 --version inside the venv.
Help me troubleshoot my pip install errors.
Download the Gold-layer data and movie metadata
Your Two-Tower model needs two data files. The Gold-layer Parquet file contains the clean user-item interactions for training. The MovieLens movies.csv provides movie titles for the API responses.
- Switch back to your Databricks workspace in the browser.
- In the left sidebar, click Catalog.
- Navigate to the volume where you saved gold_interactions.parquet in Step 1.
- Click on gold_interactions.parquet to select it.
- Click Download to save the file to your computer.
- Move the downloaded file into your rec-system project folder on your Desktop.
Where did this Parquet file come from?
In Step 1, your Databricks notebook exported the Gold-layer interactions table as a Parquet file to a Unity Catalog volume. This file contains the clean, deduplicated user-item pairs with normalized ratings that the model will train on.
- Download the MovieLens dataset for movie title lookups by running this command in your terminal (make sure you are in the rec-system folder):
curl -L -o ml-latest-small.zip https://files.grouplens.org/datasets/movielens/ml-latest-small.zip
unzip ml-latest-small.zip
What is in this zip?
The MovieLens ml-latest-small dataset contains 100,836 ratings across 9,742 movies from 610 users. You only need the movies.csv file for title lookups in the API. The extracted folder ml-latest-small/ contains it alongside other files you can ignore.
- Verify both data files are accessible by running this quick Python check:
python3 -c "
import pandas as pd
# Verify gold interactions
interactions = pd.read_parquet('gold_interactions.parquet')
print(f'Interactions: {len(interactions)} rows')
print(f'Columns: {list(interactions.columns)}')
# Verify movies
movies = pd.read_csv('ml-latest-small/movies.csv')
print(f'Movies: {len(movies)} rows')
print(f'Sample: {movies.head(2).to_string()}')
"
You should see output showing the number of interaction rows, the column names, and a sample of movie titles.
FileNotFoundError for the Parquet file?
- Confirm you moved gold_interactions.parquet into the rec-system folder (not a subfolder like Downloads).
- Run ls in your terminal to check the file is in the current directory.
- Make sure the filename matches exactly. Databricks may have added a suffix during download.
Help me find my downloaded Parquet file.
Initialize MLflow experiment tracking
MLflow tracks every training experiment so you can compare hyperparameters, metrics, and models across runs. You will start the MLflow UI server and create the experiment that the training script will log to.
- Start the MLflow tracking UI by running:
mlflow ui &
What does this command do?
The mlflow ui command launches a local web server that reads from the ./mlruns directory in your project folder. The & at the end runs it in the background so you can keep using the same terminal.
- Open your browser and navigate to http://localhost:5000.
- You should see the MLflow Tracking UI with a default experiment listed.
- Create the experiment that your training script will log to by running:
python3 -c "import mlflow; mlflow.set_experiment('two-tower-recsys'); print('Experiment created!')"
You should see Experiment created! in your terminal output.
- Refresh http://localhost:5000 in your browser. You should now see two-tower-recsys listed as an experiment in the sidebar.
MLflow UI not loading?
- Check that port 5000 is not already in use. On macOS, AirPlay Receiver sometimes uses port 5000. Disable it in System Settings > General > AirDrop & Handoff.
- If the command failed, make sure your virtual environment is still active (you should see (.venv) in your prompt).
- Try running mlflow ui --port 5001 if port 5000 is unavailable, then visit http://localhost:5001 instead.
Help me get MLflow UI running.
✔️ Awesome, I've got everything!
Great. Double check you have saved requirements.txt and that both data files are in your rec-system folder.
ⓧ I'd like to double check the full code
Here is what your requirements.txt file should contain:
torch==2.12.0
fastapi[standard]==0.136.3
pandas==2.2.3
numpy==2.2.6
Your project folder (rec-system) on your Desktop should contain:
- requirements.txt (the file above).
- gold_interactions.parquet (downloaded from Databricks).
- ml-latest-small/ folder containing movies.csv (from the MovieLens zip).
- .venv/ folder (your virtual environment).
- mlruns/ folder (created by MLflow when you set the experiment).
Your local environment is fully configured. Next up, you will build and train the Two-Tower retrieval model using this data, with every experiment tracked in MLflow.
Build and Train the Two-Tower Model with MLflow
In the previous step, you set up your local environment and verified the Gold-layer data is ready for model training.
Now you will build the core of your recommendation system: a Two-Tower retrieval model. This is the same architecture used by YouTube, Spotify, and TikTok to retrieve personalized candidates from millions of items in milliseconds. You will track every experiment in MLflow so that hyperparameters, loss curves, and model artifacts are reproducible.
In this step, get ready to:
- Define the Two-Tower model architecture in a reusable module.
- Build a training pipeline with negative sampling and MLflow tracking.
- Train for 15 epochs, evaluate Recall@10, and save the trained checkpoint.
Implement the Two-Tower architecture
A Two-Tower model has two independent neural networks: one for users and one for items. Each tower maps an ID to a dense embedding vector. At serving time, you score candidates by computing the dot product between the user embedding and each item embedding.
Why Two-Tower instead of a joint model?
Item embeddings can be precomputed offline because the item tower is independent of the user. At inference time you only compute the user embedding live and look up pre-stored item vectors. This makes retrieval scale to millions of items.
- Create a new file called model.py in your project root.
- Add the UserTower and ItemTower classes by pasting this code:
import torch
import torch.nn as nn
class UserTower(nn.Module):
def __init__(self, num_users, embedding_dim=64, hidden_dim=128):
super().__init__()
# Embed user IDs into a dense vector
self.embedding = nn.Embedding(num_users, hidden_dim)
# MLP projects to the final embedding space
self.mlp = nn.Sequential(
nn.Linear(hidden_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, embedding_dim),
)
def forward(self, user_ids):
x = self.embedding(user_ids)
return self.mlp(x)
class ItemTower(nn.Module):
def __init__(self, num_items, embedding_dim=64, hidden_dim=128):
super().__init__()
# Embed item IDs into a dense vector
self.embedding = nn.Embedding(num_items, hidden_dim)
# MLP projects to the final embedding space
self.mlp = nn.Sequential(
nn.Linear(hidden_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, embedding_dim),
)
def forward(self, item_ids):
x = self.embedding(item_ids)
return self.mlp(x)
What does this code do?
Each tower follows the same pattern: an nn.Embedding layer maps discrete IDs to a hidden_dim-sized vector, then a two-layer MLP projects it down to a 64-dimensional embedding.
The user tower and item tower are mirrors of each other. They produce embeddings in the same 64-dimensional space so you can compare them with a dot product.
- Below the ItemTower class, add the TwoTowerModel that combines both towers, plus a PolicyNetwork you will use in the next step:
class TwoTowerModel(nn.Module):
def __init__(self, num_users, num_items, embedding_dim=64, hidden_dim=128):
super().__init__()
self.user_tower = UserTower(num_users, embedding_dim, hidden_dim)
self.item_tower = ItemTower(num_items, embedding_dim, hidden_dim)
def forward(self, user_ids, item_ids):
# Get embeddings from each tower
user_emb = self.user_tower(user_ids)
item_emb = self.item_tower(item_ids)
# Score via dot product
scores = (user_emb * item_emb).sum(dim=1)
return scores
def get_user_embedding(self, user_ids):
return self.user_tower(user_ids)
def get_item_embedding(self, item_ids):
return self.item_tower(item_ids)
class PolicyNetwork(nn.Module):
def __init__(self, input_dim, hidden_dim=64):
super().__init__()
# Scores each candidate for re-ranking
self.net = nn.Sequential(
nn.Linear(input_dim, hidden_dim),
nn.ReLU(),
nn.Dropout(0.1),
nn.Linear(hidden_dim, 1),
)
def forward(self, candidate_features):
scores = self.net(candidate_features).squeeze(-1)
# Softmax converts scores to a probability distribution
probs = torch.softmax(scores, dim=0)
return probs
What does this code do?
TwoTowerModel combines both towers and scores user-item pairs via element-wise multiplication followed by a sum (dot product). It also exposes helper methods to extract embeddings independently for fast retrieval.
PolicyNetwork is a small re-ranker you will use in the next step. It takes candidate features and outputs a probability distribution over them. You are defining it here so all model architectures live in one file.
- Save model.py.
Seeing an import error?
Make sure your virtual environment is activated and torch is installed. Run pip list | grep torch to confirm.
help me debug a torch import error in model.py
✔️ Awesome, I've got everything!
Great. Double check you have saved model.py before continuing.
ⓧ I'd like to double check the full code
import torch
import torch.nn as nn
class UserTower(nn.Module):
def __init__(self, num_users, embedding_dim=64, hidden_dim=128):
super().__init__()
self.embedding = nn.Embedding(num_users, hidden_dim)
self.mlp = nn.Sequential(
nn.Linear(hidden_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, embedding_dim),
)
def forward(self, user_ids):
x = self.embedding(user_ids)
return self.mlp(x)
class ItemTower(nn.Module):
def __init__(self, num_items, embedding_dim=64, hidden_dim=128):
super().__init__()
self.embedding = nn.Embedding(num_items, hidden_dim)
self.mlp = nn.Sequential(
nn.Linear(hidden_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, embedding_dim),
)
def forward(self, item_ids):
x = self.embedding(item_ids)
return self.mlp(x)
class TwoTowerModel(nn.Module):
def __init__(self, num_users, num_items, embedding_dim=64, hidden_dim=128):
super().__init__()
self.user_tower = UserTower(num_users, embedding_dim, hidden_dim)
self.item_tower = ItemTower(num_items, embedding_dim, hidden_dim)
def forward(self, user_ids, item_ids):
user_emb = self.user_tower(user_ids)
item_emb = self.item_tower(item_ids)
scores = (user_emb * item_emb).sum(dim=1)
return scores
def get_user_embedding(self, user_ids):
return self.user_tower(user_ids)
def get_item_embedding(self, item_ids):
return self.item_tower(item_ids)
class PolicyNetwork(nn.Module):
def __init__(self, input_dim, hidden_dim=64):
super().__init__()
self.net = nn.Sequential(
nn.Linear(input_dim, hidden_dim),
nn.ReLU(),
nn.Dropout(0.1),
nn.Linear(hidden_dim, 1),
)
def forward(self, candidate_features):
scores = self.net(candidate_features).squeeze(-1)
probs = torch.softmax(scores, dim=0)
return probs
Create the training dataset with negative sampling
A recommendation model needs both positive examples (movies a user actually watched) and negative examples (movies they did not interact with). This technique is called negative sampling. It teaches the model to distinguish relevant items from irrelevant ones.
- Create a new file called train.py in your project root.
- Add the imports, dataset class, and negative sampling function by pasting this code:
import pandas as pd
import numpy as np
import torch
import mlflow
from torch.utils.data import Dataset, DataLoader
from model import TwoTowerModel
class InteractionDataset(Dataset):
def __init__(self, user_ids, item_ids, labels):
self.user_ids = torch.tensor(user_ids, dtype=torch.long)
self.item_ids = torch.tensor(item_ids, dtype=torch.long)
self.labels = torch.tensor(labels, dtype=torch.float32)
def __len__(self):
return len(self.labels)
def __getitem__(self, idx):
return self.user_ids[idx], self.item_ids[idx], self.labels[idx]
def create_negative_samples(interactions_df, num_items, neg_ratio=4):
users, items, labels = [], [], []
all_items = set(range(num_items))
# Group items each user has interacted with
user_items = interactions_df.groupby("user_idx")["item_idx"].apply(set).to_dict()
for _, row in interactions_df.iterrows():
# Positive sample: user actually watched this movie
users.append(row["user_idx"])
items.append(row["item_idx"])
labels.append(1.0)
# Negative samples: random movies the user did NOT watch
negative_pool = list(all_items - user_items.get(row["user_idx"], set()))
neg_samples = np.random.choice(negative_pool, size=min(neg_ratio, len(negative_pool)), replace=False)
for neg in neg_samples:
users.append(row["user_idx"])
items.append(neg)
labels.append(0.0)
return users, items, labels
What does this code do?
InteractionDataset wraps user-item pairs into a PyTorch Dataset so the DataLoader can batch them efficiently.
create_negative_samples generates 4 random negatives for every positive interaction. This 1:4 ratio gives the model enough signal to learn which items are relevant vs irrelevant for each user.
- Below the functions, add the data loading and setup code that prepares everything for training:
# Load gold-layer data
interactions = pd.read_parquet("gold_interactions.parquet")
# Create contiguous ID mappings
user_ids = interactions["userId"].unique()
item_ids = interactions["movieId"].unique()
user_to_idx = {uid: idx for idx, uid in enumerate(user_ids)}
item_to_idx = {mid: idx for idx, mid in enumerate(item_ids)}
idx_to_item = {v: k for k, v in item_to_idx.items()}
interactions["user_idx"] = interactions["userId"].map(user_to_idx)
interactions["item_idx"] = interactions["movieId"].map(item_to_idx)
num_users = len(user_ids)
num_items = len(item_ids)
# Train/test split
train_df = interactions.sample(frac=0.8, random_state=42)
test_df = interactions.drop(train_df.index)
# Hyperparameters
embedding_dim = 64
hidden_dim = 128
learning_rate = 0.001
neg_ratio = 4
num_epochs = 15
batch_size = 512
# Create training data with negative sampling
train_users, train_items, train_labels = create_negative_samples(train_df, num_items, neg_ratio)
train_dataset = InteractionDataset(train_users, train_items, train_labels)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
# Initialize model
model = TwoTowerModel(num_users, num_items, embedding_dim=embedding_dim, hidden_dim=hidden_dim)
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
criterion = torch.nn.BCEWithLogitsLoss()
What does this code do?
First, it loads the Gold-layer Parquet and remaps original user/movie IDs to contiguous indices starting at 0. PyTorch embeddings require contiguous integer indices.
Then it splits the data 80/20, generates negative samples for the training set, and initializes the Two-Tower model with an Adam optimizer and binary cross-entropy loss.
- Save train.py.
FileNotFoundError on gold_interactions.parquet?
Make sure gold_interactions.parquet is in your project root directory (the same folder as train.py). If you saved it in a subfolder during the previous step, move it to the root.
help me fix a FileNotFoundError for gold_interactions.parquet
Train the model and log results to MLflow
With data preparation done, you will now wrap the training loop inside an MLflow run. Every hyperparameter, per-epoch loss, final evaluation metric, and the model artifact itself gets logged automatically.
- At the bottom of train.py, add the MLflow tracking block that runs training and evaluation:
# MLflow experiment tracking
mlflow.set_experiment("two-tower-recsys")
with mlflow.start_run(run_name="two-tower-training"):
# Log hyperparameters
mlflow.log_params({
"embedding_dim": embedding_dim,
"hidden_dim": hidden_dim,
"learning_rate": learning_rate,
"neg_ratio": neg_ratio,
"num_epochs": num_epochs,
"batch_size": batch_size,
"num_users": num_users,
"num_items": num_items,
})
# Training loop
for epoch in range(num_epochs):
model.train()
total_loss = 0
for users, items, labels in train_loader:
optimizer.zero_grad()
scores = model(users, items)
loss = criterion(scores, labels)
loss.backward()
optimizer.step()
total_loss += loss.item()
avg_loss = total_loss / len(train_loader)
mlflow.log_metric("train_loss", avg_loss, step=epoch)
print(f"Epoch {epoch + 1}/{num_epochs}, Loss: {avg_loss:.4f}")
# Evaluate Recall@10
model.eval()
hits = 0
total = 0
with torch.no_grad():
all_item_ids = torch.arange(num_items, dtype=torch.long)
all_item_embs = model.get_item_embedding(all_item_ids)
for user_idx in range(min(num_users, 200)):
user_tensor = torch.tensor([user_idx], dtype=torch.long)
user_emb = model.get_user_embedding(user_tensor)
scores = (user_emb @ all_item_embs.T).squeeze(0)
top_10 = torch.topk(scores, 10).indices.tolist()
actual_items = set(test_df[test_df["user_idx"] == user_idx]["item_idx"].tolist())
if actual_items:
hits += len(set(top_10) & actual_items)
total += len(actual_items)
recall_at_10 = hits / total if total > 0 else 0
mlflow.log_metric("recall_at_10", recall_at_10)
print(f"Recall@10: {recall_at_10:.4f}")
# Log model artifact to MLflow
mlflow.pytorch.log_model(model, "two_tower_model")
What does this code do?
The mlflow.start_run() context manager creates a tracked run named two-tower-training. Inside it, log_params records all hyperparameters and log_metric records the loss after every epoch.
After training, the code evaluates Recall@10 on 200 test users. Recall@10 measures how often the model's top-10 predictions contain items the user actually watched. Finally, mlflow.pytorch.log_model() saves the full model as a reusable artifact.
- Below the MLflow block, add the final section that precomputes item embeddings and saves a local checkpoint for serving:
# Precompute item embeddings and save for serving
with torch.no_grad():
all_item_embs = model.get_item_embedding(torch.arange(num_items, dtype=torch.long))
torch.save(
{
"model_state_dict": model.state_dict(),
"item_embeddings": all_item_embs,
"num_users": num_users,
"num_items": num_items,
"user_to_idx": user_to_idx,
"item_to_idx": item_to_idx,
"idx_to_item": idx_to_item,
},
"two_tower_model.pt",
)
print("Model saved to two_tower_model.pt")
Why precompute item embeddings?
At serving time, you only need to compute the user embedding live. All item embeddings are fixed after training, so precomputing them avoids running the item tower thousands of times per request. The checkpoint also stores ID mappings so the serving layer can translate between original MovieLens IDs and model indices.
- Save train.py.
- Run the training script by executing this command:
python train.py
Training takes 1-3 minutes on CPU depending on your machine. You should see epoch loss values decreasing over 15 epochs, followed by a Recall@10 score and a confirmation that the model was saved.
What should I see?
Your terminal output will look similar to this:
Epoch 1/15, Loss: 0.6412
Epoch 2/15, Loss: 0.5873
...
Epoch 15/15, Loss: 0.3241
Recall@10: 0.0523
Model saved to two_tower_model.pt
The exact numbers will vary because of random initialization, but loss should decrease and Recall@10 should be above 0.
- Switch to your browser and open the MLflow UI at http://localhost:5000.
- Click on the two-tower-recsys experiment in the left sidebar.
- Click the two-tower-training run to see logged parameters, metrics, and the model artifact.
Training script crashes or hangs?
If you see a ModuleNotFoundError: No module named 'model' error, make sure model.py is saved in the same directory where you are running train.py.
If the script seems to hang during negative sampling, this is normal for the first run. The MovieLens dataset has around 100,000 interactions, and generating 4 negatives each takes a moment.
help me debug my train.py script that is crashing or hanging
✔️ Awesome, I've got everything!
Great. Double check you have saved train.py and confirmed the MLflow run appears in your browser.
ⓧ I'd like to double check the full code
import pandas as pd
import numpy as np
import torch
import mlflow
from torch.utils.data import Dataset, DataLoader
from model import TwoTowerModel
class InteractionDataset(Dataset):
def __init__(self, user_ids, item_ids, labels):
self.user_ids = torch.tensor(user_ids, dtype=torch.long)
self.item_ids = torch.tensor(item_ids, dtype=torch.long)
self.labels = torch.tensor(labels, dtype=torch.float32)
def __len__(self):
return len(self.labels)
def __getitem__(self, idx):
return self.user_ids[idx], self.item_ids[idx], self.labels[idx]
def create_negative_samples(interactions_df, num_items, neg_ratio=4):
users, items, labels = [], [], []
all_items = set(range(num_items))
user_items = interactions_df.groupby("user_idx")["item_idx"].apply(set).to_dict()
for _, row in interactions_df.iterrows():
users.append(row["user_idx"])
items.append(row["item_idx"])
labels.append(1.0)
negative_pool = list(all_items - user_items.get(row["user_idx"], set()))
neg_samples = np.random.choice(negative_pool, size=min(neg_ratio, len(negative_pool)), replace=False)
for neg in neg_samples:
users.append(row["user_idx"])
items.append(neg)
labels.append(0.0)
return users, items, labels
# Load gold-layer data
interactions = pd.read_parquet("gold_interactions.parquet")
# Create contiguous ID mappings
user_ids = interactions["userId"].unique()
item_ids = interactions["movieId"].unique()
user_to_idx = {uid: idx for idx, uid in enumerate(user_ids)}
item_to_idx = {mid: idx for idx, mid in enumerate(item_ids)}
idx_to_item = {v: k for k, v in item_to_idx.items()}
interactions["user_idx"] = interactions["userId"].map(user_to_idx)
interactions["item_idx"] = interactions["movieId"].map(item_to_idx)
num_users = len(user_ids)
num_items = len(item_ids)
# Train/test split
train_df = interactions.sample(frac=0.8, random_state=42)
test_df = interactions.drop(train_df.index)
# Hyperparameters
embedding_dim = 64
hidden_dim = 128
learning_rate = 0.001
neg_ratio = 4
num_epochs = 15
batch_size = 512
# Create training data with negative sampling
train_users, train_items, train_labels = create_negative_samples(train_df, num_items, neg_ratio)
train_dataset = InteractionDataset(train_users, train_items, train_labels)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
# Initialize model
model = TwoTowerModel(num_users, num_items, embedding_dim=embedding_dim, hidden_dim=hidden_dim)
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
criterion = torch.nn.BCEWithLogitsLoss()
# MLflow experiment tracking
mlflow.set_experiment("two-tower-recsys")
with mlflow.start_run(run_name="two-tower-training"):
# Log hyperparameters
mlflow.log_params({
"embedding_dim": embedding_dim,
"hidden_dim": hidden_dim,
"learning_rate": learning_rate,
"neg_ratio": neg_ratio,
"num_epochs": num_epochs,
"batch_size": batch_size,
"num_users": num_users,
"num_items": num_items,
})
# Training loop
for epoch in range(num_epochs):
model.train()
total_loss = 0
for users, items, labels in train_loader:
optimizer.zero_grad()
scores = model(users, items)
loss = criterion(scores, labels)
loss.backward()
optimizer.step()
total_loss += loss.item()
avg_loss = total_loss / len(train_loader)
mlflow.log_metric("train_loss", avg_loss, step=epoch)
print(f"Epoch {epoch + 1}/{num_epochs}, Loss: {avg_loss:.4f}")
# Evaluate Recall@10
model.eval()
hits = 0
total = 0
with torch.no_grad():
all_item_ids = torch.arange(num_items, dtype=torch.long)
all_item_embs = model.get_item_embedding(all_item_ids)
for user_idx in range(min(num_users, 200)):
user_tensor = torch.tensor([user_idx], dtype=torch.long)
user_emb = model.get_user_embedding(user_tensor)
scores = (user_emb @ all_item_embs.T).squeeze(0)
top_10 = torch.topk(scores, 10).indices.tolist()
actual_items = set(test_df[test_df["user_idx"] == user_idx]["item_idx"].tolist())
if actual_items:
hits += len(set(top_10) & actual_items)
total += len(actual_items)
recall_at_10 = hits / total if total > 0 else 0
mlflow.log_metric("recall_at_10", recall_at_10)
print(f"Recall@10: {recall_at_10:.4f}")
# Log model artifact to MLflow
mlflow.pytorch.log_model(model, "two_tower_model")
# Precompute item embeddings and save for serving
with torch.no_grad():
all_item_embs = model.get_item_embedding(torch.arange(num_items, dtype=torch.long))
torch.save(
{
"model_state_dict": model.state_dict(),
"item_embeddings": all_item_embs,
"num_users": num_users,
"num_items": num_items,
"user_to_idx": user_to_idx,
"item_to_idx": item_to_idx,
"idx_to_item": idx_to_item,
},
"two_tower_model.pt",
)
print("Model saved to two_tower_model.pt")
Your Two-Tower model is trained and tracked. Next up, you will build a REINFORCE policy network that learns to re-rank these candidates based on real-time user feedback.
Implement the REINFORCE Policy Network
Your Two-Tower model retrieves the top candidates by raw dot-product score. But raw retrieval scores treat every user the same way. There is no mechanism to learn from individual preferences over time.
In this step, you will build a REINFORCE policy network that re-ranks those candidates based on learned preferences. When a user gives feedback (thumbs up or thumbs down), the policy updates itself to serve better recommendations next time.
In this step, get ready to:
- Set up the RecommendationEngine class that loads the trained Two-Tower model and instantiates a PolicyNetwork as a re-ranker.
- Implement the REINFORCE policy gradient update that learns from user feedback signals.
- Build the recommendation pipeline that retrieves candidates via dot product, re-ranks them through the policy, and stores state for feedback.
Build the RecommendationEngine with PolicyNetwork re-ranker
The PolicyNetwork class already exists in model.py. It takes a feature vector per candidate and outputs a probability distribution over items. Now you need a class that orchestrates the full pipeline: load the trained Two-Tower model, instantiate the policy, and wire them together.
- Create a new file called engine.py in your project root.
- Add the imports and the RecommendationEngine.__init__ method by pasting this code:
import torch
from torch.distributions import Categorical
from model import TwoTowerModel, PolicyNetwork
class RecommendationEngine:
def __init__(self, checkpoint_path, num_candidates=50):
self.checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
self.num_candidates = num_candidates
# Load Two-Tower model
self.model = TwoTowerModel(
self.checkpoint["num_users"],
self.checkpoint["num_items"],
embedding_dim=64,
hidden_dim=128,
)
self.model.load_state_dict(self.checkpoint["model_state_dict"])
self.model.eval()
# Precomputed item embeddings
self.item_embeddings = self.checkpoint["item_embeddings"]
# Policy network: input is user_emb (64) + item_emb (64) + score (1) = 129
self.policy = PolicyNetwork(input_dim=129, hidden_dim=64)
self.policy_optimizer = torch.optim.Adam(self.policy.parameters(), lr=0.001)
# Store log probs for REINFORCE
self.saved_log_probs = {}
What does this code do?
The __init__ method sets up three components that work together at serving time:
- The trained TwoTowerModel is loaded from the checkpoint and set to evaluation mode. It generates user embeddings on demand.
- Precomputed item_embeddings are loaded directly from the checkpoint. Since item embeddings never change at serving time, precomputing them avoids redundant forward passes.
- The PolicyNetwork uses input_dim=129 because each candidate feature is a concatenation of the 64-dim user embedding, the 64-dim item embedding, and the 1-dim retrieval score.
- The saved_log_probs dictionary stores the log probability of each recommendation so the policy can learn from later feedback.
Why a separate re-ranker on top of retrieval?
Two-Tower retrieval is fast but static. It scores items the same way regardless of recent user behavior. A policy re-ranker adds a learnable layer that adapts in real time.
This is the same retrieve-then-rank pattern used by YouTube and Spotify. The retrieval model narrows thousands of items to a short candidate list. The policy decides the final ordering.
- Save engine.py.
Implement the REINFORCE feedback update
The core of REINFORCE is surprisingly simple. When a user provides feedback, you retrieve the log probability of the action that produced that recommendation and multiply it by the reward signal. A positive reward reinforces the action. A negative reward discourages it.
- In engine.py, add the process_feedback method below the __init__ method:
def process_feedback(self, user_id, movie_id, reward):
key = f"{user_id}_{movie_id}"
if key not in self.saved_log_probs:
return False
log_prob = self.saved_log_probs.pop(key)
# REINFORCE: policy gradient loss
policy_loss = -log_prob * reward
self.policy_optimizer.zero_grad()
policy_loss.backward()
self.policy_optimizer.step()
return True
What does this code do?
This is the REINFORCE policy gradient in four lines:
- Look up the stored log probability for this (user, movie) pair. If the recommendation was never served, return False.
- Compute the loss as -log_prob * reward. When reward is +1 (thumbs up), this increases the probability of that action. When reward is -1 (thumbs down), it decreases it.
- Backpropagate and update the policy network weights with Adam. The Two-Tower model stays frozen.
- Save engine.py.
Confused about the negative sign?
The negative in -log_prob * reward exists because optimizers minimize loss. Minimizing the negative log probability of good actions is equivalent to maximizing the probability of those actions. This is standard in policy gradient methods.
help me understand why the REINFORCE loss uses a negative sign
Create the recommendation retrieval and re-ranking pipeline
Now you need the method that actually serves recommendations. It retrieves the top-50 candidates using Two-Tower dot-product scores, then passes them through the PolicyNetwork to produce a probability distribution for re-ranking.
- In engine.py, add the get_recommendations method between __init__ and process_feedback:
def get_recommendations(self, user_id, top_n=10):
user_to_idx = self.checkpoint["user_to_idx"]
idx_to_item = self.checkpoint["idx_to_item"]
if user_id not in user_to_idx:
return []
user_idx = user_to_idx[user_id]
user_tensor = torch.tensor([user_idx], dtype=torch.long)
with torch.no_grad():
user_emb = self.model.get_user_embedding(user_tensor).squeeze(0)
# Two-Tower retrieval: dot product with all items
scores = user_emb @ self.item_embeddings.T
top_k_indices = torch.topk(scores, self.num_candidates).indices
# Policy re-ranking
candidate_features = []
for idx in top_k_indices:
item_emb = self.item_embeddings[idx]
score = scores[idx].unsqueeze(0)
feature = torch.cat([user_emb, item_emb, score])
candidate_features.append(feature)
candidate_features = torch.stack(candidate_features)
probs = self.policy(candidate_features)
# Sample top_n from the distribution
dist = Categorical(probs)
selected_indices = []
for _ in range(min(top_n, len(probs))):
action = dist.sample()
selected_indices.append(action.item())
# Store log probs for potential REINFORCE update
results = []
for rank, sel_idx in enumerate(selected_indices):
item_idx = top_k_indices[sel_idx].item()
original_id = idx_to_item[item_idx]
log_prob = torch.log(probs[sel_idx] + 1e-8)
results.append({
"movie_id": original_id,
"retrieval_score": round(scores[item_idx].item(), 3),
"rank": rank + 1,
})
key = f"{user_id}_{original_id}"
self.saved_log_probs[key] = log_prob
return results
What does this code do?
This method runs the full retrieve-then-rank pipeline in four stages:
- Compute the user embedding and score it against all precomputed item embeddings via dot product. This is the Two-Tower retrieval step.
- Take the top 50 candidates and build a feature vector for each: the user embedding (64 dims) + item embedding (64 dims) + retrieval score (1 dim) = 129 dims total.
- Pass all 50 candidate features through the PolicyNetwork to get a probability distribution. Sample from it using Categorical to produce the final ranking.
- Store each recommendation's log probability in saved_log_probs keyed by userId_movieId. This is the state the REINFORCE update needs when feedback arrives later.
- Save engine.py.
- Verify the file has no syntax errors by running:
python -c "from engine import RecommendationEngine; print('Engine loaded successfully')"
You should see Engine loaded successfully printed with no errors.
Seeing an import error?
Make sure your virtual environment is activated. The torch import will fail if the environment with your installed packages is not active.
Also confirm that model.py and two_tower_model.pt are in the same directory as engine.py. The engine imports from model.py and loads the checkpoint file on instantiation.
help me debug an import error when loading my RecommendationEngine
✔️ Awesome, I've got everything!
Great. Double check you have saved engine.py with all three methods: __init__, get_recommendations, and process_feedback.
ⓧ I'd like to double check the full code
Here is the complete engine.py file for comparison:
import torch
from torch.distributions import Categorical
from model import TwoTowerModel, PolicyNetwork
class RecommendationEngine:
def __init__(self, checkpoint_path, num_candidates=50):
self.checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
self.num_candidates = num_candidates
# Load Two-Tower model
self.model = TwoTowerModel(
self.checkpoint["num_users"],
self.checkpoint["num_items"],
embedding_dim=64,
hidden_dim=128,
)
self.model.load_state_dict(self.checkpoint["model_state_dict"])
self.model.eval()
# Precomputed item embeddings
self.item_embeddings = self.checkpoint["item_embeddings"]
# Policy network: input is user_emb (64) + item_emb (64) + score (1) = 129
self.policy = PolicyNetwork(input_dim=129, hidden_dim=64)
self.policy_optimizer = torch.optim.Adam(self.policy.parameters(), lr=0.001)
# Store log probs for REINFORCE
self.saved_log_probs = {}
def get_recommendations(self, user_id, top_n=10):
user_to_idx = self.checkpoint["user_to_idx"]
idx_to_item = self.checkpoint["idx_to_item"]
if user_id not in user_to_idx:
return []
user_idx = user_to_idx[user_id]
user_tensor = torch.tensor([user_idx], dtype=torch.long)
with torch.no_grad():
user_emb = self.model.get_user_embedding(user_tensor).squeeze(0)
# Two-Tower retrieval: dot product with all items
scores = user_emb @ self.item_embeddings.T
top_k_indices = torch.topk(scores, self.num_candidates).indices
# Policy re-ranking
candidate_features = []
for idx in top_k_indices:
item_emb = self.item_embeddings[idx]
score = scores[idx].unsqueeze(0)
feature = torch.cat([user_emb, item_emb, score])
candidate_features.append(feature)
candidate_features = torch.stack(candidate_features)
probs = self.policy(candidate_features)
# Sample top_n from the distribution
dist = Categorical(probs)
selected_indices = []
for _ in range(min(top_n, len(probs))):
action = dist.sample()
selected_indices.append(action.item())
# Store log probs for potential REINFORCE update
results = []
for rank, sel_idx in enumerate(selected_indices):
item_idx = top_k_indices[sel_idx].item()
original_id = idx_to_item[item_idx]
log_prob = torch.log(probs[sel_idx] + 1e-8)
results.append({
"movie_id": original_id,
"retrieval_score": round(scores[item_idx].item(), 3),
"rank": rank + 1,
})
key = f"{user_id}_{original_id}"
self.saved_log_probs[key] = log_prob
return results
def process_feedback(self, user_id, movie_id, reward):
key = f"{user_id}_{movie_id}"
if key not in self.saved_log_probs:
return False
log_prob = self.saved_log_probs.pop(key)
# REINFORCE: policy gradient loss
policy_loss = -log_prob * reward
self.policy_optimizer.zero_grad()
policy_loss.backward()
self.policy_optimizer.step()
return True
Your recommendation engine is ready. It retrieves candidates fast with Two-Tower dot products, then re-ranks them with a policy that gets smarter from every piece of feedback. Next up, you will wrap this engine in a FastAPI service with REST endpoints.
Build the FastAPI Service
Your recommendation engine can retrieve candidates and re-rank them with a learned policy. Now you need a way for users to actually interact with it.
In this step, you will wrap the RecommendationEngine in a FastAPI service with endpoints for getting recommendations, submitting feedback, and checking health. By the end, you will see the policy update live as you send feedback signals.
In this step, get ready to:
- Create a FastAPI app with Pydantic request/response schemas.
- Implement recommendation, feedback, and health endpoints.
- Test the full loop locally and observe the REINFORCE policy learning from feedback.
Create the FastAPI app with Pydantic schemas
The API layer is the bridge between your ML models and the outside world. You will define typed request and response schemas using Pydantic so that FastAPI can validate inputs and auto-generate documentation.
- Create a new file called main.py in your project root.
- Add the imports, app initialization, and Pydantic schemas by pasting this code:
import pandas as pd
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from engine import RecommendationEngine
app = FastAPI(title="Two-Tower + REINFORCE Recommendation API")
# Global variables populated on startup
engine = None
movies_df = None
# Request schemas define what the client sends
class RecommendRequest(BaseModel):
user_id: int
top_n: int = 10
class FeedbackRequest(BaseModel):
user_id: int
movie_id: int
reward: float
# Response schemas define what the API returns
class RecommendationResponse(BaseModel):
user_id: int
recommendations: list[dict]
class FeedbackResponse(BaseModel):
success: bool
message: str
What does this code do?
The FastAPI() constructor creates your web application. The title parameter appears in the auto-generated docs at /docs.
Each Pydantic class defines a strict contract. RecommendRequest expects a user_id integer and an optional top_n count (defaults to 10). FeedbackRequest expects the user, the movie they are reacting to, and a reward signal (+1 or -1).
Next, add the startup event handler that loads your trained model and movie metadata when the server boots. Paste this code below the Pydantic schemas:
- Add the startup event handler below the FeedbackResponse class:
@app.on_event("startup")
def startup():
global engine, movies_df
# Load the trained Two-Tower model and policy network
engine = RecommendationEngine("two_tower_model.pt", num_candidates=50)
# Load movie titles for human-readable responses
movies_df = pd.read_csv("ml-latest-small/movies.csv")
What does this code do?
The @app.on_event("startup") decorator tells FastAPI to run this function once when the server starts. It loads your trained checkpoint (two_tower_model.pt) and the movie metadata CSV into memory. Every request after that uses these pre-loaded objects, so there is no per-request loading delay.
- Save main.py.
Implement the API endpoints
Your service needs three endpoints. A health check confirms the service is running. The recommend endpoint retrieves and re-ranks movies. The feedback endpoint triggers a REINFORCE policy update.
- Add the /health and /recommend endpoints below the startup function in main.py:
@app.get("/health")
def health():
return {"status": "ok"}
@app.post("/recommend", response_model=RecommendationResponse)
def recommend(request: RecommendRequest):
results = engine.get_recommendations(request.user_id, request.top_n)
if not results:
raise HTTPException(status_code=404, detail="User not found")
# Attach human-readable movie titles to each result
for r in results:
movie_row = movies_df[movies_df["movieId"] == r["movie_id"]]
r["title"] = movie_row["title"].values[0] if len(movie_row) > 0 else "Unknown"
return RecommendationResponse(user_id=request.user_id, recommendations=results)
What does this code do?
The /health endpoint returns a simple JSON object. Load balancers and container orchestrators use this to check if the service is alive.
The /recommend endpoint calls get_recommendations() from your RecommendationEngine. It retrieves the top-K candidates via Two-Tower dot product, re-ranks them through the policy network, then attaches movie titles from the CSV before returning.
- Add the /feedback endpoint below the recommend function:
@app.post("/feedback", response_model=FeedbackResponse)
def feedback(request: FeedbackRequest):
success = engine.process_feedback(request.user_id, request.movie_id, request.reward)
if success:
return FeedbackResponse(success=True, message="Policy updated with feedback")
return FeedbackResponse(success=False, message="No stored log probability for this recommendation")
What does this code do?
When a user sends a thumbs up (reward: 1.0) or thumbs down (reward: -1.0), this endpoint calls process_feedback() on the engine. That triggers the REINFORCE gradient update, so the policy learns to rank that movie higher (or lower) for this user next time.
- Save main.py.
✔️ Awesome, I've got everything!
Great. Double check you have saved main.py with all three endpoints.
ⓧ I'd like to double check the full code
import pandas as pd
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from engine import RecommendationEngine
app = FastAPI(title="Two-Tower + REINFORCE Recommendation API")
engine = None
movies_df = None
class RecommendRequest(BaseModel):
user_id: int
top_n: int = 10
class FeedbackRequest(BaseModel):
user_id: int
movie_id: int
reward: float
class RecommendationResponse(BaseModel):
user_id: int
recommendations: list[dict]
class FeedbackResponse(BaseModel):
success: bool
message: str
@app.on_event("startup")
def startup():
global engine, movies_df
engine = RecommendationEngine("two_tower_model.pt", num_candidates=50)
movies_df = pd.read_csv("ml-latest-small/movies.csv")
@app.get("/health")
def health():
return {"status": "ok"}
@app.post("/recommend", response_model=RecommendationResponse)
def recommend(request: RecommendRequest):
results = engine.get_recommendations(request.user_id, request.top_n)
if not results:
raise HTTPException(status_code=404, detail="User not found")
for r in results:
movie_row = movies_df[movies_df["movieId"] == r["movie_id"]]
r["title"] = movie_row["title"].values[0] if len(movie_row) > 0 else "Unknown"
return RecommendationResponse(user_id=request.user_id, recommendations=results)
@app.post("/feedback", response_model=FeedbackResponse)
def feedback(request: FeedbackRequest):
success = engine.process_feedback(request.user_id, request.movie_id, request.reward)
if success:
return FeedbackResponse(success=True, message="Policy updated with feedback")
return FeedbackResponse(success=False, message="No stored log probability for this recommendation")
Test locally and observe the policy learning
The real test of a REINFORCE-based system is whether the ranking changes after feedback. You will get recommendations, tell the system you liked a specific movie, then get recommendations again and watch the re-ranking shift.
- Start the FastAPI development server by running this command:
fastapi dev main.py
What does this command do?
The fastapi dev command starts a development server with auto-reload on port 8000. It watches for file changes so you do not need to restart manually.
- Wait for the server to display a message indicating it is running.
- Open a second terminal window. Test the health endpoint by running:
curl http://localhost:8000/health
You should see {"status":"ok"} confirming the server is alive and the model loaded successfully.
- Get recommendations for user 1 by running:
curl -X POST http://localhost:8000/recommend -H "Content-Type: application/json" -d '{"user_id": 1, "top_n": 5}'
You should see a JSON response with 5 movie recommendations, each containing a movie_id, title, retrieval_score, and rank. Note the order of titles.
- Pick one of the movie_id values from the response. Submit positive feedback for it by running (replace MOVIE_ID with the actual number):
curl -X POST http://localhost:8000/feedback -H "Content-Type: application/json" -d '{"user_id": 1, "movie_id": [[MOVIE_ID="a movie_id from the recommend response"]], "reward": 1.0}'
You should see {"success":true,"message":"Policy updated with feedback"} confirming the REINFORCE gradient was applied.
- Get recommendations again to observe the re-ranking shift:
curl -X POST http://localhost:8000/recommend -H "Content-Type: application/json" -d '{"user_id": 1, "top_n": 5}'
Compare the ranking order to your first request. The movie you gave positive feedback on should now appear at a different rank, proving the policy learned from your signal.
Server not starting or seeing import errors?
- Make sure you are running the command from your project root directory (the folder containing main.py, engine.py, and model.py).
- Confirm your virtual environment is activated. The fastapi command comes from the fastapi[standard] package.
- Check that two_tower_model.pt and ml-latest-small/movies.csv exist in the project root.
Help me debug my FastAPI server startup error.
Your recommendation system is now a running API that retrieves, re-ranks, and learns from feedback. Next up, you will containerize it and deploy to AWS so anyone can access it over the internet.
Containerize and Deploy to AWS ECS Fargate
Your recommendation API works locally. Now it is time to package it into a Docker container and deploy it to the cloud so anyone can access it.
In this step, you will containerize the FastAPI service and deploy it to AWS ECS Fargate. Fargate is a serverless compute engine that runs containers without requiring you to manage servers.
In this step, get ready to:
- Write a Dockerfile and test the container locally.
- Push the image to AWS ECR and deploy an ECS Fargate service.
- Test the full deployed pipeline with the public IP.
Write a Dockerfile and test locally
A Dockerfile defines how to package your application, its dependencies, and model files into a single portable image. This image can then run identically on any machine with Docker installed.
- Create a file called Dockerfile in your project root (no file extension) with the following content:
FROM python:3.12-slim
WORKDIR /app
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
COPY model.py .
COPY engine.py .
COPY main.py .
COPY two_tower_model.pt .
COPY ml-latest-small/ ./ml-latest-small/
EXPOSE 8000
CMD ["fastapi", "run", "main.py", "--host", "0.0.0.0", "--port", "8000"]
What does this Dockerfile do?
- FROM python:3.12-slim starts from a lightweight Python base image.
- COPY requirements.txt and RUN pip install installs dependencies first. Docker caches this layer so rebuilds are fast when only code changes.
- The remaining COPY commands bring in the source code, trained model checkpoint, and movie metadata.
- CMD runs the FastAPI production server on port 8000.
- Build the Docker image by running:
docker build -t rec-api .
- Test the container locally by running:
docker run -p 8000:8000 rec-api
- Open a new terminal and verify the container is serving by running:
curl http://localhost:8000/health
You should see {"status":"ok"} in the response.
Container exiting immediately?
- Make sure two_tower_model.pt is in the same directory as your Dockerfile. The container needs this file to start the engine.
- Check that the ml-latest-small/ folder with movies.csv is in your project root.
- Run docker logs $(docker ps -lq) to see the startup error.
Why is my Docker container crashing on startup?
- Stop the container with Ctrl+C in the terminal where it is running.
Push the image to ECR and deploy to ECS Fargate
AWS ECR is a managed container registry that stores your Docker images. ECS pulls images from ECR to run as tasks on Fargate.
- Create an ECR repository by running:
aws ecr create-repository --repository-name rec-api --region us-east-1
- Authenticate Docker to ECR by running:
aws ecr get-login-password --region us-east-1 | docker login --username AWS --password-stdin [[AWS_ACCOUNT_ID="your 12-digit AWS account ID"]].dkr.ecr.us-east-1.amazonaws.com
- Tag your local image for ECR by running:
docker tag rec-api:latest [[AWS_ACCOUNT_ID="your 12-digit AWS account ID"]].dkr.ecr.us-east-1.amazonaws.com/rec-api:latest
- Push the image to ECR by running:
docker push [[AWS_ACCOUNT_ID="your 12-digit AWS account ID"]].dkr.ecr.us-east-1.amazonaws.com/rec-api:latest
This upload may take a few minutes depending on your internet connection. The PyTorch dependency makes the image fairly large.
Where do I find my AWS account ID?
Run aws sts get-caller-identity --query Account --output text in your terminal. It prints your 12-digit account ID.
- Create the CloudWatch log group for container logs by running:
aws logs create-log-group --log-group-name /ecs/rec-api --region us-east-1
- Create a file called task-definition.json in your project root with the following content:
{
"family": "rec-api",
"networkMode": "awsvpc",
"requiresCompatibilities": ["FARGATE"],
"cpu": "512",
"memory": "1024",
"executionRoleArn": "arn:aws:iam::ACCOUNT_ID:role/ecsTaskExecutionRole",
"containerDefinitions": [
{
"name": "rec-api",
"image": "ACCOUNT_ID.dkr.ecr.REGION.amazonaws.com/rec-api:latest",
"portMappings": [
{ "containerPort": 8000, "protocol": "tcp" }
],
"logConfiguration": {
"logDriver": "awslogs",
"options": {
"awslogs-group": "/ecs/rec-api",
"awslogs-region": "us-east-1",
"awslogs-stream-prefix": "ecs"
}
},
"essential": true
}
]
}
What does this task definition configure?
- cpu: 512 allocates 0.5 vCPU and memory: 1024 allocates 1 GB RAM. This is the minimum Fargate config.
- networkMode: awsvpc gives the task its own network interface with a public IP.
- executionRoleArn grants Fargate permission to pull images from ECR and send logs to CloudWatch.
- logConfiguration routes container output to the /ecs/rec-api log group you created above.
- Replace both instances of ACCOUNT_ID and REGION in the file with your actual AWS account ID and region (us-east-1).
- Create an ECS cluster by running:
aws ecs create-cluster --cluster-name rec-api-cluster
- Register the task definition by running:
aws ecs register-task-definition --cli-input-json file://task-definition.json
- Create the ECS service by running the command below. Replace the subnet and security group with values from your default VPC:
aws ecs create-service \
--cluster rec-api-cluster \
--service-name rec-api-service \
--task-definition rec-api \
--desired-count 1 \
--launch-type FARGATE \
--network-configuration "awsvpcConfiguration={subnets=[[[SUBNET_ID="your public subnet ID"]]],securityGroups=[[[SG_ID="your security group ID"]]],assignPublicIp=ENABLED}"
How do I find my subnet and security group?
Run aws ec2 describe-subnets --query "Subnets[?MapPublicIpOnLaunch==\`true\`].SubnetId" --output text to find a public subnet.
Run aws ec2 describe-security-groups --query "SecurityGroups[?GroupName=='default'].GroupId" --output text to find your default security group. Make sure the security group allows inbound traffic on port 8000.
Wait 2-3 minutes for Fargate to pull the image and start the container.
Service stuck at 0 running tasks?
- Make sure ecsTaskExecutionRole exists in your account. If not, create it with the AmazonECSTaskExecutionRolePolicy managed policy attached.
- Check that your security group allows outbound traffic so Fargate can pull the image from ECR.
- View container logs in CloudWatch under the /ecs/rec-api log group for startup errors.
Why is my ECS Fargate task failing to start?
Retrieve the public IP and test the deployed service
Your service is now running on Fargate. Each task gets its own elastic network interface with a public IP address.
- Get the task ARN by running:
aws ecs list-tasks --cluster rec-api-cluster --service-name rec-api-service
- Describe the task to get the network interface details by running (replace the task ARN with the one from the previous output):
aws ecs describe-tasks \
--cluster rec-api-cluster \
--tasks [[TASK_ARN="your task ARN from the previous command"]] \
--query "tasks[0].attachments[0].details[?name=='networkInterfaceId'].value" \
--output text
- Get the public IP from the network interface by running:
aws ec2 describe-network-interfaces \
--network-interface-ids [[ENI_ID="the ENI ID from the previous command"]] \
--query "NetworkInterfaces[0].Association.PublicIp" \
--output text
This prints the public IP address of your running service.
- Test the health endpoint by running:
curl http://[[PUBLIC_IP="your public IP from the previous command"]]:8000/health
You should see {"status":"ok"} confirming the API is live.
- Get recommendations for a user by running:
curl -X POST http://[[PUBLIC_IP="your public IP"]]:8000/recommend \
-H "Content-Type: application/json" \
-d '{"user_id": 1, "top_n": 5}'
- Submit feedback by running (use a movie_id from the recommendations response above):
curl -X POST http://[[PUBLIC_IP="your public IP"]]:8000/feedback \
-H "Content-Type: application/json" \
-d '{"user_id": 1, "movie_id": 1, "reward": 1.0}'
You should see a response confirming the policy was updated. The REINFORCE policy network has just learned from real user feedback in production.
- Request recommendations again to observe the re-ranking shift:
curl -X POST http://[[PUBLIC_IP="your public IP"]]:8000/recommend \
-H "Content-Type: application/json" \
-d '{"user_id": 1, "top_n": 5}'
The order of movies may have shifted compared to the first request. The policy is learning from your feedback to improve rankings.
Getting connection refused on the public IP?
- Check that your security group allows inbound TCP on port 8000. Run aws ec2 authorize-security-group-ingress --group-id YOUR_SG_ID --protocol tcp --port 8000 --cidr 0.0.0.0/0 to open the port.
- Make sure the task is in RUNNING state. Run aws ecs describe-tasks --cluster rec-api-cluster --tasks YOUR_TASK_ARN --query "tasks[0].lastStatus" to check.
- Wait at least 2-3 minutes after service creation for the container to fully start.
Help me connect to my deployed Fargate service.
✔️ Awesome, I've got everything!
Great. Double check you have saved both Dockerfile and task-definition.json in your project root.
ⓧ I'd like to double check the full code
Here is the Dockerfile:
FROM python:3.12-slim
WORKDIR /app
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
COPY model.py .
COPY engine.py .
COPY main.py .
COPY two_tower_model.pt .
COPY ml-latest-small/ ./ml-latest-small/
EXPOSE 8000
CMD ["fastapi", "run", "main.py", "--host", "0.0.0.0", "--port", "8000"]
Here is task-definition.json:
{
"family": "rec-api",
"networkMode": "awsvpc",
"requiresCompatibilities": ["FARGATE"],
"cpu": "512",
"memory": "1024",
"executionRoleArn": "arn:aws:iam::ACCOUNT_ID:role/ecsTaskExecutionRole",
"containerDefinitions": [
{
"name": "rec-api",
"image": "ACCOUNT_ID.dkr.ecr.REGION.amazonaws.com/rec-api:latest",
"portMappings": [
{ "containerPort": 8000, "protocol": "tcp" }
],
"logConfiguration": {
"logDriver": "awslogs",
"options": {
"awslogs-group": "/ecs/rec-api",
"awslogs-region": "us-east-1",
"awslogs-stream-prefix": "ecs"
}
},
"essential": true
}
]
}
Your recommendation system is live on the internet. Next up, the Secret Mission will challenge you to add batch feedback processing for even more stable policy updates.
Secret mission
Add Batch Feedback Replay
Your current REINFORCE implementation updates the policy one feedback signal at a time. Production recommendation systems like YouTube and Spotify batch their feedback signals together before updating, producing more stable policy gradients with lower variance. Can you implement batch feedback replay with discounted rewards?
Clean Up Your Resources
Clean Up Your Resources
Decide whether to keep your resources running, pause them to come back later, or delete them entirely. Databricks Free Edition is free, but your ECS Fargate service charges approximately $0.01/hour while running.
Cost warning
Your ECS Fargate task (0.5 vCPU, 1 GB memory) costs approximately $0.01/hour. Leaving it running 24/7 would cost around $7/month.
If you are not actively using the API, choose Pause or Delete below to stop charges.
Resources you used:
- ECS Service (running Fargate tasks on port 8000).
- ECS Cluster.
- ECR Repository (rec-api).
- CloudWatch Log Group (/ecs/rec-api).
- Local project directory (includes mlruns/, model files, and source code).
- Databricks notebook (free, no charges).
Keep everything running
No action needed. Choose this if you are still actively building or want to demo the recommendation API.
- Your ECS service stays running and accessible via the public IP.
- Fargate charges continue at approximately $0.01/hour.
- Your Databricks notebook persists for free in your workspace.
- Your local mlruns/ directory and project files remain on disk.
Pause - I'll come back to this later
Stop Fargate charges by scaling your ECS service to zero tasks. Your cluster, ECR image, and local files remain intact so you can resume later.
- Scale the ECS service to zero by running:
aws ecs update-service --cluster rec-api --service rec-api --desired-count 0
This stops all running tasks and halts Fargate billing immediately.
- To resume later, run the same command with --desired-count 1.
Delete - I don't want to use this again
Remove all project resources and start fresh. Run each of the following commands in order.
- Delete the ECS service by running:
aws ecs delete-service --cluster rec-api --service rec-api --force
- Delete the ECS cluster by running:
aws ecs delete-cluster --cluster rec-api
- Delete the ECR repository and all its images by running:
aws ecr delete-repository --repository-name rec-api --force
- Delete the CloudWatch log group by running:
aws logs delete-log-group --log-group-name /ecs/rec-api
- Delete your local project directory on macOS/Linux by running:
cd ~/Desktop
rm -rf rec-api
- On Windows, run this instead:
cd ~\Desktop
Remove-Item -Recurse -Force rec-api
- Your Databricks notebook is free and can stay in your workspace. To delete it, open your Databricks workspace, right-click the notebook, and select Move to Trash.
Nice Work!
Nice Work!
Congratulations! You just built a production-grade movie recommendation system that spans data engineering, experiment tracking, deep learning, reinforcement learning, and cloud deployment.
You've learned how to:
- Build a medallion architecture ETL pipeline (Bronze, Silver, Gold) in Databricks and train a Two-Tower retrieval model with full MLflow experiment tracking for hyperparameters, metrics, and model artifacts.
- Implement a REINFORCE policy gradient that learns from real-time user feedback to re-rank recommendations, creating a system that improves over time.
- Containerize and deploy the full system as a FastAPI service on AWS ECS Fargate with /recommend, /feedback, and /health endpoints.
- Added batch feedback replay with discounted rewards and reward normalization for more stable policy gradient updates.
Ready to quiz yourself?