Two-Tower Retrieval

## %pip install -q pandas numpy torch matplotlib
 
from pathlib import Path
from urllib.request import urlopen
 
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
 
SEED = 7
np.random.seed(SEED)
_ = torch.manual_seed(SEED)

repo_root = next((path for path in [Path.cwd(), *Path.cwd().parents] if (path / ".git").exists()), Path.cwd())
root = repo_root / "work/notebook-data/movielens-100k"
root.mkdir(parents=True, exist_ok=True)
 
ratings_url = "https://files.grouplens.org/datasets/movielens/ml-100k/u.data"
items_url = "https://files.grouplens.org/datasets/movielens/ml-100k/u.item"
 
ratings_path = root / "u.data"
items_path = root / "u.item"
 
for url, path in [(ratings_url, ratings_path), (items_url, items_path)]:
    if not path.exists():
        path.write_bytes(urlopen(url).read())
 
ratings = pd.read_csv(
    ratings_path,
    sep="\t",
    names=["user_id", "item_id", "rating", "timestamp"],
    encoding="latin-1",
)
 
items = pd.read_csv(
    items_path,
    sep="|",
    header=None,
    encoding="latin-1",
    usecols=[0, 1],
    names=["item_id", "title"],
)
 
ratings = ratings.merge(items, on="item_id", how="left")
ratings.head()

positives = ratings[ratings["rating"] >= 4].copy()
positives = positives.sort_values(["user_id", "timestamp"])
 
heldout = positives.groupby("user_id").tail(1)
train = positives.drop(index=heldout.index)
 
user_ids = sorted(positives["user_id"].unique())
item_ids = sorted(positives["item_id"].unique())
user_to_idx = {u: i for i, u in enumerate(user_ids)}
item_to_idx = {i: j for j, i in enumerate(item_ids)}
 
train_user = torch.tensor(train["user_id"].map(user_to_idx).to_numpy(), dtype=torch.long)
train_item = torch.tensor(train["item_id"].map(item_to_idx).to_numpy(), dtype=torch.long)
test_pairs = list(
    zip(
        heldout["user_id"].map(user_to_idx).to_numpy(),
        heldout["item_id"].map(item_to_idx).to_numpy(),
    )
)
 
positives_by_user = {
    user_to_idx[u]: set(group["item_id"].map(item_to_idx).tolist())
    for u, group in positives.groupby("user_id")
}
 
num_users = len(user_ids)
num_items = len(item_ids)
print(f"train positives={len(train):,} heldout users={len(test_pairs):,}")

class TwoTower(nn.Module):
    def __init__(self, num_users: int, num_items: int, dim: int = 64):
        super().__init__()
        self.user = nn.Embedding(num_users, dim)
        self.item = nn.Embedding(num_items, dim)
        self.user_bias = nn.Embedding(num_users, 1)
        self.item_bias = nn.Embedding(num_items, 1)
 
    def score(self, users, items):
        dot = (self.user(users) * self.item(items)).sum(dim=-1)
        bias = self.user_bias(users).squeeze(-1) + self.item_bias(items).squeeze(-1)
        return dot + bias

device = "cuda" if torch.cuda.is_available() else "cpu"
model = TwoTower(num_users, num_items, dim=64).to(device)
opt = torch.optim.AdamW(model.parameters(), lr=3e-3, weight_decay=1e-4)
 
batch_size = 2048
steps = 2500
losses = []
 
train_user_device = train_user.to(device)
train_item_device = train_item.to(device)
 
for step in range(steps):
    idx = torch.randint(0, len(train_user_device), (batch_size,), device=device)
    u = train_user_device[idx]
    pos_i = train_item_device[idx]
    neg_i = torch.randint(0, num_items, (batch_size,), device=device)
 
    pos_score = model.score(u, pos_i)
    neg_score = model.score(u, neg_i)
    logits = torch.cat([pos_score, neg_score], dim=0)
    labels = torch.cat([torch.ones_like(pos_score), torch.zeros_like(neg_score)], dim=0)
 
    loss = F.binary_cross_entropy_with_logits(logits, labels)
    loss_value = loss.item()
    opt.zero_grad()
    loss.backward()
    opt.step()
    losses.append(loss_value)
 
    if step % 250 == 0:
        print(f"step={step:04d} loss={loss_value:.4f}")

def evaluate(model: TwoTower, test_pairs, positives_by_user, k: int = 10):
    user_emb = model.user.weight.detach()
    item_emb = model.item.weight.detach()
    user_bias = model.user_bias.weight.detach().squeeze(-1)
    item_bias = model.item_bias.weight.detach().squeeze(-1)
 
    recalls = []
    rr = []
    for u, target in test_pairs:
        scores = user_emb[u] @ item_emb.T + user_bias[u] + item_bias
        seen = positives_by_user[u] - {target}
        if seen:
            scores[list(seen)] = -1e9
        topk = scores.topk(k).indices.tolist()
        if target in topk:
            recalls.append(1.0)
            rr.append(1.0 / (topk.index(target) + 1))
        else:
            recalls.append(0.0)
            rr.append(0.0)
    return float(np.mean(recalls)), float(np.mean(rr))
 
 
recall_at_10, mrr_at_10 = evaluate(model.cpu(), test_pairs, positives_by_user, k=10)
print({"recall@10": recall_at_10, "mrr@10": mrr_at_10})

plt.figure(figsize=(8, 4))
plt.plot(losses)
plt.title("Two-tower retrieval training loss")
plt.xlabel("step")
plt.ylabel("BCE loss")
plt.show()

title_to_item = {row.title: item_to_idx[row.item_id] for row in items.itertuples() if row.item_id in item_to_idx}
item_lookup = {item_to_idx[row.item_id]: row.title for row in items.itertuples() if row.item_id in item_to_idx}
 
anchor = "Star Wars (1977)"
anchor_idx = title_to_item[anchor]
anchor_vec = model.item.weight.detach()[anchor_idx]
sims = model.item.weight.detach() @ anchor_vec
top = sims.topk(10).indices.tolist()
[item_lookup[i] for i in top]