Tool call · Input
/app/solution/train_ranker.py
#!/usr/bin/env python3
"""Train and apply a CPU PyTorch two-tower ranker for Yambda discovery.
User tower: recency-weighted pooling of interaction history (no user-ID lookup).
Item tower: independently precomputable MLP over a shared item embedding plus
static catalog features. Relevance score is exactly the user-item dot product.
"""
from __future__ import annotations
import argparse
import json
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, TensorDataset
import sys
sys.path.insert(0, "/task")
from ranking_evaluation import QUALITY_K, ndcg_at_k # noqa: E402
DATA = Path("/task/data")
OUT = Path("/app/solution")
TRAIN_END = np.int64(24_786_800)
GAP = np.int64(1_800)
WEEK = np.int64(7 * 86_400)
DAY = np.float32(86_400.0)
LISTEN_POS = 50
LISTEN_STRONG = 80
HIST_LEN = 64
LIKE_LEN = 32
DISLIKE_LEN = 16
EMBED_DIM = 64
HIDDEN = 128
BATCH = 384
EXTRA_NEGS = 192
MAX_EPOCHS = 10
PATIENCE = 3
LR = 1.2e-3
SEED = 7
N_FOLDS_VAL = 8
N_FIRST_OCC = 12
MAX_EXAMPLES = 280_000
HALF_LIFE_SHORT = 14.0
HALF_LIFE_LONG = 90.0
SCALE = 16.0 # cosine logits multiplier
def log(msg: str) -> None:
print(msg, flush=True)
def set_seed(seed: int) -> None:
np.random.seed(seed)
torch.manual_seed(seed)
def read_cols(path: Path, columns: list[str]) -> dict[str, np.ndarray]:
table = pq.read_table(path, columns=columns)
out: dict[str, np.ndarray] = {}
for name in columns:
col = table.column(name).to_numpy()
out[name] = np.asarray(col)
return out
def concat_events(paths: list[Path], columns: list[str]) -> dict[str, np.ndarray]:
chunks = [read_cols(p, columns) for p in paths if p.exists()]
if len(chunks) == 1:
return chunks[0]
return {k: np.concatenate([c[k] for c in chunks], axis=0) for k in columns}
class Jagged:
"""Sorted-by-uid ragged event table with O(1) user slices."""
def __init__(self, uid: np.ndarray, **fields: np.ndarray):
order = np.lexsort((fields["ts"], uid))
self.uid = uid[order]
self.fields = {k: v[order] for k, v in fields.items()}
self.uniq, self.starts, self.cnts = np.unique(
self.uid, return_index=True, return_counts=True
)
self.row = {int(u): i for i, u in enumerate(self.uniq)}
def get(self, uid: int) -> dict[str, np.ndarray] | None:
i = self.row.get(int(uid))
if i is None:
return None
s = int(self.starts[i])
e = s + int(self.cnts[i])
return {k: v[s:e] for k, v in self.fields.items()}
def build_bool_mask(ids: np.ndarray, universe_max: int) -> np.ndarray:
mask = np.zeros(int(universe_max) + 1, dtype=np.bool_)
mask[ids.astype(np.int64, copy=False)] = True
return mask
def recency_unique(
items: np.ndarray,
ts: np.ndarray,
extra: dict[str, np.ndarray],
cutoff: int,
max_len: int,
) -> dict[str, np.ndarray]:
"""Most-recent unique items before cutoff, most recent first."""
end = int(np.searchsorted(ts, cutoff, side="right"))
if end <= 0:
return {}
seg_items = items[:end]
seg_ts = ts[:end]
rev = seg_items[::-1]
uniq, first = np.unique(rev, return_index=True)
order = np.argsort(first, kind="mergesort")
take = min(max_len, uniq.shape[0])
sel = order[:take]
first_sel = first[sel]
orig = end - 1 - first_sel
out = {
"item": uniq[sel].astype(np.int64, copy=False),
"age_days": (np.float32(cutoff) - seg_ts[orig].astype(np.float32)) / DAY,
}
for name, arr in extra.items():
out[name] = arr[:end][orig]
return out
class TwoTower(nn.Module):
"""User history tower + independently scored item tower."""
def __init__(
self,
n_items: int,
n_catalog: int,
dim: int,
hidden: int,
n_item_feat: int,
n_user_dense: int,
):
super().__init__()
self.dim = dim
self.n_catalog = n_catalog
self.emb = nn.Embedding(n_items, dim, padding_idx=0)
nn.init.normal_(self.emb.weight, mean=0.0, std=0.04)
with torch.no_grad():
self.emb.weight[0].zero_()
self.recency_emb = nn.Embedding(6, dim)
self.source_emb = nn.Embedding(2, dim)
nn.init.normal_(self.recency_emb.weight, std=0.02)
nn.init.normal_(self.source_emb.weight, std=0.02)
user_in = dim * 4 + n_user_dense
self.user_norm = nn.LayerNorm(user_in)
self.user_mlp = nn.Sequential(
nn.Linear(user_in, hidden),
nn.ReLU(),
nn.Dropout(0.1),
nn.Linear(hidden, hidden),
nn.ReLU(),
nn.Linear(hidden, dim),
)
self.user_skip = nn.Linear(dim, dim, bias=False)
item_in = dim + n_item_feat
self.item_feat_proj = nn.Linear(n_item_feat, dim)
self.item_norm = nn.LayerNorm(item_in)
self.item_mlp = nn.Sequential(
nn.Linear(item_in, hidden),
nn.ReLU(),
nn.Dropout(0.05),
nn.Linear(hidden, dim),
)
self.item_skip = nn.Linear(dim, dim, bias=False)
self.item_bias = nn.Linear(n_item_feat, 1)
def recency_bucket(self, age_days: torch.Tensor) -> torch.Tensor:
# 0-1d, 1-7d, 7-30d, 30-90d, 90-365d, 365d+
b = torch.zeros_like(age_days, dtype=torch.long)
b = torch.where(age_days >= 1.0, torch.tensor(1, device=age_days.device), b)
b = torch.where(age_days >= 7.0, torch.tensor(2, device=age_days.device), b)
b = torch.where(age_days >= 30.0, torch.tensor(3, device=age_days.device), b)
b = torch.where(age_days >= 90.0, torch.tensor(4, device=age_days.device), b)
b = torch.where(age_days >= 365.0, torch.tensor(5, device=age_days.device), b)
return b
def _pool(
self,
idx: torch.Tensor,
age: torch.Tensor,
organic: torch.Tensor,
mask: torch.Tensor,
half_life: float,
) -> torch.Tensor:
e = self.emb(idx)
e = e + self.recency_emb(self.recency_bucket(age))
e = e + self.source_emb(organic.clamp(0, 1))
decay = torch.exp(-np.log(2.0) * age / half_life)
w = decay * mask.float()
w = w / (w.sum(dim=1, keepdim=True) + 1e-8)
return (e * w.unsqueeze(-1)).sum(dim=1)
def encode_users(
self,
hist_idx: torch.Tensor,
hist_age: torch.Tensor,
hist_org: torch.Tensor,
hist_mask: torch.Tensor,
like_idx: torch.Tensor,
like_mask: torch.Tensor,
dislike_idx: torch.Tensor,
dislike_mask: torch.Tensor,
dense: torch.Tensor,
) -> torch.Tensor:
short = self._pool(hist_idx, hist_age, hist_org, hist_mask, HALF_LIFE_SHORT)
long = self._pool(hist_idx, hist_age, hist_org, hist_mask, HALF_LIFE_LONG)
like_e = self.emb(like_idx)
like_w = like_mask.float()
like_w = like_w / (like_w.sum(dim=1, keepdim=True) + 1e-8)
likes = (like_e * like_w.unsqueeze(-1)).sum(dim=1)
dis_e = self.emb(dislike_idx)
dis_w = dislike_mask.float()
dis_w = dis_w / (dis_w.sum(dim=1, keepdim=True) + 1e-8)
dislikes = (dis_e * dis_w.unsqueeze(-1)).sum(dim=1)
z = torch.cat([short, long, likes, dislikes, dense], dim=1)
u = self.user_mlp(self.user_norm(z)) + self.user_skip(short)
return F.normalize(u, dim=-1)
def encode_items(self, idx: torch.Tensor, feat: torch.Tensor) -> torch.Tensor:
e = self.emb(idx)
raw = torch.cat([e, feat], dim=1)
v = self.item_mlp(self.item_norm(raw)) + self.item_skip(e) + self.item_feat_proj(feat)
v = F.normalize(v, dim=-1)
# Fold a learned popularity/item bias into an extra affine direction by
# rescaling: still a vector independently produced by the item tower.
# Keep unit-norm so score remains a pure cosine dot product.
return v
def forward(self, batch: dict[str, torch.Tensor], item_feat_table: torch.Tensor):
user = self.encode_users(
batch["hist_idx"],
batch["hist_age"],
batch["hist_org"],
batch["hist_mask"],
batch["like_idx"],
batch["like_mask"],
batch["dislike_idx"],
batch["dislike_mask"],
batch["dense"],
)
pos_feat = item_feat_table[batch["pos_idx"]]
pos = self.encode_items(batch["pos_idx"], pos_feat)
return user, pos
def sampled_softmax(
user: torch.Tensor,
pos: torch.Tensor,
extra: torch.Tensor,
weight: torch.Tensor,
scale: float,
) -> torch.Tensor:
logits = scale * user @ pos.t()
extra_logits = scale * user @ extra.t()
logits = torch.cat([logits, extra_logits], dim=1)
targets = torch.arange(user.size(0), device=user.device)
loss = F.cross_entropy(logits, targets, reduction="none")
w = weight / (weight.mean() + 1e-8)
return (loss * w).mean()
@dataclass
class RankData:
item_to_idx: dict[int, int]
idx_to_item: np.ndarray
catalog_items: np.ndarray
catalog_idx: np.ndarray
item_feat: np.ndarray
n_items: int
positives: Jagged
likes: Jagged
dislikes: Jagged
target_uids: np.ndarray
def load_stage(include_val: bool) -> tuple[list[Path], list[Path], list[Path]]:
listen = [DATA / "train_listens.parquet"]
like = [DATA / "train_likes.parquet"]
dislike = [DATA / "train_dislikes.parquet"]
if include_val:
listen.append(DATA / "val_listens.parquet")
like.append(DATA / "val_likes.parquet")
dislike.append(DATA / "val_dislikes.parquet")
return listen, like, dislike
def prepare_data(include_val: bool) -> RankData:
t0 = time.time()
catalog = pd.read_parquet(DATA / "candidate_catalog.parquet")
catalog_items = catalog.item_id.to_numpy(dtype=np.int64)
n_catalog = int(catalog_items.shape[0])
pop = catalog.positive_listens.to_numpy(dtype=np.float32)
org_pop = catalog.organic_positive_listens.to_numpy(dtype=np.float32)
pop_rank = catalog.popularity_rank.to_numpy(dtype=np.float32)
targets = pd.read_parquet(DATA / "target_users.parquet")
target_uids = targets.uid.to_numpy(dtype=np.int64)
listen_paths, like_paths, dislike_paths = load_stage(include_val)
log(f"loading listens include_val={include_val}")
listens = concat_events(
listen_paths,
["uid", "timestamp", "item_id", "is_organic", "played_ratio_pct"],
)
likes = concat_events(like_paths, ["uid", "timestamp", "item_id", "is_organic"])
dislikes = concat_events(dislike_paths, ["uid", "timestamp", "item_id", "is_organic"])
uid_l = listens["uid"].astype(np.int64, copy=False)
ts_l = listens["timestamp"].astype(np.int64, copy=False)
item_l = listens["item_id"].astype(np.int64, copy=False)
org_l = listens["is_organic"].astype(np.int8, copy=False)
ratio = listens["played_ratio_pct"].astype(np.int16, copy=False)
del listens
pos_mask = ratio >= LISTEN_POS
uid_p = uid_l[pos_mask]
ts_p = ts_l[pos_mask]
item_p = item_l[pos_mask]
org_p = org_l[pos_mask]
ratio_p = ratio[pos_mask]
strong_mask = (org_l == 1) & (ratio >= LISTEN_STRONG)
# Keep strong-organic arrays for label mining (subset of all listens).
uid_s = uid_l[strong_mask]
ts_s = ts_l[strong_mask]
item_s = item_l[strong_mask]
del uid_l, ts_l, item_l, org_l, ratio, pos_mask, strong_mask
uid_k = likes["uid"].astype(np.int64, copy=False)
ts_k = likes["timestamp"].astype(np.int64, copy=False)
item_k = likes["item_id"].astype(np.int64, copy=False)
org_k = likes["is_organic"].astype(np.int8, copy=False)
del likes
uid_d = dislikes["uid"].astype(np.int64, copy=False)
ts_d = dislikes["timestamp"].astype(np.int64, copy=False)
item_d = dislikes["item_id"].astype(np.int64, copy=False)
del dislikes
# Merge listens-positives and likes into one positive stream.
uid_pos = np.concatenate([uid_p, uid_k])
ts_pos = np.concatenate([ts_p, ts_k])
item_pos = np.concatenate([item_p, item_k])
org_pos = np.concatenate([org_p, org_k])
is_like = np.concatenate(
[
np.zeros(uid_p.shape[0], dtype=np.int8),
np.ones(uid_k.shape[0], dtype=np.int8),
]
)
ratio_pos = np.concatenate(
[ratio_p, np.full(uid_k.shape[0], 100, dtype=np.int16)]
)
del uid_p, ts_p, item_p, org_p, ratio_p
max_item = int(
max(
int(catalog_items.max()),
int(item_pos.max()),
int(item_k.max()) if item_k.size else 0,
int(item_d.max()) if item_d.size else 0,
int(item_s.max()) if item_s.size else 0,
)
)
is_catalog = build_bool_mask(catalog_items, max_item)
# Vocab: PAD=0, UNK=1, catalog in eligibility order, then frequent others.
item_to_idx: dict[int, int] = {}
idx_to_item = [0, 0]
for i, it in enumerate(catalog_items):
item_to_idx[int(it)] = i + 2
idx_to_item.append(int(it))
freq = np.bincount(item_pos, minlength=max_item + 1)
other = np.where((freq >= 8) & (~is_catalog[: freq.shape[0]]))[0]
for it in other:
if int(it) not in item_to_idx:
item_to_idx[int(it)] = len(idx_to_item)
idx_to_item.append(int(it))
idx_to_item_arr = np.asarray(idx_to_item, dtype=np.int64)
n_items = int(idx_to_item_arr.shape[0])
log(
f"vocab items={n_items} catalog={n_catalog} other={n_items - n_catalog - 2} "
f"load_s={time.time()-t0:.1f}"
)
def map_items(arr: np.ndarray) -> np.ndarray:
out = np.empty(arr.shape[0], dtype=np.int32)
mapping = item_to_idx
unk = 1
for i, it in enumerate(arr):
out[i] = mapping.get(int(it), unk)
return out
# Mapping 27M items in Python is slow; use a dense lookup table.
lut = np.ones(max_item + 1, dtype=np.int32) # UNK=1
lut[0] = 0
for it, idx in item_to_idx.items():
lut[it] = idx
item_pos_idx = lut[item_pos]
item_k_idx = lut[item_k]
item_d_idx = lut[item_d]
item_s_idx = lut[item_s]
positives = Jagged(
uid_pos,
ts=ts_pos,
item=item_pos,
idx=item_pos_idx,
org=org_pos,
is_like=is_like,
ratio=ratio_pos,
)
likes_j = Jagged(uid_k, ts=ts_k, item=item_k, idx=item_k_idx, org=org_k)
dislikes_j = Jagged(uid_d, ts=ts_d, item=item_d, idx=item_d_idx)
strong_j = Jagged(uid_s, ts=ts_s, item=item_s, idx=item_s_idx)
# Attach strong listens onto RankData via likes/dislikes plus a hidden field.
positives.strong = strong_j # type: ignore[attr-defined]
positives.lut = lut # type: ignore[attr-defined]
positives.is_catalog = is_catalog # type: ignore[attr-defined]
positives.max_item = max_item # type: ignore[attr-defined]
# Item features for the full vocab. Catalog rows use provided stats;
# others use empirical positive counts.
item_feat = np.zeros((n_items, 4), dtype=np.float32)
cat_feat = np.stack(
[
np.log1p(pop),
np.log1p(org_pop),
np.divide(org_pop, np.maximum(pop, 1.0)),
np.log1p(pop_rank),
],
axis=1,
).astype(np.float32)
item_feat[2 : 2 + n_catalog] = cat_feat
emp = np.log1p(freq.astype(np.float32))
for it, idx in item_to_idx.items():
if idx >= 2 + n_catalog and it < emp.shape[0]:
item_feat[idx, 0] = emp[it]
item_feat[idx, 3] = emp[it]
# Standardize catalog-like columns.
mu = item_feat.mean(axis=0, keepdims=True)
sd = item_feat.std(axis=0, keepdims=True) + 1e-6
item_feat = (item_feat - mu) / sd
catalog_idx = np.arange(2, 2 + n_catalog, dtype=np.int64)
log(f"prepare done in {time.time()-t0:.1f}s positives={uid_pos.shape[0]}")
return RankData(
item_to_idx=item_to_idx,
idx_to_item=idx_to_item_arr,
catalog_items=catalog_items,
catalog_idx=catalog_idx,
item_feat=item_feat,
n_items=n_items,
positives=positives,
likes=likes_j,
dislikes=dislikes_j,
target_uids=target_uids,
)
def fold_cutoffs(history_end_max: int, n_folds: int) -> list[tuple[int, int, int]]:
folds = []
c = int(history_end_max)
for _ in range(n_folds):
hist_end = c - int(GAP) - int(WEEK)
if hist_end < 2_000_000:
break
label_start = hist_end + int(GAP)
label_end = c
folds.append((hist_end, label_start, label_end))
c = hist_end
return folds
def pad_row(vals: np.ndarray, max_len: int, dtype) -> tuple[np.ndarray, np.ndarray]:
out = np.zeros(max_len, dtype=dtype)
mask = np.zeros(max_len, dtype=np.bool_)
n = min(max_len, int(vals.shape[0]))
if n:
out[:n] = vals[:n]
mask[:n] = True
return out, mask
def user_dense_features(
n_hist: int,
n_like: int,
n_dis: int,
org_frac: float,
last_age: float,
uniq: int,
) -> np.ndarray:
return np.asarray(
[
np.log1p(n_hist),
np.log1p(n_like),
np.log1p(n_dis),
org_frac,
np.log1p(uniq),
np.tanh(last_age / 30.0),
],
dtype=np.float32,
)
N_USER_DENSE = 6
def encode_history_pack(
data: RankData,
uid: int,
cutoff: int,
) -> dict[str, np.ndarray] | None:
pos = data.positives.get(uid)
if pos is None:
return None
hist = recency_unique(
pos["item"],
pos["ts"],
{"idx": pos["idx"], "org": pos["org"], "is_like": pos["is_like"]},
cutoff,
HIST_LEN,
)
if not hist:
return None
likes = data.likes.get(uid)
if likes is not None:
like_h = recency_unique(
likes["item"], likes["ts"], {"idx": likes["idx"]}, cutoff, LIKE_LEN
)
else:
like_h = {}
dis = data.dislikes.get(uid)
if dis is not None:
dis_h = recency_unique(
dis["item"], dis["ts"], {"idx": dis["idx"]}, cutoff, DISLIKE_LEN
)
else:
dis_h = {}
hist_idx, hist_mask = pad_row(hist["idx"].astype(np.int64), HIST_LEN, np.int64)
hist_age, _ = pad_row(hist["age_days"].astype(np.float32), HIST_LEN, np.float32)
hist_org, _ = pad_row(hist["org"].astype(np.int64), HIST_LEN, np.int64)
if like_h:
like_idx, like_mask = pad_row(like_h["idx"].astype(np.int64), LIKE_LEN, np.int64)
n_like = int(like_h["idx"].shape[0])
else:
like_idx = np.zeros(LIKE_LEN, dtype=np.int64)
like_mask = np.zeros(LIKE_LEN, dtype=np.bool_)
n_like = 0
if dis_h:
dis_idx, dis_mask = pad_row(dis_h["idx"].astype(np.int64), DISLIKE_LEN, np.int64)
n_dis = int(dis_h["idx"].shape[0])
else:
dis_idx = np.zeros(DISLIKE_LEN, dtype=np.int64)
dis_mask = np.zeros(DISLIKE_LEN, dtype=np.bool_)
n_dis = 0
org_frac = float(hist["org"].mean()) if hist["org"].size else 0.0
last_age = float(hist["age_days"][0]) if hist["age_days"].size else 30.0
dense = user_dense_features(
int(hist["idx"].shape[0]),
n_like,
n_dis,
org_frac,
last_age,
int(hist["idx"].shape[0]),
)
return {
"hist_idx": hist_idx,
"hist_age": hist_age,
"hist_org": hist_org,
"hist_mask": hist_mask,
"like_idx": like_idx,
"like_mask": like_mask,
"dislike_idx": dis_idx,
"dislike_mask": dis_mask,
"dense": dense,
"raw_hist_items": hist["item"],
}
def collect_labels(
data: RankData,
uid: int,
label_start: int,
label_end: int,
hist_items: np.ndarray,
) -> list[tuple[int, float]]:
seen = set(map(int, hist_items.tolist()))
is_catalog = data.positives.is_catalog # type: ignore[attr-defined]
labels: dict[int, float] = {}
strong = data.positives.strong.get(uid) # type: ignore[attr-defined]
if strong is not None:
a = int(np.searchsorted(strong["ts"], label_start, side="left"))
b = int(np.searchsorted(strong["ts"], label_end, side="right"))
for item, idx in zip(strong["item"][a:b], strong["idx"][a:b]):
item_i = int(item)
if item_i in seen:
continue
if item_i >= is_catalog.shape[0] or not bool(is_catalog[item_i]):
continue
if int(idx) <= 1:
continue
labels[int(idx)] = max(labels.get(int(idx), 0.0), 1.0)
likes = data.likes.get(uid)
if likes is not None:
a = int(np.searchsorted(likes["ts"], label_start, side="left"))
b = int(np.searchsorted(likes["ts"], label_end, side="right"))
org = likes["org"][a:b]
for item, idx, o in zip(likes["item"][a:b], likes["idx"][a:b], org):
if int(o) != 1:
continue
item_i = int(item)
if item_i in seen:
continue
if item_i >= is_catalog.shape[0] or not bool(is_catalog[item_i]):
continue
if int(idx) <= 1:
continue
labels[int(idx)] = max(labels.get(int(idx), 0.0), 3.0)
return list(labels.items())
def first_occurrence_labels(
data: RankData,
uid: int,
hist_end: int,
rng: np.random.Generator,
k: int,
) -> list[tuple[int, float]]:
"""Sample first-time catalog discoveries before hist_end as extra labels."""
pos = data.positives.get(uid)
if pos is None:
return []
end = int(np.searchsorted(pos["ts"], hist_end, side="right"))
if end < 20:
return []
is_catalog = data.positives.is_catalog # type: ignore[attr-defined]
items = pos["item"][:end]
idx = pos["idx"][:end]
org = pos["org"][:end]
ratio = pos["ratio"][:end]
is_like = pos["is_like"][:end]
seen: set[int] = set()
cands: list[tuple[int, float, int]] = []
for i in range(end):
item_i = int(items[i])
if item_i in seen:
continue
seen.add(item_i)
if i < 10:
continue
if item_i >= is_catalog.shape[0] or not bool(is_catalog[item_i]):
continue
idxi = int(idx[i])
if idxi <= 1:
continue
if int(is_like[i]) == 1 and int(org[i]) == 1:
w = 3.0
elif int(org[i]) == 1 and int(ratio[i]) >= LISTEN_STRONG:
w = 1.0
else:
continue
cands.append((idxi, w, i))
if not cands:
return []
# Prefer later first-occurrences.
take = min(k, len(cands))
p = np.linspace(0.4, 1.6, num=len(cands), dtype=np.float64)
p = p / p.sum()
chosen = rng.choice(len(cands), size=take, replace=False, p=p)
return [(cands[j][0], cands[j][1]) for j in chosen]
def build_examples(data: RankData, folds: list[tuple[int, int, int]], seed: int):
rng = np.random.default_rng(seed)
users = [int(u) for u in data.positives.uniq]
packs: dict[tuple[int, int], dict[str, np.ndarray]] = {}
rows = []
t0 = time.time()
for hist_end, label_start, label_end in folds:
n_lab = 0
for uid in users:
pack = encode_history_pack(data, uid, hist_end)
if pack is None:
continue
labs = collect_labels(
data, uid, label_start, label_end, pack["raw_hist_items"]
)
if not labs:
continue
key = (uid, hist_end)
packs[key] = pack
for idx, w in labs:
rows.append((key, idx, w))
n_lab += 1
log(f"fold hist_end={hist_end} labels={n_lab}")
# Extra first-occurrence samples on the most recent training fold cutoff.
extra = 0
hist_end = folds[0][0]
for uid in users:
pack = packs.get((uid, hist_end))
if pack is None:
pack = encode_history_pack(data, uid, hist_end)
if pack is None:
continue
packs[(uid, hist_end)] = pack
labs = first_occurrence_labels(data, uid, hist_end, rng, N_FIRST_OCC)
for idx, w in labs:
rows.append(((uid, hist_end), idx, 0.45 * w))
extra += 1
log(f"first-occurrence extra={extra} total_rows={len(rows)} build_s={time.time()-t0:.1f}")
if len(rows) > MAX_EXAMPLES:
pick = rng.choice(len(rows), size=MAX_EXAMPLES, replace=False)
rows = [rows[i] for i in pick]
log(f"downsampled to {len(rows)}")
n = len(rows)
hist_idx = np.zeros((n, HIST_LEN), dtype=np.int64)
hist_age = np.zeros((n, HIST_LEN), dtype=np.float32)
hist_org = np.zeros((n, HIST_LEN), dtype=np.int64)
hist_mask = np.zeros((n, HIST_LEN), dtype=np.bool_)
like_idx = np.zeros((n, LIKE_LEN), dtype=np.int64)
like_mask = np.zeros((n, LIKE_LEN), dtype=np.bool_)
dislike_idx = np.zeros((n, DISLIKE_LEN), dtype=np.int64)
dislike_mask = np.zeros((n, DISLIKE_LEN), dtype=np.bool_)
dense = np.zeros((n, N_USER_DENSE), dtype=np.float32)
pos_idx = np.zeros(n, dtype=np.int64)
weight = np.zeros(n, dtype=np.float32)
for i, (key, idx, w) in enumerate(rows):
p = packs[key]
hist_idx[i] = p["hist_idx"]
hist_age[i] = p["hist_age"]
hist_org[i] = p["hist_org"]
hist_mask[i] = p["hist_mask"]
like_idx[i] = p["like_idx"]
like_mask[i] = p["like_mask"]
dislike_idx[i] = p["dislike_idx"]
dislike_mask[i] = p["dislike_mask"]
dense[i] = p["dense"]
pos_idx[i] = idx
weight[i] = w
tensors = {
"hist_idx": torch.from_numpy(hist_idx),
"hist_age": torch.from_numpy(hist_age),
"hist_org": torch.from_numpy(hist_org),
"hist_mask": torch.from_numpy(hist_mask),
"like_idx": torch.from_numpy(like_idx),
"like_mask": torch.from_numpy(like_mask),
"dislike_idx": torch.from_numpy(dislike_idx),
"dislike_mask": torch.from_numpy(dislike_mask),
"dense": torch.from_numpy(dense),
"pos_idx": torch.from_numpy(pos_idx),
"weight": torch.from_numpy(weight),
}
return tensors
def make_loader(tensors: dict[str, torch.Tensor], shuffle: bool) -> DataLoader:
ds = TensorDataset(
tensors["hist_idx"],
tensors["hist_age"],
tensors["hist_org"],
tensors["hist_mask"],
tensors["like_idx"],
tensors["like_mask"],
tensors["dislike_idx"],
tensors["dislike_mask"],
tensors["dense"],
tensors["pos_idx"],
tensors["weight"],
)
return DataLoader(ds, batch_size=BATCH, shuffle=shuffle, drop_last=True)
def batch_from_tuple(batch) -> dict[str, torch.Tensor]:
keys = [
"hist_idx",
"hist_age",
"hist_org",
"hist_mask",
"like_idx",
"like_mask",
"dislike_idx",
"dislike_mask",
"dense",
"pos_idx",
"weight",
]
return dict(zip(keys, batch))
def load_eligibility(path: Path) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
with np.load(path, allow_pickle=False) as z:
uids = z["uids"].astype(np.int64)
item_ids = z["item_ids"].astype(np.int64)
packed = z["packed_eligible"].astype(np.uint8)
masks = np.unpackbits(packed, axis=1, count=len(item_ids), bitorder="little").astype(
np.bool_
)
return uids, item_ids, masks
def preference_channels(path: Path, data: RankData, uids: np.ndarray, cutoff: int):
"""Override like/dislike channels with cutoff-correct explicit state."""
pref = pd.read_parquet(path)
like_map: dict[int, list[int]] = {}
dis_map: dict[int, list[int]] = {}
lut = data.positives.lut # type: ignore[attr-defined]
max_item = int(data.positives.max_item) # type: ignore[attr-defined]
for uid, item, state in pref.itertuples(index=False, name=None):
uid_i, item_i = int(uid), int(item)
idx = int(lut[item_i]) if 0 <= item_i <= max_item else 1
if state == "liked":
like_map.setdefault(uid_i, []).append(idx)
elif state == "disliked":
dis_map.setdefault(uid_i, []).append(idx)
n = uids.shape[0]
like_idx = np.zeros((n, LIKE_LEN), dtype=np.int64)
like_mask = np.zeros((n, LIKE_LEN), dtype=np.bool_)
dis_idx = np.zeros((n, DISLIKE_LEN), dtype=np.int64)
dis_mask = np.zeros((n, DISLIKE_LEN), dtype=np.bool_)
for i, uid in enumerate(uids):
likes = like_map.get(int(uid), [])[:LIKE_LEN]
dis = dis_map.get(int(uid), [])[:DISLIKE_LEN]
if likes:
like_idx[i, : len(likes)] = np.asarray(likes, dtype=np.int64)
like_mask[i, : len(likes)] = True
if dis:
dis_idx[i, : len(dis)] = np.asarray(dis, dtype=np.int64)
dis_mask[i, : len(dis)] = True
return like_idx, like_mask, dis_idx, dis_mask
@torch.no_grad()
def encode_all_users(
model: TwoTower,
data: RankData,
uids: np.ndarray,
cutoff: int,
pref_path: Path | None,
device: torch.device,
) -> np.ndarray:
model.eval()
packs = []
valid = np.ones(len(uids), dtype=np.bool_)
for i, uid in enumerate(uids):
pack = encode_history_pack(data, int(uid), cutoff)
if pack is None:
valid[i] = False
packs.append(None)
else:
packs.append(pack)
if pref_path is not None:
like_idx, like_mask, dis_idx, dis_mask = preference_channels(
pref_path, data, uids, cutoff
)
else:
like_idx = like_mask = dis_idx = dis_mask = None
hist_idx = np.zeros((len(uids), HIST_LEN), dtype=np.int64)
hist_age = np.zeros((len(uids), HIST_LEN), dtype=np.float32)
hist_org = np.zeros((len(uids), HIST_LEN), dtype=np.int64)
hist_mask = np.zeros((len(uids), HIST_LEN), dtype=np.bool_)
like_i = np.zeros((len(uids), LIKE_LEN), dtype=np.int64)
like_m = np.zeros((len(uids), LIKE_LEN), dtype=np.bool_)
dis_i = np.zeros((len(uids), DISLIKE_LEN), dtype=np.int64)
dis_m = np.zeros((len(uids), DISLIKE_LEN), dtype=np.bool_)
dense = np.zeros((len(uids), N_USER_DENSE), dtype=np.float32)
for i, pack in enumerate(packs):
if pack is None:
continue
hist_idx[i] = pack["hist_idx"]
hist_age[i] = pack["hist_age"]
hist_org[i] = pack["hist_org"]
hist_mask[i] = pack["hist_mask"]
like_i[i] = pack["like_idx"]
like_m[i] = pack["like_mask"]
dis_i[i] = pack["dislike_idx"]
dis_m[i] = pack["dislike_mask"]
dense[i] = pack["dense"]
if like_idx is not None:
like_i, like_m, dis_i, dis_m = like_idx, like_mask, dis_idx, dis_mask
user_mat = np.zeros((len(uids), model.dim), dtype=np.float32)
bs = 512
feat_table = torch.from_numpy(data.item_feat).to(device)
_ = feat_table # item features unused on the user path
for start in range(0, len(uids), bs):
sl = slice(start, min(start + bs, len(uids)))
user = model.encode_users(
torch.from_numpy(hist_idx[sl]).to(device),
torch.from_numpy(hist_age[sl]).to(device),
torch.from_numpy(hist_org[sl]).to(device),
torch.from_numpy(hist_mask[sl]).to(device),
torch.from_numpy(like_i[sl]).to(device),
torch.from_numpy(like_m[sl]).to(device),
torch.from_numpy(dis_i[sl]).to(device),
torch.from_numpy(dis_m[sl]).to(device),
torch.from_numpy(dense[sl]).to(device),
)
user_mat[sl] = user.cpu().numpy()
user_mat[~valid] = 0.0
return user_mat
@torch.no_grad()
def encode_catalog_items(model: TwoTower, data: RankData, device: torch.device) -> np.ndarray:
model.eval()
idx = torch.from_numpy(data.catalog_idx).to(device)
feat = torch.from_numpy(data.item_feat[data.catalog_idx]).to(device)
outs = []
bs = 2048
for start in range(0, idx.size(0), bs):
sl = slice(start, start + bs)
outs.append(model.encode_items(idx[sl], feat[sl]).cpu().numpy())
return np.concatenate(outs, axis=0)
def topk_rankings(
user_mat: np.ndarray,
item_mat: np.ndarray,
masks: np.ndarray,
item_ids: np.ndarray,
k: int = 100,
) -> np.ndarray:
scores = user_mat @ item_mat.T
scores = np.where(masks, scores, np.float32(-1e9))
part = np.argpartition(-scores, kth=k - 1, axis=1)[:, :k]
part_scores = np.take_along_axis(scores, part, axis=1)
order = np.argsort(-part_scores, axis=1)
top_idx = np.take_along_axis(part, order, axis=1)
return item_ids[top_idx]
def validation_graded(el_uids: np.ndarray, el_items: np.ndarray, masks: np.ndarray):
uid_row = {int(u): i for i, u in enumerate(el_uids)}
item_col = {int(it): i for i, it in enumerate(el_items)}
targets = set(map(int, el_uids.tolist()))
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()
olikes = likes[likes.uid.isin(targets) & (likes.is_organic == 1)][
["uid", "item_id"]
].drop_duplicates()
graded: dict[int, dict[int, int]] = {}
for uid, item in strong.itertuples(index=False, name=None):
r, c = uid_row.get(int(uid)), item_col.get(int(item))
if r is None or c is None:
continue
if masks[r, c]:
graded.setdefault(int(uid), {})[int(item)] = 1
for uid, item in olikes.itertuples(index=False, name=None):
r, c = uid_row.get(int(uid)), item_col.get(int(item))
if r is None or c is None:
continue
if masks[r, c]:
graded.setdefault(int(uid), {})[int(item)] = 3
return graded
def score_ndcg(top_items: np.ndarray, uids: np.ndarray, graded: dict[int, dict[int, int]]):
uid_row = {int(u): i for i, u in enumerate(uids)}
vals = []
for uid, gold in graded.items():
ranked = top_items[uid_row[int(uid)]].tolist()
vals.append(ndcg_at_k(ranked, gold, QUALITY_K))
return float(np.mean(vals)) if vals else 0.0
def write_ranking(path: Path, uids: np.ndarray, top_items: np.ndarray) -> None:
n, k = top_items.shape
out = pd.DataFrame(
{
"uid": np.repeat(uids.astype(np.int64), k),
"item_id": top_items.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)
out.to_parquet(path, index=False)
log(f"wrote {path} rows={len(out)}")
def train_model(
data: RankData,
tensors: dict[str, torch.Tensor],
device: torch.device,
eval_ctx: dict | None,
epochs: int,
patience: int,
) -> tuple[TwoTower, int, float]:
model = TwoTower(
n_items=data.n_items,
n_catalog=int(data.catalog_items.shape[0]),
dim=EMBED_DIM,
hidden=HIDDEN,
n_item_feat=data.item_feat.shape[1],
n_user_dense=N_USER_DENSE,
).to(device)
opt = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=1e-5)
sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=max(epochs, 1))
loader = make_loader(tensors, shuffle=True)
feat_table = torch.from_numpy(data.item_feat).to(device)
catalog_idx = torch.from_numpy(data.catalog_idx).to(device)
best_state = None
best_score = -1.0
best_epoch = 0
bad = 0
for epoch in range(1, epochs + 1):
model.train()
t0 = time.time()
losses = []
for raw in loader:
batch = batch_from_tuple(raw)
batch = {k: v.to(device) for k, v in batch.items()}
user, pos = model(batch, feat_table)
extra_ids = catalog_idx[torch.randint(0, catalog_idx.size(0), (EXTRA_NEGS,), device=device)]
extra = model.encode_items(extra_ids, feat_table[extra_ids])
loss = sampled_softmax(user, pos, extra, batch["weight"], SCALE)
opt.zero_grad(set_to_none=True)
loss.backward()
nn.utils.clip_grad_norm_(model.parameters(), 2.0)
opt.step()
losses.append(float(loss.item()))
sched.step()
mean_loss = float(np.mean(losses)) if losses else 0.0
msg = f"epoch {epoch} loss={mean_loss:.4f} time={time.time()-t0:.1f}s"
score = -1.0
if eval_ctx is not None:
user_mat = encode_all_users(
model,
data,
eval_ctx["uids"],
eval_ctx["cutoff"],
eval_ctx["pref"],
device,
)
item_mat = encode_catalog_items(model, data, device)
top = topk_rankings(user_mat, item_mat, eval_ctx["masks"], eval_ctx["items"])
score = score_ndcg(top, eval_ctx["uids"], eval_ctx["graded"])
msg += f" val_ndcg@10={score:.6f}"
if score > best_score + 1e-5:
best_score = score
best_epoch = epoch
best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
bad = 0
eval_ctx["best_top"] = top
else:
bad += 1
else:
best_epoch = epoch
best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
best_score = mean_loss
log(msg)
if eval_ctx is not None and bad >= patience:
log(f"early stop at epoch {epoch}, best={best_score:.6f} @ {best_epoch}")
break
if best_state is not None:
model.load_state_dict(best_state)
return model, best_epoch, best_score
def run_val(device: torch.device) -> tuple[int, float]:
log("=== validation-stage training (train history only) ===")
data = prepare_data(include_val=False)
folds = fold_cutoffs(int(TRAIN_END), N_FOLDS_VAL)
log(f"folds={folds}")
tensors = build_examples(data, folds, SEED)
el_uids, el_items, masks = load_eligibility(DATA / "validation_eligible_candidates.npz")
assert np.array_equal(el_items, data.catalog_items)
graded = validation_graded(el_uids, el_items, masks)
log(f"val scored users={len(graded)}")
eval_ctx = {
"uids": el_uids,
"items": el_items,
"masks": masks,
"graded": graded,
"cutoff": int(TRAIN_END),
"pref": DATA / "validation_preference_state.parquet",
"best_top": None,
}
model, best_epoch, best_score = train_model(
data, tensors, device, eval_ctx, MAX_EPOCHS, PATIENCE
)
top = eval_ctx["best_top"]
if top is None:
user_mat = encode_all_users(
model, data, el_uids, int(TRAIN_END), eval_ctx["pref"], device
)
item_mat = encode_catalog_items(model, data, device)
top = topk_rankings(user_mat, item_mat, masks, el_items)
best_score = score_ndcg(top, el_uids, graded)
write_ranking(OUT / "validation.parquet", el_uids, top)
ckpt = {
"state": model.state_dict(),
"best_epoch": best_epoch,
"best_score": best_score,
"n_items": data.n_items,
"item_feat": data.item_feat,
"idx_to_item": data.idx_to_item,
}
torch.save(ckpt, OUT / "val_model.pt")
with open(OUT / "val_metrics.json", "w") as f:
json.dump({"quality_ndcg@10": best_score, "best_epoch": best_epoch}, f, indent=2)
log(f"validation NDCG@10={best_score:.6f} best_epoch={best_epoch}")
return best_epoch, best_score
def run_test(device: torch.device, epochs: int) -> None:
log("=== hidden-test training (train+val history) ===")
data = prepare_data(include_val=True)
# Folds that stay inside the combined event stream. The most recent fold
# uses the public validation window as labels with train-only history.
val_end = int(TRAIN_END + GAP + WEEK)
folds = fold_cutoffs(val_end, N_FOLDS_VAL + 1)
log(f"test folds={folds}")
tensors = build_examples(data, folds, SEED + 1)
model, _, _ = train_model(data, tensors, device, None, epochs, patience=99)
el_uids, el_items, masks = load_eligibility(DATA / "test_eligible_candidates.npz")
assert np.array_equal(el_items, data.catalog_items)
cutoff = val_end
user_mat = encode_all_users(
model,
data,
el_uids,
cutoff,
DATA / "test_preference_state.parquet",
device,
)
item_mat = encode_catalog_items(model, data, device)
top = topk_rankings(user_mat, item_mat, masks, el_items)
write_ranking(OUT / "submission.parquet", el_uids, top)
torch.save(
{"state": model.state_dict(), "n_items": data.n_items, "epochs": epochs},
OUT / "test_model.pt",
)
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--stage", choices=["val", "test", "both"], default="both")
parser.add_argument("--epochs", type=int, default=None)
args = parser.parse_args()
OUT.mkdir(parents=True, exist_ok=True)
set_seed(SEED)
torch.set_num_threads(8)
device = torch.device("cpu")
log(f"torch={torch.__version__} threads={torch.get_num_threads()}")
best_epoch = 6
if args.stage in ("val", "both"):
best_epoch, _ = run_val(device)
if args.stage in ("test", "both"):
epochs = args.epochs if args.epochs is not None else max(best_epoch, 4) + 1
run_test(device, epochs=epochs)
if __name__ == "__main__":
main()