Tool call · Input
/app/solution/train.py
#!/usr/bin/env python3
"""Train a CPU PyTorch two-tower ranker and emit validation/test rankings."""
from __future__ import annotations
import json
import math
import sys
import time
from collections import defaultdict
from pathlib import Path
import numpy as np
import pandas as pd
import pyarrow.parquet as pq
import torch
import torch.nn.functional as F
from torch.optim import AdamW
SOLUTION = Path("/app/solution")
sys.path.insert(0, str(SOLUTION))
sys.path.insert(0, "/task")
from model import TwoTower # noqa: E402
from ranking_evaluation import ndcg_at_k, score_rankings # noqa: E402
DATA = Path("/task/data")
TRAIN_CUTOFF = 24_786_800
MAX_ITEM_ID = 10_000_000
H_HIST = 64
L_LIKES = 64
D_DISLIKES = 32
DIM = 64
HIDDEN = 128
BATCH = 512
N_EXTRA_NEG = 96
MAX_LIKE_TARGETS = 48
MAX_LISTEN_TARGETS = 28
MIN_HIST = 3
SEED = 42
LISTEN_STRONG = 80
QUALITY_K = 10
TOP_N = 100
N_USER_FEAT = 8
N_ITEM_FEAT = 9
def log(msg: str) -> None:
print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True)
def recency_bucket(age_days: np.ndarray) -> np.ndarray:
buckets = np.full(age_days.shape, 6, dtype=np.int64)
buckets[age_days < 90] = 5
buckets[age_days < 30] = 4
buckets[age_days < 14] = 3
buckets[age_days < 7] = 2
buckets[age_days < 3] = 1
buckets[age_days < 1] = 0
return buckets
def listen_action(organic: np.ndarray, played: np.ndarray) -> np.ndarray:
strong = played >= 80
weak = (played >= 50) & ~strong
org = organic.astype(bool)
act = np.zeros(len(played), dtype=np.int16)
act[strong & org] = 1
act[weak & org] = 2
act[strong & ~org] = 3
act[weak & ~org] = 4
return act
def unique_recent(items: np.ndarray, limit: int) -> list[int]:
seen: set[int] = set()
out: list[int] = []
for i in range(len(items) - 1, -1, -1):
item = int(items[i])
if item <= 0 or item in seen:
continue
seen.add(item)
out.append(i)
if len(out) >= limit:
break
return out
def pad_int(values: np.ndarray, limit: int, dtype=np.int64) -> tuple[np.ndarray, np.ndarray]:
out = np.zeros(limit, dtype=dtype)
mask = np.zeros(limit, dtype=bool)
n = min(len(values), limit)
if n:
out[:n] = values[:n]
mask[:n] = True
return out, mask
def pad_float(values: np.ndarray, limit: int) -> np.ndarray:
out = np.zeros(limit, dtype=np.float32)
n = min(len(values), limit)
if n:
out[:n] = values[:n]
return out
def make_item_lookup(catalog_ids: np.ndarray) -> np.ndarray:
lookup = np.zeros(MAX_ITEM_ID, dtype=np.int32)
lookup[catalog_ids.astype(np.int64)] = np.arange(1, len(catalog_ids) + 1, dtype=np.int32)
return lookup
def load_listens(path: Path, lookup: np.ndarray) -> dict[str, np.ndarray]:
log(f"loading {path.name}")
table = pq.read_table(
path,
columns=["uid", "timestamp", "item_id", "is_organic", "played_ratio_pct"],
)
item = table["item_id"].to_numpy()
mapped = lookup[item]
keep = mapped != 0
played = table["played_ratio_pct"].to_numpy()[keep]
organic = table["is_organic"].to_numpy()[keep]
keep_pos = played >= 50
uid = table["uid"].to_numpy()[keep][keep_pos].astype(np.int64, copy=False)
ts = table["timestamp"].to_numpy()[keep][keep_pos].astype(np.int64, copy=False)
idx = mapped[keep][keep_pos].astype(np.int32, copy=False)
organic = organic[keep_pos].astype(np.uint8, copy=False)
played = played[keep_pos].astype(np.int16, copy=False)
del table
action = listen_action(organic, played)
log(f" catalog listens >=50% kept {len(uid):,}")
return {
"uid": uid,
"ts": ts,
"idx": idx,
"action": action,
"organic": organic,
"played": played,
}
def load_explicit(path: Path, lookup: np.ndarray, action: int) -> dict[str, np.ndarray]:
df = pd.read_parquet(path, columns=["uid", "timestamp", "item_id", "is_organic"])
mapped = lookup[df.item_id.to_numpy()]
keep = mapped != 0
uid = df.uid.to_numpy()[keep].astype(np.int64, copy=False)
ts = df.timestamp.to_numpy()[keep].astype(np.int64, copy=False)
idx = mapped[keep].astype(np.int32, copy=False)
organic = df.is_organic.to_numpy()[keep].astype(np.uint8, copy=False)
return {
"uid": uid,
"ts": ts,
"idx": idx,
"action": np.full(len(uid), action, dtype=np.int16),
"organic": organic,
"played": np.full(len(uid), 100, dtype=np.int16),
}
def concat_events(parts: list[dict[str, np.ndarray]]) -> dict[str, np.ndarray]:
keys = ["uid", "ts", "idx", "action", "organic", "played"]
return {k: np.concatenate([p[k] for p in parts if len(p[k])]) for k in keys}
def group_users(events: dict[str, np.ndarray]) -> dict[int, dict[str, np.ndarray]]:
uid = events["uid"]
order = np.lexsort((events["action"], events["ts"], uid))
uid = uid[order]
ts = events["ts"][order]
idx = events["idx"][order]
action = events["action"][order]
organic = events["organic"][order]
played = events["played"][order]
uniq, starts, counts = np.unique(uid, return_index=True, return_counts=True)
users: dict[int, dict[str, np.ndarray]] = {}
for user, start, count in zip(uniq.tolist(), starts.tolist(), counts.tolist()):
sl = slice(start, start + count)
users[int(user)] = {
"ts": ts[sl],
"idx": idx[sl],
"action": action[sl],
"organic": organic[sl],
"played": played[sl],
}
return users
def user_features_from_counts(
n_events: int,
n_unique: int,
n_likes: int,
n_dislikes: int,
mean_organic: float,
age_last: float,
span: float,
) -> np.ndarray:
feat = np.array(
[
math.log1p(n_events),
math.log1p(n_unique),
math.log1p(n_likes),
math.log1p(n_dislikes),
mean_organic,
math.log1p(max(age_last, 0.0)),
math.log1p(max(span, 0.0)),
n_likes / (n_events + 1.0),
],
dtype=np.float32,
)
return feat
def pack_channels(
idx: np.ndarray,
action: np.ndarray,
ts: np.ndarray,
cutoff: int,
like_items: np.ndarray | None,
dislike_items: np.ndarray | None,
) -> dict[str, np.ndarray]:
age_days = np.clip((cutoff - ts).astype(np.float64) / 86400.0, 0.0, 4000.0)
buckets = recency_bucket(age_days)
age_w = np.exp(-age_days / 21.0).astype(np.float32)
hist_sel = unique_recent(idx, H_HIST)
hist_idx, hist_mask = pad_int(idx[hist_sel] if hist_sel else np.zeros(0, np.int32), H_HIST)
hist_act, _ = pad_int(action[hist_sel] if hist_sel else np.zeros(0, np.int16), H_HIST)
hist_rec, _ = pad_int(buckets[hist_sel] if hist_sel else np.zeros(0, np.int64), H_HIST)
hist_agew = pad_float(age_w[hist_sel] if hist_sel else np.zeros(0, np.float32), H_HIST)
if like_items is None:
like_pos = np.flatnonzero(action == 5)
like_sel = unique_recent(idx[like_pos], L_LIKES) if len(like_pos) else []
like_src = idx[like_pos][like_sel] if like_sel else np.zeros(0, np.int32)
else:
like_src = like_items[-L_LIKES:][::-1] if len(like_items) else np.zeros(0, np.int32)
like_idx, like_mask = pad_int(like_src, L_LIKES)
if dislike_items is None:
dis_pos = np.flatnonzero(action == 6)
dis_sel = unique_recent(idx[dis_pos], D_DISLIKES) if len(dis_pos) else []
dis_src = idx[dis_pos][dis_sel] if dis_sel else np.zeros(0, np.int32)
else:
dis_src = dislike_items[-D_DISLIKES:][::-1] if len(dislike_items) else np.zeros(0, np.int32)
dislike_idx, dislike_mask = pad_int(dis_src, D_DISLIKES)
n_events = int(len(idx))
n_unique = int(len(np.unique(idx))) if n_events else 0
n_likes = int((action == 5).sum()) if like_items is None else int(len(like_items))
n_dislikes = int((action == 6).sum()) if dislike_items is None else int(len(dislike_items))
mean_org = float(np.mean(action[np.isin(action, [1, 2])] > 0)) if n_events else 0.0
# organic share among listens
listen_m = np.isin(action, [1, 2, 3, 4])
mean_org = float(np.mean(np.isin(action[listen_m], [1, 2]))) if listen_m.any() else 0.0
age_last = float(age_days[-1]) if n_events else 30.0
span = float(age_days[0] - age_days[-1]) if n_events else 0.0
ufeat = user_features_from_counts(
n_events, n_unique, n_likes, n_dislikes, mean_org, age_last, span
)
return {
"hist_idx": hist_idx,
"hist_act": hist_act,
"hist_rec": hist_rec,
"hist_agew": hist_agew,
"hist_mask": hist_mask,
"like_idx": like_idx,
"like_mask": like_mask,
"dislike_idx": dislike_idx,
"dislike_mask": dislike_mask,
"user_feat": ufeat,
}
def build_examples(users: dict[int, dict[str, np.ndarray]]) -> dict[str, np.ndarray]:
bags: dict[str, list[np.ndarray]] = defaultdict(list)
n_kept = 0
for rec in users.values():
ts = rec["ts"]
idx = rec["idx"]
action = rec["action"]
like_targets = np.flatnonzero(action == 5)
listen_targets = np.flatnonzero(action == 1)
if len(like_targets) > MAX_LIKE_TARGETS:
like_targets = like_targets[-MAX_LIKE_TARGETS:]
if len(listen_targets) > MAX_LISTEN_TARGETS:
listen_targets = listen_targets[-MAX_LISTEN_TARGETS:]
targets = np.unique(np.concatenate([like_targets, listen_targets]))
for pos in targets.tolist():
if pos < MIN_HIST:
continue
hist = {k: rec[k][:pos] for k in rec}
packed = pack_channels(
hist["idx"], hist["action"], hist["ts"], int(ts[pos]), None, None
)
if not packed["hist_mask"].any() and not packed["like_mask"].any():
continue
packed["pos_idx"] = np.int64(idx[pos])
packed["pos_is_like"] = np.bool_(action[pos] == 5)
for key, value in packed.items():
bags[key].append(value)
n_kept += 1
log(f"training examples {n_kept:,}")
out: dict[str, np.ndarray] = {}
for key, values in bags.items():
if key in {"pos_idx", "pos_is_like"}:
out[key] = np.asarray(values)
else:
out[key] = np.stack(values, axis=0)
return out
def zscore_user_features(
examples: dict[str, np.ndarray],
) -> tuple[np.ndarray, np.ndarray]:
feat = examples["user_feat"]
mean = feat.mean(axis=0)
std = feat.std(axis=0).clip(1e-6)
examples["user_feat"] = ((feat - mean) / std).astype(np.float32)
return mean.astype(np.float32), std.astype(np.float32)
def apply_user_zscore(
feat: np.ndarray, mean: np.ndarray, std: np.ndarray
) -> np.ndarray:
return ((feat - mean) / std).astype(np.float32)
def build_item_feature_table(
catalog: pd.DataFrame,
events: dict[str, np.ndarray],
n_items: int,
) -> np.ndarray:
table = np.zeros((n_items, N_ITEM_FEAT), dtype=np.float32)
pos = catalog.positive_listens.to_numpy(dtype=np.float64)
org = catalog.organic_positive_listens.to_numpy(dtype=np.float64)
rank = catalog.popularity_rank.to_numpy(dtype=np.float64)
org_ratio = org / np.clip(pos, 1.0, None)
item_idx = events["idx"]
likes = np.bincount(item_idx[events["action"] == 5], minlength=n_items)
dislikes = np.bincount(item_idx[events["action"] == 6], minlength=n_items)
listens = np.bincount(item_idx[np.isin(events["action"], [1, 2, 3, 4])], minlength=n_items)
org_listen = np.bincount(item_idx[np.isin(events["action"], [1, 2])], minlength=n_items)
played_sum = np.bincount(
item_idx, weights=events["played"].astype(np.float64), minlength=n_items
)
played_cnt = np.bincount(item_idx, minlength=n_items)
mean_played = played_sum / np.clip(played_cnt, 1, None)
users_per_item = np.zeros(n_items, dtype=np.float32)
# unique users approximated by listen count compression
users_per_item[1:] = np.sqrt(np.clip(listens[1:], 0, None)).astype(np.float32)
raw = np.zeros((n_items - 1, N_ITEM_FEAT), dtype=np.float32)
raw[:, 0] = np.log1p(pos).astype(np.float32)
raw[:, 1] = np.log1p(org).astype(np.float32)
raw[:, 2] = org_ratio.astype(np.float32)
raw[:, 3] = (rank / len(catalog)).astype(np.float32)
raw[:, 4] = np.log1p(likes[1:]).astype(np.float32)
raw[:, 5] = np.log1p(dislikes[1:]).astype(np.float32)
raw[:, 6] = (mean_played[1:] / 100.0).astype(np.float32)
raw[:, 7] = np.log1p(listens[1:]).astype(np.float32)
raw[:, 8] = (org_listen[1:] / np.clip(listens[1:], 1, None)).astype(np.float32)
mean = raw.mean(axis=0)
std = raw.std(axis=0).clip(1e-6)
table[1:] = (raw - mean) / std
return table
def load_preference_maps(
path: Path, lookup: np.ndarray
) -> tuple[dict[int, np.ndarray], dict[int, np.ndarray]]:
df = pd.read_parquet(path)
mapped = lookup[df.item_id.to_numpy()]
uids = df.uid.to_numpy()
states = df.state.to_numpy()
likes: dict[int, list[int]] = defaultdict(list)
dislikes: dict[int, list[int]] = defaultdict(list)
for uid, item, state in zip(uids, mapped, states):
if item == 0:
continue
if state == "liked":
likes[int(uid)].append(int(item))
elif state == "disliked":
dislikes[int(uid)].append(int(item))
like_map = {u: np.asarray(v, dtype=np.int32) for u, v in likes.items()}
dislike_map = {u: np.asarray(v, dtype=np.int32) for u, v in dislikes.items()}
return like_map, dislike_map
def build_gold(
targets: set[int], eligibility_bits: np.ndarray, uid_row: dict[int, int], item_ids: np.ndarray
) -> dict[int, dict[int, int]]:
item_col = {int(item): i for i, item in enumerate(item_ids)}
listens = pd.read_parquet(DATA / "val_listens.parquet")
likes = pd.read_parquet(DATA / "val_likes.parquet")
strong = listens[
listens.uid.isin(targets)
& (listens.is_organic == 1)
& (listens.played_ratio_pct >= LISTEN_STRONG)
][["uid", "item_id"]].drop_duplicates()
organic_likes = likes[
likes.uid.isin(targets) & (likes.is_organic == 1)
][["uid", "item_id"]].drop_duplicates()
gold: dict[int, dict[int, int]] = defaultdict(dict)
def add(frame: pd.DataFrame, grade: int) -> None:
for uid, item in frame.itertuples(index=False, name=None):
uid_i, item_i = int(uid), int(item)
row = uid_row.get(uid_i)
col = item_col.get(item_i)
if row is None or col is None:
continue
if eligibility_bits[row, col]:
gold[uid_i][item_i] = grade
add(strong, 1)
add(organic_likes, 3)
return dict(gold)
def examples_to_tensors(examples: dict[str, np.ndarray]) -> dict[str, torch.Tensor]:
tensors: dict[str, torch.Tensor] = {}
for key, value in examples.items():
if value.dtype == np.bool_:
tensors[key] = torch.from_numpy(np.ascontiguousarray(value))
elif value.dtype == np.float32:
tensors[key] = torch.from_numpy(np.ascontiguousarray(value))
else:
tensors[key] = torch.from_numpy(np.ascontiguousarray(value.astype(np.int64)))
return tensors
def encode_item_matrix(
model: TwoTower, item_feat: torch.Tensor, n_items: int
) -> torch.Tensor:
idx = torch.arange(n_items, dtype=torch.long)
chunks = []
model.eval()
with torch.inference_mode():
for start in range(0, n_items, 4096):
sl = idx[start : start + 4096]
chunks.append(model.encode_items(sl, item_feat[sl]))
return torch.cat(chunks, dim=0)
def encode_user_batch(
model: TwoTower,
item_vecs: torch.Tensor,
batch: dict[str, torch.Tensor],
) -> torch.Tensor:
hist_vecs = item_vecs[batch["hist_idx"]]
like_vecs = item_vecs[batch["like_idx"]]
dislike_vecs = item_vecs[batch["dislike_idx"]]
return model.encode_users(
hist_vecs,
batch["hist_act"],
batch["hist_rec"],
batch["hist_agew"],
batch["hist_mask"],
like_vecs,
like_mask=batch["like_mask"],
dislike_vecs=dislike_vecs,
dislike_mask=batch["dislike_mask"],
user_features=batch["user_feat"],
)
def pack_inference_users(
uids: np.ndarray,
users: dict[int, dict[str, np.ndarray]],
like_map: dict[int, np.ndarray],
dislike_map: dict[int, np.ndarray],
cutoff: int,
feat_mean: np.ndarray,
feat_std: np.ndarray,
) -> dict[str, torch.Tensor]:
bags: dict[str, list[np.ndarray]] = defaultdict(list)
empty = {
"ts": np.zeros(0, np.int64),
"idx": np.zeros(0, np.int32),
"action": np.zeros(0, np.int16),
"organic": np.zeros(0, np.uint8),
"played": np.zeros(0, np.int16),
}
for uid in uids.tolist():
rec = users.get(int(uid), empty)
packed = pack_channels(
rec["idx"],
rec["action"],
rec["ts"],
cutoff,
like_map.get(int(uid)),
dislike_map.get(int(uid)),
)
packed["user_feat"] = apply_user_zscore(packed["user_feat"], feat_mean, feat_std)
for key, value in packed.items():
bags[key].append(value)
stacked = {k: np.stack(v, axis=0) for k, v in bags.items()}
return examples_to_tensors(stacked)
def rank_users(
user_vecs: np.ndarray,
item_vecs: np.ndarray,
item_ids: np.ndarray,
elig_bits: np.ndarray,
uids: np.ndarray,
) -> pd.DataFrame:
scores = user_vecs @ item_vecs.T
scores[:, 0] = -1e9
# item_vecs includes pad at row 0; catalog vectors are rows 1..
catalog_scores = scores[:, 1:]
catalog_scores = np.where(elig_bits, catalog_scores, -1e9)
catalog_scores = catalog_scores - np.arange(catalog_scores.shape[1], dtype=np.float32) * 1e-7
part = np.argpartition(-catalog_scores, TOP_N - 1, axis=1)[:, :TOP_N]
part_scores = np.take_along_axis(catalog_scores, part, axis=1)
order = np.argsort(-part_scores, axis=1)
top = np.take_along_axis(part, order, axis=1)
top_items = item_ids[top]
n_users = len(uids)
out_uid = np.repeat(uids.astype(np.int64), TOP_N)
out_item = top_items.reshape(-1).astype(np.int64)
out_rank = np.tile(np.arange(1, TOP_N + 1, dtype=np.int64), n_users)
return pd.DataFrame({"uid": out_uid, "item_id": out_item, "rank": out_rank})
def evaluate_ranking(
frame: pd.DataFrame, gold: dict[int, dict[int, int]]
) -> dict[str, float | int]:
ranked = {
int(uid): list(map(int, grp.item_id.to_numpy()))
for uid, grp in frame.sort_values(["uid", "rank"]).groupby("uid")
}
return score_rankings(ranked, gold)
def train_loop(
model: TwoTower,
examples: dict[str, torch.Tensor],
item_feat: torch.Tensor,
n_items: int,
seconds_budget: float,
max_epochs: int,
lr: float,
eval_fn,
) -> tuple[dict[str, torch.Tensor], float, int]:
opt = AdamW(model.parameters(), lr=lr, weight_decay=1e-5)
n = examples["pos_idx"].shape[0]
steps_per_epoch = math.ceil(n / BATCH)
best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
best_score = -1.0
best_steps = 0
t0 = time.time()
total_steps = 0
for epoch in range(1, max_epochs + 1):
model.train()
order = torch.randperm(n)
running = 0.0
seen = 0
for start in range(0, n, BATCH):
ix = order[start : start + BATCH]
pos_idx = examples["pos_idx"][ix]
extra = torch.randint(1, n_items, (ix.shape[0], N_EXTRA_NEG))
extra = torch.where(
extra.eq(pos_idx.unsqueeze(1)), (extra % (n_items - 1)) + 1, extra
)
hist_idx = examples["hist_idx"][ix]
like_idx = examples["like_idx"][ix]
dislike_idx = examples["dislike_idx"][ix]
hist_vecs = model.encode_items(hist_idx, item_feat[hist_idx])
like_vecs = model.encode_items(like_idx, item_feat[like_idx])
dislike_vecs = model.encode_items(dislike_idx, item_feat[dislike_idx])
users = model.encode_users(
hist_vecs,
examples["hist_act"][ix],
examples["hist_rec"][ix],
examples["hist_agew"][ix],
examples["hist_mask"][ix],
like_vecs,
examples["like_mask"][ix],
dislike_vecs,
examples["dislike_mask"][ix],
examples["user_feat"][ix],
)
pos_vecs = model.encode_items(pos_idx, item_feat[pos_idx])
neg_vecs = model.encode_items(extra, item_feat[extra])
temp = model.temperature()
logits_in = users @ pos_vecs.t() / temp
logits_ex = torch.bmm(neg_vecs, users.unsqueeze(-1)).squeeze(-1) / temp
logits = torch.cat([logits_in, logits_ex], dim=1)
labels = torch.arange(ix.shape[0])
w = torch.where(examples["pos_is_like"][ix], 2.5, 1.0)
loss_u = F.cross_entropy(logits, labels, reduction="none")
loss_i = F.cross_entropy(logits_in.t(), labels, reduction="none")
loss = (loss_u * w + 0.5 * loss_i).mean()
opt.zero_grad(set_to_none=True)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
opt.step()
running += float(loss.item()) * ix.shape[0]
seen += ix.shape[0]
total_steps += 1
if time.time() - t0 > seconds_budget:
break
elapsed = time.time() - t0
log(
f"epoch {epoch} loss {running / max(seen, 1):.4f} "
f"steps {total_steps} elapsed {elapsed:.0f}s"
)
score = eval_fn(model)
log(f" val NDCG@10 {score:.6f}")
if score >= best_score:
best_score = score
best_steps = total_steps
best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
log(" new best")
if time.time() - t0 > seconds_budget:
log("hit training time budget")
break
# keep going if improving; otherwise still finish remaining epochs in budget
_ = steps_per_epoch
return best_state, best_score, best_steps
def infer_rank_frame(
model: TwoTower,
item_feat: torch.Tensor,
n_items: int,
infer_pack: dict[str, torch.Tensor],
item_ids: np.ndarray,
elig_bits: np.ndarray,
uids: np.ndarray,
) -> pd.DataFrame:
model.eval()
with torch.inference_mode():
item_vecs = encode_item_matrix(model, item_feat, n_items)
user_chunks = []
n = infer_pack["hist_idx"].shape[0]
for start in range(0, n, 512):
sl = {k: v[start : start + 512] for k, v in infer_pack.items()}
user_chunks.append(encode_user_batch(model, item_vecs, sl))
user_vecs = torch.cat(user_chunks, dim=0)
item_np = item_vecs.cpu().numpy()
user_np = user_vecs.cpu().numpy()
return rank_users(user_np, item_np, item_ids, elig_bits, uids)
def write_readme(val_score: float) -> None:
text = f"""# Made-for-You two-tower ranker
## Architecture
CPU PyTorch two-tower retrieval model. The **item tower** maps a catalog item
id plus static/count features through an embedding, linear feature projection,
LayerNorm, and a residual GELU MLP to a 64-dimensional L2-normalized vector.
Item vectors are independent of the request user, so they can be computed
offline and stored in a MIPS index.
The **user tower** never looks up a user id. It builds a request-time vector
from interaction history:
- attention-pooled recent catalog events (action + recency embeddings)
- recency-weighted mean of those item vectors
- mean of active liked-item vectors
- mean of active disliked-item vectors
- history summary features (counts, organic share, recency)
A two-layer GELU MLP produces a 64-dimensional L2-normalized user vector.
The relevance score is exactly the **dot product** of the two vectors
(cosine similarity because both are L2-normalized).
## Training objective
In-batch sampled softmax (InfoNCE) plus 96 uniform catalog negatives.
Organic likes are up-weighted 2.5× versus organic ≥80% listens. A lighter
reverse (item→user) in-batch loss is added. History is cutoff-correct:
each positive is predicted from strictly earlier events. The validation
ranker is trained on `train_*` only. The hidden-test ranker is refit on
`train_* + val_*` and uses `test_preference_state` plus test eligibility.
## Public validation
graded NDCG@10 = **{val_score:.6f}**
## Submission command
```sh
python /app/solution/train.py
```
The script writes `/app/solution/validation.parquet` and
`/app/solution/submission.parquet`.
"""
(SOLUTION / "README.md").write_text(text)
def main() -> None:
t_all = time.time()
torch.manual_seed(SEED)
np.random.seed(SEED)
torch.set_num_threads(8)
catalog = pd.read_parquet(DATA / "candidate_catalog.parquet")
targets_df = pd.read_parquet(DATA / "target_users.parquet")
target_uids = targets_df.uid.to_numpy(dtype=np.int64)
target_set = set(map(int, target_uids))
catalog_ids = catalog.item_id.to_numpy(dtype=np.int64)
n_items = len(catalog_ids) + 1
lookup = make_item_lookup(catalog_ids)
val_elig = np.load(DATA / "validation_eligible_candidates.npz")
test_elig = np.load(DATA / "test_eligible_candidates.npz")
val_bits = np.unpackbits(val_elig["packed_eligible"], axis=1, bitorder="little")[
:, : len(catalog_ids)
].astype(bool)
test_bits = np.unpackbits(test_elig["packed_eligible"], axis=1, bitorder="little")[
:, : len(catalog_ids)
].astype(bool)
uid_row = {int(u): i for i, u in enumerate(val_elig["uids"])}
gold = build_gold(target_set, val_bits, uid_row, catalog_ids)
log(f"gold users {len(gold)}")
train_listens = load_listens(DATA / "train_listens.parquet", lookup)
train_likes = load_explicit(DATA / "train_likes.parquet", lookup, 5)
train_dislikes = load_explicit(DATA / "train_dislikes.parquet", lookup, 6)
train_events = concat_events([train_listens, train_likes, train_dislikes])
del train_listens, train_likes, train_dislikes
train_users = group_users(train_events)
log(f"train users with catalog events {len(train_users)}")
item_feat_np = build_item_feature_table(catalog, train_events, n_items)
item_feat = torch.from_numpy(item_feat_np)
examples_np = build_examples(train_users)
feat_mean, feat_std = zscore_user_features(examples_np)
examples = examples_to_tensors(examples_np)
del examples_np
val_like_map, val_dis_map = load_preference_maps(
DATA / "validation_preference_state.parquet", lookup
)
infer_val = pack_inference_users(
target_uids, train_users, val_like_map, val_dis_map, TRAIN_CUTOFF, feat_mean, feat_std
)
model = TwoTower(
n_items=n_items,
n_item_features=N_ITEM_FEAT,
n_user_features=N_USER_FEAT,
dim=DIM,
hidden=HIDDEN,
dropout=0.1,
)
log(f"params {sum(p.numel() for p in model.parameters()):,}")
def eval_model(m: TwoTower) -> float:
frame = infer_rank_frame(
m, item_feat, n_items, infer_val, catalog_ids, val_bits, target_uids
)
metrics = evaluate_ranking(frame, gold)
return float(metrics["quality_ndcg@10"])
remaining = 58 * 60 - (time.time() - t_all)
stage1_budget = max(8 * 60.0, min(20 * 60.0, remaining * 0.45))
log(f"stage-1 train budget {stage1_budget:.0f}s")
best_state, best_score, best_steps = train_loop(
model,
examples,
item_feat,
n_items,
seconds_budget=stage1_budget,
max_epochs=8,
lr=8e-4,
eval_fn=eval_model,
)
model.load_state_dict(best_state)
val_frame = infer_rank_frame(
model, item_feat, n_items, infer_val, catalog_ids, val_bits, target_uids
)
val_path = SOLUTION / "validation.parquet"
val_frame.to_parquet(val_path, index=False)
val_metrics = evaluate_ranking(val_frame, gold)
best_score = float(val_metrics["quality_ndcg@10"])
log(f"wrote {val_path} NDCG@10={best_score:.6f}")
torch.save(best_state, SOLUTION / "two_tower_val.pt")
write_readme(best_score)
# Refit on train + val for the hidden test.
remaining = 58 * 60 - (time.time() - t_all)
log(f"remaining before refit {remaining:.0f}s")
val_listens = load_listens(DATA / "val_listens.parquet", lookup)
val_likes = load_explicit(DATA / "val_likes.parquet", lookup, 5)
val_dislikes = load_explicit(DATA / "val_dislikes.parquet", lookup, 6)
full_events = concat_events(
[
train_events,
val_listens,
val_likes,
val_dislikes,
]
)
del val_listens, val_likes, val_dislikes, train_events
full_users = group_users(full_events)
item_feat_full_np = build_item_feature_table(catalog, full_events, n_items)
item_feat_full = torch.from_numpy(item_feat_full_np)
full_examples_np = build_examples(full_users)
full_mean, full_std = zscore_user_features(full_examples_np)
full_examples = examples_to_tensors(full_examples_np)
del full_examples_np
test_like_map, test_dis_map = load_preference_maps(
DATA / "test_preference_state.parquet", lookup
)
# validation window end ≈ 25393395; 30-minute gap then test
test_cutoff = 25_395_200
infer_test = pack_inference_users(
target_uids, full_users, test_like_map, test_dis_map, test_cutoff, full_mean, full_std
)
remaining = 57 * 60 - (time.time() - t_all)
if remaining > 120:
log(f"stage-2 refit budget {remaining:.0f}s")
model_full = TwoTower(
n_items=n_items,
n_item_features=N_ITEM_FEAT,
n_user_features=N_USER_FEAT,
dim=DIM,
hidden=HIDDEN,
dropout=0.1,
)
model_full.load_state_dict(best_state)
def no_eval(_m: TwoTower) -> float:
return best_score
# Continue training; keep last state (no hidden labels).
t_refit = time.time()
opt = AdamW(model_full.parameters(), lr=3e-4, weight_decay=1e-5)
n = full_examples["pos_idx"].shape[0]
total_steps = 0
max_epochs = 6
model_full.train()
last_state = {k: v.detach().cpu().clone() for k, v in model_full.state_dict().items()}
for epoch in range(1, max_epochs + 1):
order = torch.randperm(n)
running = 0.0
seen = 0
model_full.train()
for start in range(0, n, BATCH):
ix = order[start : start + BATCH]
pos_idx = full_examples["pos_idx"][ix]
extra = torch.randint(1, n_items, (ix.shape[0], N_EXTRA_NEG))
extra = torch.where(
extra.eq(pos_idx.unsqueeze(1)), (extra % (n_items - 1)) + 1, extra
)
hist_idx = full_examples["hist_idx"][ix]
like_idx = full_examples["like_idx"][ix]
dislike_idx = full_examples["dislike_idx"][ix]
hist_vecs = model_full.encode_items(hist_idx, item_feat_full[hist_idx])
like_vecs = model_full.encode_items(like_idx, item_feat_full[like_idx])
dislike_vecs = model_full.encode_items(
dislike_idx, item_feat_full[dislike_idx]
)
users = model_full.encode_users(
hist_vecs,
full_examples["hist_act"][ix],
full_examples["hist_rec"][ix],
full_examples["hist_agew"][ix],
full_examples["hist_mask"][ix],
like_vecs,
full_examples["like_mask"][ix],
dislike_vecs,
full_examples["dislike_mask"][ix],
full_examples["user_feat"][ix],
)
pos_vecs = model_full.encode_items(pos_idx, item_feat_full[pos_idx])
neg_vecs = model_full.encode_items(extra, item_feat_full[extra])
temp = model_full.temperature()
logits_in = users @ pos_vecs.t() / temp
logits_ex = torch.bmm(neg_vecs, users.unsqueeze(-1)).squeeze(-1) / temp
logits = torch.cat([logits_in, logits_ex], dim=1)
labels = torch.arange(ix.shape[0])
w = torch.where(full_examples["pos_is_like"][ix], 2.5, 1.0)
loss_u = F.cross_entropy(logits, labels, reduction="none")
loss_i = F.cross_entropy(logits_in.t(), labels, reduction="none")
loss = (loss_u * w + 0.5 * loss_i).mean()
opt.zero_grad(set_to_none=True)
loss.backward()
torch.nn.utils.clip_grad_norm_(model_full.parameters(), 1.0)
opt.step()
running += float(loss.item()) * ix.shape[0]
seen += ix.shape[0]
total_steps += 1
if time.time() - t_refit > remaining - 45:
break
log(f"refit epoch {epoch} loss {running / max(seen, 1):.4f}")
last_state = {
k: v.detach().cpu().clone() for k, v in model_full.state_dict().items()
}
if time.time() - t_refit > remaining - 45:
break
model_full.load_state_dict(last_state)
test_model = model_full
test_feat = item_feat_full
else:
log("skipping refit; using stage-1 weights with full history")
test_model = model
test_feat = item_feat_full
infer_test = pack_inference_users(
target_uids,
full_users,
test_like_map,
test_dis_map,
test_cutoff,
feat_mean,
feat_std,
)
sub = infer_rank_frame(
test_model, test_feat, n_items, infer_test, catalog_ids, test_bits, target_uids
)
sub_path = SOLUTION / "submission.parquet"
sub.to_parquet(sub_path, index=False)
torch.save(test_model.state_dict(), SOLUTION / "two_tower_test.pt")
log(f"wrote {sub_path} rows={len(sub)}")
write_readme(best_score)
log(f"done in {time.time() - t_all:.0f}s")
print(json.dumps({"validation_ndcg@10": best_score, "rows": int(len(sub))}, indent=2))
if __name__ == "__main__":
main()