Tool call · Input
/app/solution/two_tower_ranker.py
#!/usr/bin/env python3
"""Latency-efficient two-tower ranker for the Made-for-You shelf.
User tower: weighted pooling of interaction-history embeddings (no user-ID
lookup). Item tower: independently precomputable catalog vectors. The relevance
score is the dot product, i.e. request-time MIPS over stored item embeddings.
"""
from __future__ import annotations
import argparse
import json
import math
import os
import sys
import time
from dataclasses import dataclass
from pathlib import Path
import numpy as np
import pandas as pd
import pyarrow.parquet as pq
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset
sys.path.insert(0, "/task")
from ranking_evaluation import score_rankings # noqa: E402
DATA = Path("/task/data")
SOLUTION = Path("/app/solution")
CKPT_DIR = SOLUTION / "checkpoints"
TRAIN_CUTOFF = 24_786_800
LISTEN_STRONG = 80
HALF_LIFE_DAYS = 28.0
DAY = 86_400.0
PAD, LISTEN, ORG_STRONG, LIKE, DISLIKE = 0, 1, 2, 3, 4
def log(msg: str) -> None:
print(msg, flush=True)
def set_seeds(seed: int) -> None:
np.random.seed(seed)
torch.manual_seed(seed)
def load_catalog() -> tuple[np.ndarray, np.ndarray, np.ndarray]:
cat = pd.read_parquet(DATA / "candidate_catalog.parquet")
item_ids = cat.item_id.to_numpy(np.int64)
pos = cat.positive_listens.to_numpy(np.float32)
org = cat.organic_positive_listens.to_numpy(np.float32)
rank = cat.popularity_rank.to_numpy(np.float32)
n = len(item_ids)
feats = np.zeros((n + 1, 4), dtype=np.float32)
feats[1:, 0] = np.log1p(pos)
feats[1:, 1] = np.log1p(org)
feats[1:, 2] = np.log(np.clip(rank, 1.0, None))
feats[1:, 3] = org / np.clip(pos, 1.0, None)
# Standardize feature columns (pad row stays 0).
for c in range(4):
col = feats[1:, c]
mu, sd = float(col.mean()), float(col.std() + 1e-6)
feats[1:, c] = (col - mu) / sd
pop = np.zeros(n + 1, dtype=np.float32)
pop[1:] = (np.log1p(pos) - np.log1p(pos).mean()) / (np.log1p(pos).std() + 1e-6)
return item_ids, feats, pop
def make_mapper(item_ids: np.ndarray) -> np.ndarray:
mapper = np.full(int(item_ids.max()) + 1, -1, dtype=np.int32)
mapper[item_ids] = np.arange(1, len(item_ids) + 1, dtype=np.int32)
return mapper
def _map_items(raw: np.ndarray, mapper: np.ndarray) -> np.ndarray:
idx = np.full(len(raw), -1, dtype=np.int32)
valid = (raw >= 0) & (raw < len(mapper))
idx[valid] = mapper[raw[valid]]
return idx
def _append(buckets: dict[str, list[np.ndarray]], uid, tsv, idx, kind, st) -> None:
if len(uid) == 0:
return
buckets["uid"].append(np.asarray(uid, dtype=np.int32))
buckets["ts"].append(np.asarray(tsv, dtype=np.int32))
buckets["item"].append(np.asarray(idx, dtype=np.int32))
buckets["type"].append(np.asarray(kind, dtype=np.int8))
buckets["strength"].append(np.asarray(st, dtype=np.float32))
def load_events(splits: list[str], mapper: np.ndarray) -> dict[str, np.ndarray]:
buckets = {"uid": [], "ts": [], "item": [], "type": [], "strength": []}
for split in splits:
pf = pq.ParquetFile(DATA / f"{split}_listens.parquet")
for rg in range(pf.num_row_groups):
table = pf.read_row_group(
rg,
columns=["uid", "timestamp", "item_id", "is_organic", "played_ratio_pct"],
)
uid = table["uid"].to_numpy()
tsv = table["timestamp"].to_numpy()
iid = table["item_id"].to_numpy()
org = table["is_organic"].to_numpy()
ratio = table["played_ratio_pct"].to_numpy()
del table
keep = ratio >= 50
uid, tsv, iid, org, ratio = uid[keep], tsv[keep], iid[keep], org[keep], ratio[keep]
idx = _map_items(iid.astype(np.int64, copy=False), mapper)
keep = idx > 0
uid, tsv, idx, org, ratio = uid[keep], tsv[keep], idx[keep], org[keep], ratio[keep]
kind = np.full(len(uid), LISTEN, dtype=np.int8)
strong = (ratio >= LISTEN_STRONG) & (org == 1)
kind[strong] = ORG_STRONG
st = np.where(
strong,
2.5,
np.where(ratio >= LISTEN_STRONG, 1.4, np.where(org == 1, 1.1, 0.8)),
).astype(np.float32)
_append(buckets, uid, tsv, idx, kind, st)
likes = pd.read_parquet(DATA / f"{split}_likes.parquet")
uid, tsv, iid, org = (
likes.uid.to_numpy(),
likes.timestamp.to_numpy(),
likes.item_id.to_numpy(),
likes.is_organic.to_numpy(),
)
del likes
idx = _map_items(iid.astype(np.int64, copy=False), mapper)
keep = idx > 0
uid, tsv, idx, org = uid[keep], tsv[keep], idx[keep], org[keep]
_append(
buckets,
uid,
tsv,
idx,
np.full(len(uid), LIKE, dtype=np.int8),
np.where(org == 1, 5.5, 4.0).astype(np.float32),
)
dis = pd.read_parquet(DATA / f"{split}_dislikes.parquet")
uid, tsv, iid = dis.uid.to_numpy(), dis.timestamp.to_numpy(), dis.item_id.to_numpy()
del dis
idx = _map_items(iid.astype(np.int64, copy=False), mapper)
keep = idx > 0
uid, tsv, idx = uid[keep], tsv[keep], idx[keep]
_append(
buckets,
uid,
tsv,
idx,
np.full(len(uid), DISLIKE, dtype=np.int8),
np.full(len(uid), -2.5, dtype=np.float32),
)
uid = np.concatenate(buckets["uid"])
tsv = np.concatenate(buckets["ts"])
item = np.concatenate(buckets["item"])
kind = np.concatenate(buckets["type"])
st = np.concatenate(buckets["strength"])
order = np.lexsort((tsv, uid))
return {
"uid": uid[order],
"ts": tsv[order],
"item": item[order],
"type": kind[order],
"strength": st[order],
}
@dataclass
class UserTables:
uids: np.ndarray
uid_to_row: dict[int, int]
items: list[np.ndarray]
ts: list[np.ndarray]
types: list[np.ndarray]
strength: list[np.ndarray]
pos_item: list[np.ndarray]
pos_w: list[np.ndarray]
def build_user_tables(events: dict[str, np.ndarray]) -> UserTables:
uid = events["uid"]
uniq, starts = np.unique(uid, return_index=True)
ends = np.append(starts[1:], len(uid))
n = len(uniq)
items: list[np.ndarray] = [None] * n # type: ignore
ts: list[np.ndarray] = [None] * n # type: ignore
types: list[np.ndarray] = [None] * n # type: ignore
strength: list[np.ndarray] = [None] * n # type: ignore
pos_item: list[np.ndarray] = [None] * n # type: ignore
pos_w: list[np.ndarray] = [None] * n # type: ignore
ev_item, ev_ts, ev_type, ev_st = (
events["item"],
events["ts"],
events["type"],
events["strength"],
)
for i, (a, b) in enumerate(zip(starts, ends)):
it, tsv, tp, st = ev_item[a:b], ev_ts[a:b], ev_type[a:b], ev_st[a:b]
items[i], ts[i], types[i], strength[i] = it, tsv, tp, st
pos_mask = st > 0
if not np.any(pos_mask):
pos_item[i] = np.zeros(0, dtype=np.int32)
pos_w[i] = np.zeros(0, dtype=np.float32)
continue
it_p, tp_p = it[pos_mask], tp[pos_mask]
_, first_idx = np.unique(it_p, return_index=True)
first_idx.sort()
it_f, tp_f = it_p[first_idx], tp_p[first_idx]
w = np.full(len(it_f), 0.5, dtype=np.float32)
w[tp_f == LIKE] = 3.2
w[tp_f == ORG_STRONG] = 1.7
pos_item[i] = it_f.astype(np.int32, copy=False)
pos_w[i] = w
return UserTables(
uids=uniq.astype(np.int64),
uid_to_row={int(u): i for i, u in enumerate(uniq)},
items=items,
ts=ts,
types=types,
strength=strength,
pos_item=pos_item,
pos_w=pos_w,
)
def compute_idf(tables: UserTables, n_items: int) -> np.ndarray:
df = np.zeros(n_items, dtype=np.float32)
for arr in tables.pos_item:
if len(arr):
df[np.unique(arr)] += 1.0
idf = np.log((len(tables.uids) + 1.0) / (df + 1.0)) + 1.0
idf[0] = 0.0
return idf.astype(np.float32)
def recency_weights(ts: np.ndarray, now: int, half_life: float = HALF_LIFE_DAYS) -> np.ndarray:
age_days = np.clip(now - ts.astype(np.int64), 0, None).astype(np.float32) / DAY
return np.exp(-math.log(2.0) * age_days / half_life).astype(np.float32)
def aggregate_profile(
items: np.ndarray,
types: np.ndarray,
strength: np.ndarray,
ts: np.ndarray,
now: int,
idf: np.ndarray,
max_items: int,
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
if len(items) == 0:
z = np.zeros(0, dtype=np.int32)
return z, np.zeros(0, dtype=np.int8), np.zeros(0, dtype=np.float32)
rec = recency_weights(ts, now)
w = strength * rec * idf[items]
order = np.argsort(items, kind="mergesort")
it, wt, tp = items[order], w[order], types[order]
change = np.ones(len(it), dtype=bool)
if len(it) > 1:
change[1:] = it[1:] != it[:-1]
starts = np.flatnonzero(change)
wsum = np.add.reduceat(wt, starts)
cnt = np.add.reduceat(np.ones(len(it), dtype=np.float32), starts)
wsum = wsum * np.log1p(cnt)
tmax = np.maximum.reduceat(tp.astype(np.int16), starts).astype(np.int8)
uniq = it[starts]
if len(uniq) > max_items:
keep = np.argpartition(-wsum, max_items - 1)[:max_items]
# Keep high-weight items; sort by weight desc for stable padding.
keep = keep[np.argsort(-wsum[keep])]
uniq, tmax, wsum = uniq[keep], tmax[keep], wsum[keep]
else:
keep = np.argsort(-wsum)
uniq, tmax, wsum = uniq[keep], tmax[keep], wsum[keep]
return uniq.astype(np.int32), tmax, wsum.astype(np.float32)
@dataclass
class Profiles:
items: np.ndarray # [U, L]
types: np.ndarray
weights: np.ndarray
lengths: np.ndarray
row_of_target: np.ndarray # [n_targets] index into U or -1
def build_profiles(
tables: UserTables,
target_uids: np.ndarray,
idf: np.ndarray,
now: int,
max_items: int,
) -> Profiles:
n_u = len(tables.uids)
items = np.zeros((n_u, max_items), dtype=np.int32)
types = np.zeros((n_u, max_items), dtype=np.int8)
weights = np.zeros((n_u, max_items), dtype=np.float32)
lengths = np.zeros(n_u, dtype=np.int16)
for i in range(n_u):
it, tp, w = aggregate_profile(
tables.items[i],
tables.types[i],
tables.strength[i],
tables.ts[i],
now,
idf,
max_items,
)
L = len(it)
if L == 0:
continue
items[i, :L] = it
types[i, :L] = tp
weights[i, :L] = w
lengths[i] = L
row_of_target = np.full(len(target_uids), -1, dtype=np.int32)
for r, uid in enumerate(target_uids):
ur = tables.uid_to_row.get(int(uid))
if ur is not None:
row_of_target[r] = ur
log(
f"profiles: users={n_u} mean_len={float(lengths.mean()):.1f} "
f"targets_mapped={int((row_of_target >= 0).sum())}"
)
return Profiles(items, types, weights, lengths, row_of_target)
class LOODataset(Dataset):
def __init__(self, user_row: np.ndarray, pos: np.ndarray, pos_w: np.ndarray):
self.user_row = user_row
self.pos = pos
self.pos_w = pos_w
def __len__(self) -> int:
return int(self.pos.shape[0])
def __getitem__(self, i: int):
return (
int(self.user_row[i]),
int(self.pos[i]),
float(self.pos_w[i]),
)
def build_examples(
tables: UserTables, profiles: Profiles, max_pos_per_user: int, min_hist: int, rng: np.random.Generator
) -> LOODataset:
rows = []
pos = []
wts = []
for i in range(len(tables.uids)):
if int(profiles.lengths[i]) < min_hist:
continue
p_it = tables.pos_item[i]
p_w = tables.pos_w[i]
if len(p_it) == 0:
continue
n_take = min(len(p_it), max_pos_per_user)
if len(p_it) > n_take:
# Prefer recent/stronger first-time items (they are time-ordered).
# Mix last n_take*2 then sample by weight.
sl = slice(-min(len(p_it), n_take * 2), None)
cand_i = p_it[sl]
cand_w = p_w[sl]
pw = cand_w / cand_w.sum()
pick = rng.choice(len(cand_i), size=n_take, replace=False, p=pw)
cand_i, cand_w = cand_i[pick], cand_w[pick]
else:
cand_i, cand_w = p_it, p_w
for item, w in zip(cand_i.tolist(), cand_w.tolist()):
rows.append(i)
pos.append(item)
wts.append(w)
log(f"training examples: {len(rows)}")
return LOODataset(
np.asarray(rows, dtype=np.int32),
np.asarray(pos, dtype=np.int32),
np.asarray(wts, dtype=np.float32),
)
class TwoTower(nn.Module):
"""Separate user/item towers; score = dot(user_vec, item_vec).
The last coordinate is an explicit popularity channel so MIPS can recover
a popularity prior without a request-time bias term.
"""
def __init__(self, n_items: int, item_feats: torch.Tensor, pop: torch.Tensor, d: int = 80):
super().__init__()
self.d = d
self.emb = nn.Embedding(n_items, d, padding_idx=PAD)
nn.init.normal_(self.emb.weight, mean=0.0, std=0.01)
with torch.no_grad():
self.emb.weight[PAD].zero_()
self.register_buffer("item_feats", item_feats)
self.register_buffer("pop", pop)
# Item tower: residual MLP over ID embedding + popularity features.
self.item_mlp = nn.Sequential(
nn.Linear(d + 4, d),
nn.GELU(),
nn.Linear(d, d),
)
for m in self.item_mlp.modules():
if isinstance(m, nn.Linear):
nn.init.xavier_uniform_(m.weight, gain=0.3)
nn.init.zeros_(m.bias)
# User tower: residual MLP over pooled history embeddings + stats.
self.user_mlp = nn.Sequential(
nn.Linear(d + 4, d),
nn.GELU(),
nn.Linear(d, d),
)
for m in self.user_mlp.modules():
if isinstance(m, nn.Linear):
nn.init.xavier_uniform_(m.weight, gain=0.3)
nn.init.zeros_(m.bias)
self.pop_gate = nn.Parameter(torch.tensor(1.0))
self.log_tau = nn.Parameter(torch.log(torch.tensor(0.35)))
def item_vectors(self, idx: torch.Tensor) -> torch.Tensor:
e = self.emb(idx)
feats = self.item_feats[idx]
h = e + self.item_mlp(torch.cat([e, feats], dim=-1))
pop = self.pop[idx].unsqueeze(-1)
return torch.cat([h, pop], dim=-1)
def user_vectors(
self,
hist_idx: torch.Tensor,
hist_w: torch.Tensor,
hist_mask: torch.Tensor,
hist_type: torch.Tensor,
) -> torch.Tensor:
e = self.emb(hist_idx)
mask_f = hist_mask.float()
w = hist_w * mask_f
den = w.sum(dim=1, keepdim=True).clamp_min(1e-6)
pooled = (e * w.unsqueeze(-1)).sum(dim=1) / den
hlen = mask_f.sum(dim=1, keepdim=True).clamp_min(1.0)
frac_like = ((hist_type == LIKE) & hist_mask).float().sum(dim=1, keepdim=True) / hlen
mean_pop = (self.pop[hist_idx] * mask_f).sum(dim=1, keepdim=True) / hlen
stats = torch.cat(
[torch.log1p(hlen) / 5.0, frac_like, mean_pop, torch.log1p(den) / 5.0], dim=1
)
h = pooled + self.user_mlp(torch.cat([pooled, stats], dim=-1))
gate = F.softplus(self.pop_gate).expand(h.size(0), 1)
return torch.cat([h, gate], dim=-1)
def tau(self) -> torch.Tensor:
return self.log_tau.exp().clamp(0.05, 2.0)
def collate_batch(batch, profiles: Profiles, device: torch.device):
user_rows = torch.tensor([b[0] for b in batch], dtype=torch.long)
pos = torch.tensor([b[1] for b in batch], dtype=torch.long)
pos_w = torch.tensor([b[2] for b in batch], dtype=torch.float32)
hist_idx = torch.from_numpy(profiles.items[user_rows.numpy()]).long()
hist_w = torch.from_numpy(profiles.weights[user_rows.numpy()])
hist_type = torch.from_numpy(profiles.types[user_rows.numpy()]).long()
hist_len = torch.from_numpy(profiles.lengths[user_rows.numpy()].astype(np.int64))
L = hist_idx.size(1)
ar = torch.arange(L)
hist_mask = ar.unsqueeze(0) < hist_len.unsqueeze(1)
# Leave-one-out: drop the positive from the user profile.
hist_w = hist_w.masked_fill(hist_idx == pos.unsqueeze(1), 0.0)
hist_mask = hist_mask & (hist_idx != pos.unsqueeze(1)) & (hist_idx != 0)
return (
hist_idx.to(device),
hist_w.to(device),
hist_mask.to(device),
hist_type.to(device),
pos.to(device),
pos_w.to(device),
)
def info_nce(user: torch.Tensor, item: torch.Tensor, extra: torch.Tensor, pos_w: torch.Tensor, tau: torch.Tensor):
logits = user @ item.t() / tau
if extra is not None:
logits = torch.cat([logits, user @ extra.t() / tau], dim=1)
labels = torch.arange(user.size(0), device=user.device)
per = F.cross_entropy(logits, labels, reduction="none")
w = pos_w / pos_w.mean().clamp_min(1e-6)
loss = (per * w).mean()
inv = F.cross_entropy(item @ user.t() / tau, labels)
return loss + 0.25 * inv
def train_one_epoch(
model: TwoTower,
loader: DataLoader,
opt: torch.optim.Optimizer,
n_items: int,
n_extra: int,
device: torch.device,
) -> float:
model.train()
total, n = 0.0, 0
t0 = time.time()
for step, batch in enumerate(loader, 1):
hist_idx, hist_w, hist_mask, hist_type, pos, pos_w = batch
opt.zero_grad(set_to_none=True)
user = model.user_vectors(hist_idx, hist_w, hist_mask, hist_type)
item = model.item_vectors(pos)
extra = model.item_vectors(torch.randint(1, n_items, (n_extra,), device=device))
loss = info_nce(user, item, extra, pos_w, model.tau())
loss.backward()
nn.utils.clip_grad_norm_(model.parameters(), 2.0)
opt.step()
total += float(loss.detach())
n += 1
if step % 50 == 0:
log(
f" step {step:4d} loss={total / n:.4f} "
f"tau={float(model.tau().detach()):.3f} "
f"pop_gate={float(F.softplus(model.pop_gate).detach()):.3f} "
f"{time.time() - t0:.1f}s"
)
return total / max(n, 1)
@torch.inference_mode()
def rank_all(
model: TwoTower,
profiles: Profiles,
target_uids: np.ndarray,
catalog_ids: np.ndarray,
elig_mask: np.ndarray,
device: torch.device,
) -> np.ndarray:
model.eval()
n_cat = len(catalog_ids)
item_vecs = model.item_vectors(torch.arange(1, n_cat + 1, device=device))
n_t = len(target_uids)
user_vecs = torch.zeros((n_t, model.d + 1), device=device)
rows = profiles.row_of_target
valid = rows >= 0
if valid.any():
idx_np = np.flatnonzero(valid)
# Encode in chunks.
bs = 512
for s in range(0, len(idx_np), bs):
chunk = idx_np[s : s + bs]
ur = rows[chunk]
hist_idx = torch.from_numpy(profiles.items[ur]).long().to(device)
hist_w = torch.from_numpy(profiles.weights[ur]).to(device)
hist_type = torch.from_numpy(profiles.types[ur]).long().to(device)
hist_len = torch.from_numpy(profiles.lengths[ur].astype(np.int64))
hist_mask = (torch.arange(hist_idx.size(1)) < hist_len.unsqueeze(1)).to(device)
hist_mask = hist_mask & (hist_idx != 0)
user_vecs[torch.from_numpy(chunk).to(device)] = model.user_vectors(
hist_idx, hist_w, hist_mask, hist_type
)
# Targets with no history: popularity-only vector.
if (~valid).any():
gate = float(F.softplus(model.pop_gate).detach())
empty = torch.zeros(model.d + 1, device=device)
empty[-1] = gate
user_vecs[torch.from_numpy(np.flatnonzero(~valid)).to(device)] = empty
scores = user_vecs @ item_vecs.t()
scores = scores.masked_fill(~torch.from_numpy(elig_mask).to(device), -1e9)
_, topi = torch.topk(scores, k=100, dim=1)
return catalog_ids[topi.cpu().numpy()]
def load_eligibility(path: Path) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
with np.load(path) as z:
uids = z["uids"].astype(np.int64)
item_ids = z["item_ids"].astype(np.int64)
packed = z["packed_eligible"]
mask = np.unpackbits(packed, axis=1, bitorder="little")[:, : len(item_ids)].astype(bool)
return uids, item_ids, mask
def align_eligibility(target_uids: np.ndarray, elig_uids: np.ndarray, mask: np.ndarray) -> np.ndarray:
row = {int(u): i for i, u in enumerate(elig_uids)}
out = np.empty((len(target_uids), mask.shape[1]), dtype=bool)
for i, u in enumerate(target_uids):
out[i] = mask[row[int(u)]]
return out
def validation_graded(
target_set: set[int],
uid_row: dict[int, int],
item_col: dict[int, int],
elig_mask: np.ndarray,
) -> dict[int, dict[int, int]]:
listens = pd.read_parquet(DATA / "val_listens.parquet")
likes = pd.read_parquet(DATA / "val_likes.parquet")
strong = listens[
listens.uid.isin(target_set)
& (listens.is_organic == 1)
& (listens.played_ratio_pct >= LISTEN_STRONG)
][["uid", "item_id"]].drop_duplicates()
org_likes = likes[likes.uid.isin(target_set) & (likes.is_organic == 1)][
["uid", "item_id"]
].drop_duplicates()
del listens, likes
def eligible(u: int, i: int) -> bool:
r, c = uid_row.get(u), item_col.get(i)
return r is not None and c is not None and bool(elig_mask[r, c])
graded: dict[int, dict[int, int]] = {}
for u, i in strong.itertuples(index=False, name=None):
u, i = int(u), int(i)
if eligible(u, i):
graded.setdefault(u, {})[i] = 1
for u, i in org_likes.itertuples(index=False, name=None):
u, i = int(u), int(i)
if eligible(u, i):
graded.setdefault(u, {})[i] = 3
return graded
def write_ranking(path: Path, uids: np.ndarray, ranked: np.ndarray) -> None:
n, k = ranked.shape
df = pd.DataFrame(
{
"uid": np.repeat(uids.astype(np.int64), k),
"item_id": ranked.reshape(-1).astype(np.int64),
"rank": np.tile(np.arange(1, k + 1, dtype=np.int64), n),
}
)
path.parent.mkdir(parents=True, exist_ok=True)
df.to_parquet(path, index=False)
log(f"wrote {path} rows={len(df)}")
def pop_ranking(target_uids: np.ndarray, catalog_ids: np.ndarray, elig_mask: np.ndarray) -> np.ndarray:
n = len(target_uids)
ranked = np.zeros((n, 100), dtype=np.int64)
for i in range(n):
elig = np.flatnonzero(elig_mask[i])
ranked[i] = catalog_ids[elig[:100]]
return ranked
def score_ranked(uids: np.ndarray, ranked: np.ndarray, graded: dict[int, dict[int, int]]) -> dict:
mapping = {int(u): ranked[i].tolist() for i, u in enumerate(uids)}
return score_rankings(mapping, graded)
def make_optimizer(model: TwoTower, lr: float, wd: float) -> torch.optim.Optimizer:
emb, other = [], []
for name, p in model.named_parameters():
if not p.requires_grad:
continue
(emb if "emb.weight" in name else other).append(p)
return torch.optim.AdamW(
[
{"params": emb, "lr": lr, "weight_decay": 0.0},
{"params": other, "lr": lr, "weight_decay": wd},
]
)
def run(args: argparse.Namespace) -> None:
t_all = time.time()
os.makedirs(CKPT_DIR, exist_ok=True)
set_seeds(args.seed)
torch.set_num_threads(max(1, min(8, os.cpu_count() or 4)))
device = torch.device("cpu")
rng = np.random.default_rng(args.seed)
catalog_ids, item_feats_np, pop_np = load_catalog()
n_cat = len(catalog_ids)
n_items = n_cat + 1
mapper = make_mapper(catalog_ids)
target_uids = pd.read_parquet(DATA / "target_users.parquet").uid.to_numpy(np.int64)
target_set = set(map(int, target_uids))
log(f"catalog={n_cat} target_users={len(target_uids)}")
val_eu, val_ei, val_pm = load_eligibility(DATA / "validation_eligible_candidates.npz")
test_eu, test_ei, test_pm = load_eligibility(DATA / "test_eligible_candidates.npz")
assert np.array_equal(val_ei, catalog_ids) and np.array_equal(test_ei, catalog_ids)
val_mask = align_eligibility(target_uids, val_eu, val_pm)
test_mask = align_eligibility(target_uids, test_eu, test_pm)
uid_row = {int(u): i for i, u in enumerate(target_uids)}
item_col = {int(it): i for i, it in enumerate(catalog_ids)}
graded = validation_graded(target_set, uid_row, item_col, val_mask)
log(f"validation graded users={len(graded)} pairs={sum(len(v) for v in graded.values())}")
# Always persist a popularity ranking so a valid artifact exists early.
pop_val = pop_ranking(target_uids, catalog_ids, val_mask)
pop_test = pop_ranking(target_uids, catalog_ids, test_mask)
pop_scores = score_ranked(target_uids, pop_val, graded)
log(f"popularity baseline NDCG@10={pop_scores['quality_ndcg@10']:.6f}")
write_ranking(SOLUTION / "validation.parquet", target_uids, pop_val)
write_ranking(SOLUTION / "submission.parquet", target_uids, pop_test)
with open(SOLUTION / "val_score.json", "w") as f:
json.dump({"quality_ndcg@10": pop_scores["quality_ndcg@10"], "model": "popularity"}, f)
item_feats = torch.from_numpy(item_feats_np)
pop_t = torch.from_numpy(pop_np)
def fit(splits: list[str], epochs: int, ckpt_name: str, warmup: Path | None, now_override: int | None):
log(f"=== load events {splits} ===")
events = load_events(splits, mapper)
now = int(now_override) if now_override is not None else int(events["ts"].max()) + 1
log(f"events={len(events['uid'])} now={now}")
tables = build_user_tables(events)
del events
idf = compute_idf(tables, n_items)
profiles = build_profiles(tables, target_uids, idf, now, args.profile_len)
ds = build_examples(tables, profiles, args.max_pos_per_user, args.min_hist, rng)
loader = DataLoader(
ds,
batch_size=args.batch_size,
shuffle=True,
num_workers=0,
drop_last=True,
collate_fn=lambda batch: collate_batch(batch, profiles, device),
)
model = TwoTower(n_items, item_feats, pop_t, d=args.d).to(device)
if warmup is not None and warmup.exists():
log(f"warm-start {warmup}")
model.load_state_dict(torch.load(warmup, map_location="cpu"))
opt = make_optimizer(model, args.lr, args.weight_decay)
best_ndcg = -1.0
best_state = None
is_val_stage = "val" not in splits
for epoch in range(1, epochs + 1):
log(f"--- epoch {epoch}/{epochs} examples={len(ds)} ---")
loss = train_one_epoch(model, loader, opt, n_items, args.n_extra_neg, device)
log(f"epoch {epoch} mean_loss={loss:.4f}")
if is_val_stage:
ranked = rank_all(model, profiles, target_uids, catalog_ids, val_mask, device)
scores = score_ranked(target_uids, ranked, graded)
ndcg = float(scores["quality_ndcg@10"])
log(f"public val NDCG@10={ndcg:.6f}")
if ndcg > best_ndcg:
best_ndcg = ndcg
best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
torch.save(best_state, CKPT_DIR / ckpt_name)
write_ranking(SOLUTION / "validation.parquet", target_uids, ranked)
with open(SOLUTION / "val_score.json", "w") as f:
json.dump({"quality_ndcg@10": ndcg, "epoch": epoch, "model": "two_tower"}, f)
log(f"new best {ndcg:.6f}")
else:
torch.save(model.state_dict(), CKPT_DIR / ckpt_name)
if best_state is not None:
model.load_state_dict(best_state)
return model, profiles
if args.stage in ("val", "all"):
model, _ = fit(["train"], args.val_epochs, "best_val.pt", None, TRAIN_CUTOFF)
log("=== backup test ranking (val weights, train+val history) ===")
events_tv = load_events(["train", "val"], mapper)
now_tv = int(events_tv["ts"].max()) + 1
tables_tv = build_user_tables(events_tv)
del events_tv
idf_tv = compute_idf(tables_tv, n_items)
profiles_tv = build_profiles(tables_tv, target_uids, idf_tv, now_tv, args.profile_len)
ranked_test = rank_all(model, profiles_tv, target_uids, catalog_ids, test_mask, device)
write_ranking(SOLUTION / "submission.parquet", target_uids, ranked_test)
if args.stage in ("test", "all"):
warmup = CKPT_DIR / "best_val.pt"
events_tv = None
model, profiles_tv = fit(
["train", "val"],
args.test_epochs,
"final_test.pt",
warmup if warmup.exists() else None,
None,
)
ranked_test = rank_all(model, profiles_tv, target_uids, catalog_ids, test_mask, device)
write_ranking(SOLUTION / "submission.parquet", target_uids, ranked_test)
log(f"done in {time.time() - t_all:.1f}s")
def parse_args() -> argparse.Namespace:
p = argparse.ArgumentParser()
p.add_argument("--stage", choices=["val", "test", "all"], default="all")
p.add_argument("--seed", type=int, default=42)
p.add_argument("--d", type=int, default=80)
p.add_argument("--batch-size", type=int, default=1024)
p.add_argument("--lr", type=float, default=2e-3)
p.add_argument("--weight-decay", type=float, default=1e-4)
p.add_argument("--val-epochs", type=int, default=4)
p.add_argument("--test-epochs", type=int, default=3)
p.add_argument("--profile-len", type=int, default=256)
p.add_argument("--max-pos-per-user", type=int, default=64)
p.add_argument("--min-hist", type=int, default=5)
p.add_argument("--n-extra-neg", type=int, default=512)
return p.parse_args()
if __name__ == "__main__":
run(parse_args())