Tool call · Input
/app/solution/two_tower_ranker.py
#!/usr/bin/env python3
"""CPU two-tower ranker for the Made-for-You catalog.
User vectors come from interaction history (no user-ID lookup). Item vectors
are produced by a separate item tower and are precomputable. The relevance
score is the dot product of the two vectors.
"""
from __future__ import annotations
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 as pa
import pyarrow.parquet as pq
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset
os.environ.setdefault("OMP_NUM_THREADS", "8")
os.environ.setdefault("MKL_NUM_THREADS", "8")
SOLUTION = Path("/app/solution")
DATA = Path("/task/data")
sys.path.insert(0, "/task")
from evaluate_public import validation_truth # noqa: E402
from ranking_evaluation import ( # noqa: E402
CandidateSet,
ndcg_at_k,
score_rankings,
)
TRAIN_CUTOFF = 24_786_800
GAP_SECONDS = 1_800
DAY = 86_400
WEEK = 7 * DAY
SEED = 42
EMB_DIM = 64
HIDDEN = 160
OUT_DIM = 64
HIST_LONG = 64
HIST_SHORT = 8
MAX_LIKES = 32
MAX_DISLIKES = 16
MAX_SKIPS = 8
PAD = 0
N_ITEM_FEAT = 9
N_USER_FEAT = 10
BATCH_SIZE = 384
LR = 1.2e-3
WEIGHT_DECAY = 1e-4
MAX_EPOCHS_VAL = 8
MAX_EPOCHS_FINAL = 7
LOG_EVERY = 50
MAX_POS_PER_USER_WINDOW = 10
MAX_SEQ_POS_PER_USER = 12
N_TRAIN_WINDOWS = 8
TAU_INIT = 20.0
DEVICE = torch.device("cpu")
def log(msg: str) -> None:
stamp = time.strftime("%H:%M:%S")
line = f"[{stamp}] {msg}"
print(line, flush=True)
with open(SOLUTION / "train.log", "a", encoding="utf-8") as handle:
handle.write(line + "\n")
def set_seeds(seed: int) -> None:
np.random.seed(seed)
torch.manual_seed(seed)
torch.set_num_threads(8)
torch.set_num_interop_threads(2)
def pad_1d(values: np.ndarray, length: int, dtype) -> np.ndarray:
out = np.zeros(length, dtype=dtype)
n = min(length, int(values.shape[0]))
if n:
out[:n] = values[:n]
return out
class MLP(nn.Module):
def __init__(self, in_dim: int, hidden: int, out_dim: int, dropout: float = 0.1):
super().__init__()
self.fc1 = nn.Linear(in_dim, hidden)
self.n1 = nn.LayerNorm(hidden)
self.fc2 = nn.Linear(hidden, hidden)
self.n2 = nn.LayerNorm(hidden)
self.fc3 = nn.Linear(hidden, out_dim)
self.skip = nn.Linear(in_dim, out_dim)
self.drop = nn.Dropout(dropout)
def forward(self, x: torch.Tensor) -> torch.Tensor:
h = self.drop(F.gelu(self.n1(self.fc1(x))))
h = self.drop(F.gelu(self.n2(self.fc2(h))))
return self.fc3(h) + self.skip(x)
class TwoTower(nn.Module):
"""History user tower + independently evaluable item tower."""
def __init__(self, n_catalog: int, n_item_feat: int, n_user_feat: int):
super().__init__()
self.n_catalog = n_catalog
vocab = n_catalog + 1
self.emb = nn.Embedding(vocab, EMB_DIM, padding_idx=PAD)
nn.init.normal_(self.emb.weight, mean=0.0, std=0.05)
self.emb.weight.data[PAD].zero_()
self.pos_emb = nn.Embedding(HIST_SHORT, EMB_DIM)
nn.init.normal_(self.pos_emb.weight, mean=0.0, std=0.02)
self.org_emb = nn.Embedding(2, EMB_DIM)
nn.init.normal_(self.org_emb.weight, mean=0.0, std=0.02)
self.item_feat_proj = nn.Linear(n_item_feat, EMB_DIM)
self.user_mlp = MLP(EMB_DIM * 6 + n_user_feat, HIDDEN, OUT_DIM, dropout=0.12)
self.item_mlp = MLP(EMB_DIM * 2, HIDDEN, OUT_DIM, dropout=0.05)
self.log_tau = nn.Parameter(torch.tensor(math.log(TAU_INIT), dtype=torch.float32))
self.register_buffer("item_features", torch.zeros(n_catalog, n_item_feat))
def set_item_features(self, feat: np.ndarray) -> None:
tensor = torch.from_numpy(feat.astype(np.float32))
self.item_features = tensor
def tau(self) -> torch.Tensor:
return self.log_tau.exp().clamp(1.0, 80.0)
def _finalize(self, raw: torch.Tensor) -> torch.Tensor:
taste = F.normalize(raw[:, :-1], dim=-1)
pop = torch.tanh(raw[:, -1:])
return torch.cat([taste, pop], dim=-1)
@staticmethod
def _mean(emb: torch.Tensor, mask: torch.Tensor, weights: torch.Tensor | None = None) -> torch.Tensor:
if weights is None:
weights = mask.float()
else:
weights = weights * mask.float()
denom = weights.sum(dim=1, keepdim=True).clamp_min(1e-6)
return (emb * weights.unsqueeze(-1)).sum(dim=1) / denom
def encode_users(
self,
long_ids: torch.Tensor,
long_w: torch.Tensor,
long_org: torch.Tensor,
short_ids: torch.Tensor,
like_ids: torch.Tensor,
dislike_ids: torch.Tensor,
skip_ids: torch.Tensor,
user_feat: torch.Tensor,
) -> torch.Tensor:
long_e = self.emb(long_ids) + self.org_emb(long_org.clamp(0, 1))
long_mask = long_ids.ne(PAD)
long_vec = self._mean(long_e, long_mask, long_w)
short_e = self.emb(short_ids) + self.pos_emb.weight.unsqueeze(0)
short_vec = self._mean(short_e, short_ids.ne(PAD))
like_vec = self._mean(self.emb(like_ids), like_ids.ne(PAD))
dislike_vec = self._mean(self.emb(dislike_ids), dislike_ids.ne(PAD))
skip_vec = self._mean(self.emb(skip_ids), skip_ids.ne(PAD))
rec_mask = long_mask & long_org.eq(0)
rec_vec = self._mean(long_e, rec_mask, long_w)
x = torch.cat(
[long_vec, short_vec, like_vec, dislike_vec, skip_vec, rec_vec, user_feat],
dim=-1,
)
return self._finalize(self.user_mlp(x))
def encode_items(self, catalog_idx: torch.Tensor) -> torch.Tensor:
emb = self.emb(catalog_idx + 1)
feat = self.item_feat_proj(self.item_features[catalog_idx])
raw = self.item_mlp(torch.cat([emb, feat], dim=-1))
return self._finalize(raw)
def encode_all_items(self) -> torch.Tensor:
idx = torch.arange(self.n_catalog, device=self.emb.weight.device)
return self.encode_items(idx)
@dataclass
class UserArrays:
pos_item: np.ndarray
pos_ts: np.ndarray
pos_org: np.ndarray
skip_item: np.ndarray
skip_ts: np.ndarray
disc_item: np.ndarray
disc_ts: np.ndarray
disc_grade: np.ndarray
expl_ts: np.ndarray
expl_item: np.ndarray
expl_kind: np.ndarray
uniq_item: np.ndarray
uniq_last_ts: np.ndarray
uniq_org: np.ndarray
def load_catalog() -> tuple[pd.DataFrame, dict[int, int], np.ndarray]:
catalog = pd.read_parquet(DATA / "candidate_catalog.parquet")
item_ids = catalog.item_id.to_numpy(dtype=np.int64)
mapping = {int(item): idx for idx, item in enumerate(item_ids)}
return catalog, mapping, item_ids
def _concat_parquets(paths: list[Path], columns: list[str]) -> pa.Table:
tables = [pq.read_table(path, columns=columns) for path in paths]
return pa.concat_tables(tables) if len(tables) > 1 else tables[0]
def _catalog_mask(item_ids: np.ndarray, catalog_ids: np.ndarray) -> np.ndarray:
lookup_size = int(max(int(item_ids.max()), int(catalog_ids.max()))) + 1
is_cat = np.zeros(lookup_size, dtype=np.uint8)
is_cat[catalog_ids] = 1
return is_cat[item_ids].astype(bool)
def load_listen_tables(listen_paths: list[Path], catalog_ids: np.ndarray) -> dict[str, np.ndarray]:
table = _concat_parquets(
listen_paths,
["uid", "timestamp", "item_id", "is_organic", "played_ratio_pct", "track_length_seconds"],
)
uid = table.column("uid").to_numpy().astype(np.int64)
ts = table.column("timestamp").to_numpy().astype(np.int32)
item = table.column("item_id").to_numpy().astype(np.int64)
organic = table.column("is_organic").to_numpy().astype(np.int8)
play = table.column("played_ratio_pct").to_numpy().astype(np.int16)
length = table.column("track_length_seconds").to_numpy().astype(np.int32)
mask = _catalog_mask(item, catalog_ids)
return {
"uid": uid[mask],
"ts": ts[mask],
"item": item[mask],
"organic": organic[mask],
"play": play[mask],
"length": length[mask],
}
def load_explicit(paths: list[Path], catalog_ids: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
if not paths:
return (
np.zeros(0, dtype=np.int64),
np.zeros(0, dtype=np.int32),
np.zeros(0, dtype=np.int64),
)
table = _concat_parquets(paths, ["uid", "timestamp", "item_id"])
uid = table.column("uid").to_numpy().astype(np.int64)
ts = table.column("timestamp").to_numpy().astype(np.int32)
item = table.column("item_id").to_numpy().astype(np.int64)
mask = _catalog_mask(item, catalog_ids)
return uid[mask], ts[mask], item[mask]
def map_items(items: np.ndarray, mapping: dict[int, int]) -> np.ndarray:
out = np.empty(items.shape[0], dtype=np.int32)
for i, item in enumerate(items):
out[i] = mapping[int(item)]
return out
def _group_splits(uids: np.ndarray) -> dict[int, tuple[int, int]]:
if uids.size == 0:
return {}
bounds = np.flatnonzero(uids[1:] != uids[:-1]) + 1
starts = np.concatenate([[0], bounds])
ends = np.concatenate([bounds, [uids.size]])
return {int(uids[s]): (int(s), int(e)) for s, e in zip(starts, ends)}
def _sort_uid_ts(uid: np.ndarray, ts: np.ndarray) -> np.ndarray:
return np.lexsort((ts, uid))
def compute_item_features(
listens: dict[str, np.ndarray],
mapping: dict[int, int],
n_items: int,
like_items: np.ndarray,
dislike_items: np.ndarray,
catalog: pd.DataFrame,
) -> tuple[np.ndarray, np.ndarray]:
item_idx = np.fromiter((mapping[int(x)] for x in listens["item"]), dtype=np.int32, count=len(listens["item"]))
play = listens["play"]
organic = listens["organic"]
length = listens["length"].astype(np.float64)
n_listens = np.bincount(item_idx, minlength=n_items).astype(np.float64)
n_pos = np.bincount(item_idx[play >= 50], minlength=n_items).astype(np.float64)
n_org = np.bincount(item_idx[(organic == 1) & (play >= 50)], minlength=n_items).astype(np.float64)
n_skip = np.bincount(item_idx[play < 20], minlength=n_items).astype(np.float64)
n_strong = np.bincount(item_idx[(organic == 1) & (play >= 80)], minlength=n_items).astype(np.float64)
sum_play = np.bincount(item_idx, weights=play.astype(np.float64), minlength=n_items)
sum_len = np.bincount(item_idx, weights=length, minlength=n_items)
like_idx = np.fromiter((mapping[int(x)] for x in like_items), dtype=np.int32, count=len(like_items)) if len(like_items) else np.zeros(0, dtype=np.int32)
dis_idx = np.fromiter((mapping[int(x)] for x in dislike_items), dtype=np.int32, count=len(dislike_items)) if len(dislike_items) else np.zeros(0, dtype=np.int32)
n_like = np.bincount(like_idx, minlength=n_items).astype(np.float64) if len(like_idx) else np.zeros(n_items)
n_dis = np.bincount(dis_idx, minlength=n_items).astype(np.float64) if len(dis_idx) else np.zeros(n_items)
pop = catalog.positive_listens.to_numpy(dtype=np.float64)
feat = np.column_stack(
[
np.log1p(n_pos),
np.log1p(n_org),
n_org / np.clip(n_pos, 1.0, None),
np.log1p(n_strong),
np.log1p(n_like),
np.log1p(n_dis),
n_skip / np.clip(n_listens, 1.0, None),
(sum_play / np.clip(n_listens, 1.0, None)) / 100.0,
np.log1p(pop),
]
).astype(np.float32)
mean = feat.mean(axis=0)
std = feat.std(axis=0)
std[std < 1e-6] = 1.0
feat = (feat - mean) / std
n_users_est = max(1, len(np.unique(listens["uid"])))
df = np.bincount(item_idx[play >= 50], minlength=n_items).astype(np.float64)
idf = np.log((n_users_est + 1.0) / (df + 1.0)).astype(np.float32)
idf = np.clip(idf, 0.05, 6.0)
return feat, idf
def _unique_last(items: np.ndarray, ts: np.ndarray, org: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
if items.size == 0:
return (
np.zeros(0, dtype=np.int32),
np.zeros(0, dtype=np.int32),
np.zeros(0, dtype=np.int8),
)
seen: dict[int, int] = {}
for i in range(items.shape[0]):
seen[int(items[i])] = i
order = np.fromiter((seen[k] for k in seen), dtype=np.int32, count=len(seen))
order = order[np.argsort(ts[order], kind="stable")]
return items[order], ts[order], org[order]
def build_user_store(
listens: dict[str, np.ndarray],
mapping: dict[int, int],
likes: tuple[np.ndarray, np.ndarray, np.ndarray],
unlikes: tuple[np.ndarray, np.ndarray, np.ndarray],
dislikes: tuple[np.ndarray, np.ndarray, np.ndarray],
undislikes: tuple[np.ndarray, np.ndarray, np.ndarray],
like_org: np.ndarray | None,
) -> dict[int, UserArrays]:
item_idx = np.fromiter((mapping[int(x)] for x in listens["item"]), dtype=np.int32, count=len(listens["item"]))
uid = listens["uid"]
ts = listens["ts"]
org = listens["organic"]
play = listens["play"]
order = _sort_uid_ts(uid, ts)
uid, ts, org, play, item_idx = uid[order], ts[order], org[order], play[order], item_idx[order]
pos_mask = play >= 50
skip_mask = play < 20
disc_mask = (org == 1) & (play >= 80)
store: dict[int, dict[str, list]] = {}
def ensure(u: int) -> dict[str, list]:
rec = store.get(u)
if rec is None:
rec = {
"pos_item": [],
"pos_ts": [],
"pos_org": [],
"skip_item": [],
"skip_ts": [],
"disc_item": [],
"disc_ts": [],
"disc_grade": [],
"expl_ts": [],
"expl_item": [],
"expl_kind": [],
}
store[u] = rec
return rec
splits = _group_splits(uid)
for u, (s, e) in splits.items():
rec = ensure(u)
sl = slice(s, e)
pm = pos_mask[sl]
rec["pos_item"] = item_idx[sl][pm]
rec["pos_ts"] = ts[sl][pm]
rec["pos_org"] = org[sl][pm]
sm = skip_mask[sl]
rec["skip_item"] = item_idx[sl][sm]
rec["skip_ts"] = ts[sl][sm]
dm = disc_mask[sl]
rec["disc_item"] = item_idx[sl][dm]
rec["disc_ts"] = ts[sl][dm]
rec["disc_grade"] = np.ones(int(dm.sum()), dtype=np.int8)
def add_expl(triple, kind: int, organic_filter: np.ndarray | None = None) -> None:
eu, et, ei = triple
if eu.size == 0:
return
mask = _catalog_mask(ei, np.array(list(mapping.keys()), dtype=np.int64)) if False else np.ones(len(ei), dtype=bool)
if organic_filter is not None:
mask = mask & (organic_filter == 1)
eu, et, ei = eu[mask], et[mask], ei[mask]
mapped = np.fromiter((mapping[int(x)] for x in ei), dtype=np.int32, count=len(ei))
for i in range(len(eu)):
rec = ensure(int(eu[i]))
rec["expl_ts"].append(int(et[i]))
rec["expl_item"].append(int(mapped[i]))
rec["expl_kind"].append(kind)
add_expl(likes, 0)
add_expl(unlikes, 1)
add_expl(dislikes, 2)
add_expl(undislikes, 3)
like_uid, like_ts, like_item = likes
if like_org is not None and like_uid.size:
org_mask = like_org == 1
lu, lt, li = like_uid[org_mask], like_ts[org_mask], like_item[org_mask]
cat_ok = np.fromiter((int(x) in mapping for x in li), dtype=bool, count=len(li))
lu, lt, li = lu[cat_ok], lt[cat_ok], li[cat_ok]
mapped = np.fromiter((mapping[int(x)] for x in li), dtype=np.int32, count=len(li))
for i in range(len(lu)):
rec = ensure(int(lu[i]))
if isinstance(rec["disc_item"], list):
rec["disc_item"].append(int(mapped[i]))
rec["disc_ts"].append(int(lt[i]))
rec["disc_grade"].append(3)
else:
rec["disc_item"] = np.concatenate([rec["disc_item"], np.array([mapped[i]], dtype=np.int32)])
rec["disc_ts"] = np.concatenate([rec["disc_ts"], np.array([lt[i]], dtype=np.int32)])
rec["disc_grade"] = np.concatenate([rec["disc_grade"], np.array([3], dtype=np.int8)])
users: dict[int, UserArrays] = {}
for u, rec in store.items():
def as_arr(key, dtype, default_list=True):
val = rec[key]
if isinstance(val, np.ndarray):
return val.astype(dtype, copy=False)
if not val:
return np.zeros(0, dtype=dtype)
return np.asarray(val, dtype=dtype)
pos_item = as_arr("pos_item", np.int32)
pos_ts = as_arr("pos_ts", np.int32)
pos_org = as_arr("pos_org", np.int8)
disc_item = as_arr("disc_item", np.int32)
disc_ts = as_arr("disc_ts", np.int32)
disc_grade = as_arr("disc_grade", np.int8)
if disc_item.size:
d_order = np.argsort(disc_ts, kind="stable")
disc_item, disc_ts, disc_grade = disc_item[d_order], disc_ts[d_order], disc_grade[d_order]
expl_ts = as_arr("expl_ts", np.int32)
expl_item = as_arr("expl_item", np.int32)
expl_kind = as_arr("expl_kind", np.int8)
if expl_ts.size:
e_order = np.argsort(expl_ts, kind="stable")
expl_ts, expl_item, expl_kind = expl_ts[e_order], expl_item[e_order], expl_kind[e_order]
uniq_item, uniq_last_ts, uniq_org = _unique_last(pos_item, pos_ts, pos_org)
users[u] = UserArrays(
pos_item=pos_item,
pos_ts=pos_ts,
pos_org=pos_org,
skip_item=as_arr("skip_item", np.int32),
skip_ts=as_arr("skip_ts", np.int32),
disc_item=disc_item,
disc_ts=disc_ts,
disc_grade=disc_grade,
expl_ts=expl_ts,
expl_item=expl_item,
expl_kind=expl_kind,
uniq_item=uniq_item,
uniq_last_ts=uniq_last_ts,
uniq_org=uniq_org,
)
return users
def explicit_state(ud: UserArrays, cutoff: int) -> tuple[np.ndarray, np.ndarray]:
liked: set[int] = set()
disliked: set[int] = set()
ts = ud.expl_ts
for i in range(ts.shape[0]):
if ts[i] >= cutoff:
break
item = int(ud.expl_item[i])
kind = int(ud.expl_kind[i])
if kind == 0:
liked.add(item)
elif kind == 1:
liked.discard(item)
elif kind == 2:
disliked.add(item)
else:
disliked.discard(item)
like_arr = np.fromiter(liked, dtype=np.int32) if liked else np.zeros(0, dtype=np.int32)
dis_arr = np.fromiter(disliked, dtype=np.int32) if disliked else np.zeros(0, dtype=np.int32)
return like_arr, dis_arr
def pack_user_inputs(
ud: UserArrays,
cutoff: int,
idf: np.ndarray,
like_override: np.ndarray | None = None,
dislike_override: np.ndarray | None = None,
) -> dict[str, np.ndarray]:
if ud.uniq_item.size:
keep = ud.uniq_last_ts < cutoff
items = ud.uniq_item[keep]
ts = ud.uniq_last_ts[keep]
org = ud.uniq_org[keep]
else:
items = ud.uniq_item
ts = ud.uniq_last_ts
org = ud.uniq_org
if items.size:
recent_first_i = items[::-1]
recent_first_t = ts[::-1]
recent_first_o = org[::-1]
else:
recent_first_i = items
recent_first_t = ts
recent_first_o = org
long_i = pad_1d(recent_first_i, HIST_LONG, np.int32)
long_o = pad_1d(recent_first_o, HIST_LONG, np.int64)
long_ids = np.where(long_i > 0, long_i + 1, 0).astype(np.int64)
if recent_first_i.size:
recency = (cutoff - recent_first_t[:HIST_LONG].astype(np.float32)) / DAY
recency = np.clip(recency, 0.0, 365.0)
w = np.exp(-recency / 18.0) * idf[recent_first_i[:HIST_LONG]]
long_w = pad_1d(w.astype(np.float32), HIST_LONG, np.float32)
else:
long_w = np.zeros(HIST_LONG, dtype=np.float32)
short_raw = pad_1d(recent_first_i, HIST_SHORT, np.int32)
short_ids = np.where(short_raw > 0, short_raw + 1, 0).astype(np.int64)
if like_override is None or dislike_override is None:
likes, dislikes = explicit_state(ud, cutoff)
else:
likes, dislikes = like_override, dislike_override
if likes.size:
likes = likes[::-1]
if dislikes.size:
dislikes = dislikes[::-1]
like_ids = np.where(
pad_1d(likes, MAX_LIKES, np.int32) > 0,
pad_1d(likes, MAX_LIKES, np.int32) + 1,
0,
).astype(np.int64)
dislike_ids = np.where(
pad_1d(dislikes, MAX_DISLIKES, np.int32) > 0,
pad_1d(dislikes, MAX_DISLIKES, np.int32) + 1,
0,
).astype(np.int64)
if ud.skip_item.size:
skip_keep = ud.skip_ts < cutoff
skip_items = ud.skip_item[skip_keep][::-1]
else:
skip_items = ud.skip_item
skip_ids = np.where(
pad_1d(skip_items, MAX_SKIPS, np.int32) > 0,
pad_1d(skip_items, MAX_SKIPS, np.int32) + 1,
0,
).astype(np.int64)
n_hist = int(items.size)
n_like = int(likes.size)
n_dis = int(dislikes.size)
last_age = 0.0 if not ts.size else max(0.0, (cutoff - float(ts[-1])) / DAY)
mean_idf = float(idf[items[-HIST_LONG:]].mean()) if n_hist else 0.0
org_frac = float(org[-HIST_LONG:].mean()) if n_hist else 0.0
user_feat = np.array(
[
math.log1p(n_hist),
math.log1p(n_like),
math.log1p(n_dis),
org_frac,
math.log1p(last_age),
mean_idf,
1.0 if n_like else 0.0,
1.0 if n_dis else 0.0,
1.0 if n_hist else 0.0,
min(n_hist, HIST_LONG) / float(HIST_LONG),
],
dtype=np.float32,
)
return {
"long_ids": long_ids,
"long_w": long_w,
"long_org": np.clip(long_o, 0, 1).astype(np.int64),
"short_ids": short_ids,
"like_ids": like_ids,
"dislike_ids": dislike_ids,
"skip_ids": skip_ids,
"user_feat": user_feat,
}
def _select_window_positives(ud: UserArrays, start: int, end: int, hist_items: set[int]) -> list[tuple[int, float]]:
if ud.disc_item.size == 0:
return []
ts = ud.disc_ts
lo = int(np.searchsorted(ts, start, side="left"))
hi = int(np.searchsorted(ts, end, side="left"))
best: dict[int, int] = {}
for i in range(lo, hi):
item = int(ud.disc_item[i])
if item in hist_items:
continue
grade = int(ud.disc_grade[i])
prev = best.get(item, 0)
if grade > prev:
best[item] = grade
ranked = sorted(best.items(), key=lambda kv: (-kv[1], kv[0]))
out = []
for item, grade in ranked[:MAX_POS_PER_USER_WINDOW]:
weight = 3.0 if grade >= 3 else 1.0
out.append((item, weight))
return out
def build_samples(users: dict[int, UserArrays], idf: np.ndarray, label_ends: list[int]) -> dict[str, np.ndarray]:
buckets = {
"long_ids": [],
"long_w": [],
"long_org": [],
"short_ids": [],
"like_ids": [],
"dislike_ids": [],
"skip_ids": [],
"user_feat": [],
"pos": [],
"weight": [],
}
n_seq = 0
n_win = 0
for label_end in label_ends:
label_start = label_end - WEEK
hist_cutoff = label_start - GAP_SECONDS
for ud in users.values():
if ud.uniq_item.size:
hist_keep = ud.uniq_last_ts < hist_cutoff
hist_items = set(int(x) for x in ud.uniq_item[hist_keep])
else:
hist_items = set()
positives = _select_window_positives(ud, label_start, label_end, hist_items)
if not positives:
continue
packed = pack_user_inputs(ud, hist_cutoff, idf)
for item, weight in positives:
for key in (
"long_ids",
"long_w",
"long_org",
"short_ids",
"like_ids",
"dislike_ids",
"skip_ids",
"user_feat",
):
buckets[key].append(packed[key])
buckets["pos"].append(item)
buckets["weight"].append(weight)
n_win += 1
for ud in users.values():
if ud.uniq_item.size < 6:
continue
last_ts = ud.uniq_last_ts
items = ud.uniq_item
disc_set: dict[int, int] = {}
if ud.disc_item.size:
for item, grade in zip(ud.disc_item.tolist(), ud.disc_grade.tolist()):
g = disc_set.get(int(item), 0)
if int(grade) > g:
disc_set[int(item)] = int(grade)
taken = 0
start = max(5, int(items.size) - 80)
for i in range(items.size - 1, start - 1, -1):
if taken >= MAX_SEQ_POS_PER_USER:
break
cutoff = int(last_ts[i])
item = int(items[i])
grade = disc_set.get(item, 0)
if grade <= 0:
continue
packed = pack_user_inputs(ud, cutoff, idf)
for key in (
"long_ids",
"long_w",
"long_org",
"short_ids",
"like_ids",
"dislike_ids",
"skip_ids",
"user_feat",
):
buckets[key].append(packed[key])
buckets["pos"].append(item)
buckets["weight"].append(3.0 if grade >= 3 else 1.0)
taken += 1
n_seq += 1
if not buckets["pos"]:
raise RuntimeError("no training samples constructed")
log(f"samples window={n_win} sequential={n_seq} total={len(buckets['pos'])}")
arrays = {
"long_ids": np.stack(buckets["long_ids"]).astype(np.int64),
"long_w": np.stack(buckets["long_w"]).astype(np.float32),
"long_org": np.stack(buckets["long_org"]).astype(np.int64),
"short_ids": np.stack(buckets["short_ids"]).astype(np.int64),
"like_ids": np.stack(buckets["like_ids"]).astype(np.int64),
"dislike_ids": np.stack(buckets["dislike_ids"]).astype(np.int64),
"skip_ids": np.stack(buckets["skip_ids"]).astype(np.int64),
"user_feat": np.stack(buckets["user_feat"]).astype(np.float32),
"pos": np.asarray(buckets["pos"], dtype=np.int64),
"weight": np.asarray(buckets["weight"], dtype=np.float32),
}
return arrays
def make_loader(samples: dict[str, np.ndarray], batch_size: int, shuffle: bool) -> DataLoader:
tensors = [
torch.from_numpy(samples["long_ids"]),
torch.from_numpy(samples["long_w"]),
torch.from_numpy(samples["long_org"]),
torch.from_numpy(samples["short_ids"]),
torch.from_numpy(samples["like_ids"]),
torch.from_numpy(samples["dislike_ids"]),
torch.from_numpy(samples["skip_ids"]),
torch.from_numpy(samples["user_feat"]),
torch.from_numpy(samples["pos"]),
torch.from_numpy(samples["weight"]),
]
ds = TensorDataset(*tensors)
return DataLoader(ds, batch_size=batch_size, shuffle=shuffle, drop_last=False)
class EMA:
def __init__(self, model: nn.Module, decay: float = 0.995):
self.decay = decay
self.shadow = {k: v.detach().clone() for k, v in model.state_dict().items()}
def update(self, model: nn.Module) -> None:
for key, value in model.state_dict().items():
if not value.dtype.is_floating_point:
self.shadow[key] = value.detach().clone()
continue
self.shadow[key].mul_(self.decay).add_(value.detach(), alpha=1.0 - self.decay)
def copy_to(self, model: nn.Module) -> dict[str, torch.Tensor]:
current = {k: v.detach().clone() for k, v in model.state_dict().items()}
model.load_state_dict(self.shadow, strict=True)
return current
def train_one_epoch(model: TwoTower, loader: DataLoader, opt: torch.optim.Optimizer, ema: EMA) -> float:
model.train()
total = 0.0
n = 0
for step, batch in enumerate(loader, start=1):
(
long_ids,
long_w,
long_org,
short_ids,
like_ids,
dislike_ids,
skip_ids,
user_feat,
pos,
weight,
) = batch
user_vec = model.encode_users(
long_ids, long_w, long_org, short_ids, like_ids, dislike_ids, skip_ids, user_feat
)
item_vec = model.encode_all_items()
logits = user_vec @ item_vec.T * model.tau()
per = F.cross_entropy(logits, pos, reduction="none")
loss = (per * weight).mean()
opt.zero_grad(set_to_none=True)
loss.backward()
nn.utils.clip_grad_norm_(model.parameters(), 1.0)
opt.step()
ema.update(model)
total += float(loss.item()) * int(pos.size(0))
n += int(pos.size(0))
if step % LOG_EVERY == 0:
log(f" step {step}/{len(loader)} loss={loss.item():.4f} tau={float(model.tau()):.2f}")
return total / max(n, 1)
def preference_index_sets(path: Path, mapping: dict[int, int]) -> dict[int, tuple[np.ndarray, np.ndarray]]:
frame = pd.read_parquet(path)
out: dict[int, tuple[list[int], list[int]]] = {}
for uid, item, state in frame.itertuples(index=False, name=None):
idx = mapping.get(int(item))
if idx is None:
continue
rec = out.get(int(uid))
if rec is None:
rec = ([], [])
out[int(uid)] = rec
if state == "liked":
rec[0].append(idx)
elif state == "disliked":
rec[1].append(idx)
packed = {}
for uid, (likes, dislikes) in out.items():
packed[uid] = (
np.asarray(likes, dtype=np.int32),
np.asarray(dislikes, dtype=np.int32),
)
return packed
def encode_target_users(
model: TwoTower,
target_uids: np.ndarray,
users: dict[int, UserArrays],
idf: np.ndarray,
cutoff: int,
pref: dict[int, tuple[np.ndarray, np.ndarray]],
) -> np.ndarray:
model.eval()
rows = []
with torch.inference_mode():
batch_acc: dict[str, list] = {k: [] for k in (
"long_ids", "long_w", "long_org", "short_ids", "like_ids", "dislike_ids", "skip_ids", "user_feat"
)}
empty = UserArrays(
pos_item=np.zeros(0, np.int32),
pos_ts=np.zeros(0, np.int32),
pos_org=np.zeros(0, np.int8),
skip_item=np.zeros(0, np.int32),
skip_ts=np.zeros(0, np.int32),
disc_item=np.zeros(0, np.int32),
disc_ts=np.zeros(0, np.int32),
disc_grade=np.zeros(0, np.int8),
expl_ts=np.zeros(0, np.int32),
expl_item=np.zeros(0, np.int32),
expl_kind=np.zeros(0, np.int8),
uniq_item=np.zeros(0, np.int32),
uniq_last_ts=np.zeros(0, np.int32),
uniq_org=np.zeros(0, np.int8),
)
uids = []
def flush():
if not uids:
return
tensors = {k: torch.from_numpy(np.stack(v)) for k, v in batch_acc.items()}
vec = model.encode_users(
tensors["long_ids"],
tensors["long_w"],
tensors["long_org"],
tensors["short_ids"],
tensors["like_ids"],
tensors["dislike_ids"],
tensors["skip_ids"],
tensors["user_feat"],
)
rows.append(vec.numpy())
for k in batch_acc:
batch_acc[k].clear()
uids.clear()
for uid in target_uids:
ud = users.get(int(uid), empty)
like_ov = dis_ov = None
if int(uid) in pref:
like_ov, dis_ov = pref[int(uid)]
packed = pack_user_inputs(ud, cutoff, idf, like_ov, dis_ov)
for k, v in packed.items():
batch_acc[k].append(v)
uids.append(int(uid))
if len(uids) == 512:
flush()
flush()
return np.concatenate(rows, axis=0)
def eligibility_mask(path: Path) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
with np.load(path, allow_pickle=False) as art:
uids = art["uids"].astype(np.int64)
item_ids = art["item_ids"].astype(np.int64)
packed = art["packed_eligible"]
n_items = item_ids.shape[0]
mask = np.unpackbits(packed, axis=1, bitorder="little")[:, :n_items].astype(bool)
return uids, item_ids, mask
def apply_pref_mask(
mask: np.ndarray,
uids: np.ndarray,
pref: dict[int, tuple[np.ndarray, np.ndarray]],
) -> np.ndarray:
masked = mask.copy()
uid_row = {int(u): i for i, u in enumerate(uids)}
for uid, (likes, dislikes) in pref.items():
row = uid_row.get(int(uid))
if row is None:
continue
if likes.size:
masked[row, likes] = False
if dislikes.size:
masked[row, dislikes] = False
return masked
def topk_ranking(
user_vecs: np.ndarray,
item_vecs: np.ndarray,
mask: np.ndarray,
uids: np.ndarray,
item_ids: np.ndarray,
k: int = 100,
) -> pd.DataFrame:
scores = user_vecs @ item_vecs.T
scores = np.where(mask, scores, -np.inf)
top_idx = np.argpartition(-scores, kth=k - 1, axis=1)[:, :k]
top_scores = np.take_along_axis(scores, top_idx, axis=1)
order = np.argsort(-top_scores, axis=1)
top_idx = np.take_along_axis(top_idx, order, axis=1)
n_users = uids.shape[0]
out_uid = np.repeat(uids, k)
out_item = item_ids[top_idx.reshape(-1)]
out_rank = np.tile(np.arange(1, k + 1, dtype=np.int64), n_users)
return pd.DataFrame({"uid": out_uid, "item_id": out_item, "rank": out_rank})
def save_ranking(frame: pd.DataFrame, path: Path) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
frame.to_parquet(path, index=False)
def score_frame(frame: pd.DataFrame, graded: dict[int, dict[int, int]]) -> dict:
ranked = {
int(uid): list(map(int, grp.sort_values("rank").item_id.tolist()))
for uid, grp in frame.groupby("uid")
}
return score_rankings(ranked, graded)
def parameter_groups(model: nn.Module):
decay, no_decay = [], []
for name, param in model.named_parameters():
if not param.requires_grad:
continue
if param.ndim == 1 or name.endswith("bias"):
no_decay.append(param)
else:
decay.append(param)
return [
{"params": decay, "weight_decay": WEIGHT_DECAY},
{"params": no_decay, "weight_decay": 0.0},
]
def train_model(
samples: dict[str, np.ndarray],
n_items: int,
item_feat: np.ndarray,
eval_fn,
max_epochs: int,
time_limit: float,
start_time: float,
) -> tuple[TwoTower, dict, float]:
model = TwoTower(n_items, N_ITEM_FEAT, N_USER_FEAT)
model.set_item_features(item_feat)
model.to(DEVICE)
opt = torch.optim.AdamW(parameter_groups(model), lr=LR)
sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=max(1, max_epochs))
ema = EMA(model, decay=0.995)
loader = make_loader(samples, BATCH_SIZE, shuffle=True)
best_score = -1.0
best_state = None
stale = 0
for epoch in range(1, max_epochs + 1):
if time.time() - start_time > time_limit:
log("time limit reached before epoch start")
break
log(f"epoch {epoch}/{max_epochs} batches={len(loader)}")
loss = train_one_epoch(model, loader, opt, ema)
sched.step()
backup = ema.copy_to(model)
score = eval_fn(model)
model.load_state_dict(backup, strict=True)
log(f"epoch {epoch} loss={loss:.4f} val_ndcg={score:.6f}")
if score > best_score + 1e-6:
best_score = score
best_state = {k: v.detach().clone() for k, v in ema.shadow.items()}
stale = 0
else:
stale += 1
if stale >= 2 and epoch >= 3:
log("early stop")
break
if best_state is None:
best_state = {k: v.detach().clone() for k, v in model.state_dict().items()}
best_score = 0.0
model.load_state_dict(best_state, strict=True)
model.set_item_features(item_feat)
return model, best_state, best_score
def load_stage_events(include_val: bool) -> tuple:
catalog, mapping, item_ids = load_catalog()
listen_paths = [DATA / "train_listens.parquet"]
like_paths = [DATA / "train_likes.parquet"]
unlike_paths = [DATA / "train_unlikes.parquet"]
dislike_paths = [DATA / "train_dislikes.parquet"]
undislike_paths = [DATA / "train_undislikes.parquet"]
if include_val:
listen_paths.append(DATA / "val_listens.parquet")
like_paths.append(DATA / "val_likes.parquet")
unlike_paths.append(DATA / "val_unlikes.parquet")
dislike_paths.append(DATA / "val_dislikes.parquet")
undislike_paths.append(DATA / "val_undislikes.parquet")
log(f"loading listens include_val={include_val}")
listens = load_listen_tables(listen_paths, catalog.item_id.to_numpy(dtype=np.int64))
likes = load_explicit(like_paths, catalog.item_id.to_numpy(dtype=np.int64))
unlikes = load_explicit(unlike_paths, catalog.item_id.to_numpy(dtype=np.int64))
dislikes = load_explicit(dislike_paths, catalog.item_id.to_numpy(dtype=np.int64))
undislikes = load_explicit(undislike_paths, catalog.item_id.to_numpy(dtype=np.int64))
like_tables = _concat_parquets(like_paths, ["uid", "timestamp", "item_id", "is_organic"])
like_org_all = like_tables.column("is_organic").to_numpy().astype(np.int8)
like_items_all = like_tables.column("item_id").to_numpy().astype(np.int64)
like_org_mask = _catalog_mask(like_items_all, catalog.item_id.to_numpy(dtype=np.int64))
like_org = like_org_all[like_org_mask]
log(f"catalog listens={len(listens['uid'])} likes={len(likes[0])}")
item_feat, idf = compute_item_features(
listens, mapping, len(item_ids), likes[2], dislikes[2], catalog
)
log("building user store")
users = build_user_store(listens, mapping, likes, unlikes, dislikes, undislikes, like_org)
log(f"users in store={len(users)}")
return catalog, mapping, item_ids, users, item_feat, idf
def window_ends_for_train(include_val_window: bool) -> list[int]:
ends = [TRAIN_CUTOFF - i * WEEK for i in range(N_TRAIN_WINDOWS)]
if include_val_window:
val_end = TRAIN_CUTOFF + GAP_SECONDS + WEEK
ends = [val_end] + ends
return ends
def main() -> None:
SOLUTION.mkdir(parents=True, exist_ok=True)
(SOLUTION / "train.log").write_text("", encoding="utf-8")
set_seeds(SEED)
start = time.time()
log("start two-tower training")
targets = pd.read_parquet(DATA / "target_users.parquet")
target_uids = targets.uid.to_numpy(dtype=np.int64)
target_set = set(map(int, target_uids))
catalog, mapping, item_ids, users, item_feat, idf = load_stage_events(include_val=False)
n_items = len(item_ids)
samples = build_samples(users, idf, window_ends_for_train(include_val_window=False))
val_pref = preference_index_sets(DATA / "validation_preference_state.parquet", mapping)
val_uids, val_item_ids, val_el = eligibility_mask(DATA / "validation_eligible_candidates.npz")
val_mask = apply_pref_mask(val_el, val_uids, val_pref)
elig = CandidateSet(DATA / "validation_eligible_candidates.npz")
graded = validation_truth(elig, target_set)
log(f"validation graded users={len(graded)}")
uid_order = val_uids
row_of = {int(u): i for i, u in enumerate(uid_order)}
target_in_order = uid_order
def eval_model(model: TwoTower) -> float:
model.eval()
with torch.inference_mode():
item_vecs = model.encode_all_items().numpy()
user_vecs = encode_target_users(
model, target_in_order, users, idf, TRAIN_CUTOFF, val_pref
)
frame = topk_ranking(user_vecs, item_vecs, val_mask, uid_order, val_item_ids)
metrics = score_frame(frame, graded)
return float(metrics["quality_ndcg@10"])
def eval_and_save(model: TwoTower, path: Path) -> float:
model.eval()
with torch.inference_mode():
item_vecs = model.encode_all_items().numpy()
user_vecs = encode_target_users(
model, target_in_order, users, idf, TRAIN_CUTOFF, val_pref
)
frame = topk_ranking(user_vecs, item_vecs, val_mask, uid_order, val_item_ids)
save_ranking(frame, path)
metrics = score_frame(frame, graded)
log(f"saved {path} {metrics}")
return float(metrics["quality_ndcg@10"])
val_time_limit = 28 * 60
model, best_state, best_score = train_model(
samples,
n_items,
item_feat,
eval_fn=lambda m: eval_and_save(m, SOLUTION / "validation.parquet"),
max_epochs=MAX_EPOCHS_VAL,
time_limit=val_time_limit,
start_time=start,
)
log(f"best public val ndcg={best_score:.6f}")
torch.save(best_state, SOLUTION / "val_model.pt")
remaining = 54 * 60 - (time.time() - start)
log(f"seconds remaining before final fit budget snapshot={remaining:.0f}")
catalog, mapping, item_ids, users, item_feat, idf = load_stage_events(include_val=True)
samples = build_samples(users, idf, window_ends_for_train(include_val_window=True))
test_pref = preference_index_sets(DATA / "test_preference_state.parquet", mapping)
test_uids, test_item_ids, test_el = eligibility_mask(DATA / "test_eligible_candidates.npz")
test_mask = apply_pref_mask(test_el, test_uids, test_pref)
test_cutoff = TRAIN_CUTOFF + GAP_SECONDS + WEEK
def dummy_eval(model: TwoTower) -> float:
return 0.0
final_epochs = MAX_EPOCHS_FINAL
if remaining < 12 * 60:
final_epochs = 3
final_model, final_state, _ = train_model(
samples,
n_items,
item_feat,
eval_fn=dummy_eval,
max_epochs=final_epochs,
time_limit=52 * 60,
start_time=start,
)
torch.save(final_state, SOLUTION / "final_model.pt")
final_model.eval()
with torch.inference_mode():
item_vecs = final_model.encode_all_items().numpy()
user_vecs = encode_target_users(
final_model, test_uids, users, idf, test_cutoff, test_pref
)
submission = topk_ranking(user_vecs, item_vecs, test_mask, test_uids, test_item_ids)
save_ranking(submission, SOLUTION / "submission.parquet")
log(f"wrote submission rows={len(submission)}")
readme = f"""# Made-for-You two-tower ranker
## Architecture
Separate PyTorch user and item towers produce 64-dimensional vectors. The
request-time score is their dot product (maximum inner-product search over
precomputed item vectors).
**Item tower.** Catalog ID embedding (64d) concatenated with projected
offline item statistics (log positive/organic/like/dislike counts, skip rate,
mean play ratio, popularity). A 2-layer residual GELU MLP emits a vector
whose first 63 dimensions are L2-normalized taste and whose last dimension is
a bounded popularity residual.
**User tower.** No user-ID lookup. Inputs are cutoff-correct interaction
history:
- recency- and IDF-weighted mean of the last 64 unique positive catalog tracks
(organic flag added to the token embedding)
- mean of the last 8 tracks plus learned positional embeddings (short-term)
- mean of active likes, dislikes, skips, and non-organic positives
- scalar history summaries (log counts, organic fraction, recency, mean IDF)
A residual GELU MLP maps the concatenated user features to the same 64d space.
Candidate eligibility is applied as a retrieval mask. Active liked and
disliked items at the cutoff are excluded because they never receive discovery
labels.
## Training objective
Sampled-softmax-style full-catalog softmax (catalog size 30k, so every item is
a negative; the same trainer can swap in sampled negatives for a million-item
catalog). Positives are:
1. Organic discovery events in rolling 7-day windows (organic like = weight 3,
organic listen ≥80% = weight 1), using only history strictly before a 30
minute gap.
2. Sequential first-time strong organic engagements.
Vectors are trained with temperature-scaled cross entropy and an EMA of
weights. No global user-item matrix is materialized or factorized.
## Public validation
Graded NDCG@10 on the public window, history = `train_*` only:
**{best_score:.6f}**
## Submission command
```sh
python /app/solution/two_tower_ranker.py
```
The script writes `solution/validation.parquet` during model selection, then
refits on `train_* + val_*` and writes `solution/submission.parquet`.
"""
(SOLUTION / "README.md").write_text(readme, encoding="utf-8")
log(f"done in {(time.time() - start) / 60:.1f} min")
if __name__ == "__main__":
main()