Agent trace

{
  "todos": [
    {
      "id": "explore",
      "content": "Explore data, evaluator, constraints, and environment",
      "status": "in_progress"
    },
    {
      "id": "design",
      "content": "Design two-tower architecture and training pipeline",
      "status": "pending"
    },
    {
      "id": "train_val",
      "content": "Train on train_* and evaluate public validation NDCG@10",
      "status": "pending"
    },
    {
      "id": "refit",
      "content": "Refit on train+val and produce test submission",
      "status": "pending"
    },
    {
      "id": "docs",
      "content": "Write README and keep best submission saved",
      "status": "pending"
    }
  ],
  "merge": false
}

{
  "target_directory": "/task"
}

{
  "target_directory": "/app"
}

{
  "target_file": "/task/requirements.txt"
}

{
  "target_file": "/task/evaluate_public.py"
}

{
  "target_file": "/task/ranking_evaluation.py"
}

{
  "target_directory": "/app/solution"
}

python3 - <<'PY'
import pandas as pd
import numpy as np
from pathlib import Path

data = Path('/task/data')

# catalog
cat = pd.read_parquet(data/'candidate_catalog.parquet')
print('=== candidate_catalog ===')
print(cat.head())
print(cat.dtypes)
print('shape', cat.shape)
print('cols', cat.columns.tolist())
print()

# target users
tu = pd.read_parquet(data/'target_users.parquet')
print('=== target_users ===')
print(tu.head())
print(tu.dtypes)
print('shape', tu.shape)
print('cols', tu.columns.tolist())
print()

# preference states
for name in ['validation_preference_state.parquet', 'test_preference_state.parquet']:
    df = pd.read_parquet(data/name)
    print(f'=== {name} ===')
    print(df.head())
    print(df.dtypes)
    print('shape', df.shape)
    print('cols', df.columns.tolist())
    if 'state' in df.columns:
        print(df['state'].value_counts())
    print()

# eligible
for name in ['validation_eligible_candidates.npz', 'test_eligible_candidates.npz']:
    z = np.load(data/name)
    print(f'=== {name} ===')
    print('keys', list(z.keys()))
    for k in z.keys():
        a = z[k]
        print(f'  {k}: shape={a.shape} dtype={a.dtype}')
    print()

# events
for fam in ['listens','likes','dislikes','unlikes','undislikes']:
    for split in ['train','val']:
        p = data/f'{split}_{fam}.parquet'
        df = pd.read_parquet(p)
        print(f'=== {split}_{fam} === shape={df.shape} cols={df.columns.tolist()}')
        print(df.dtypes)
        print(df.head(2))
        print('ts min/max', df['timestamp'].min(), df['timestamp'].max())
        print('n uid', df['uid'].nunique(), 'n item', df['item_id'].nunique())
        print()
PY

which python python3; ls /usr/bin/python*; echo '---'; python3 -c 'import sys; print(sys.executable, sys.path)'; echo '---'; ls /usr/local/bin 2>/dev/null; ls /opt 2>/dev/null; find / -name 'pandas' -type d 2>/dev/null | head; echo PATH=$PATH

/opt/task-venv/bin/python - <<'PY'
import pandas as pd
import numpy as np
from pathlib import Path

data = Path('/task/data')

# catalog
cat = pd.read_parquet(data/'candidate_catalog.parquet')
print('=== candidate_catalog ===')
print(cat.head())
print(cat.dtypes)
print('shape', cat.shape)
print('cols', cat.columns.tolist())
print()

# target users
tu = pd.read_parquet(data/'target_users.parquet')
print('=== target_users ===')
print(tu.head())
print(tu.dtypes)
print('shape', tu.shape)
print('cols', tu.columns.tolist())
print()

# preference states
for name in ['validation_preference_state.parquet', 'test_preference_state.parquet']:
    df = pd.read_parquet(data/name)
    print(f'=== {name} ===')
    print(df.head())
    print(df.dtypes)
    print('shape', df.shape)
    print('cols', df.columns.tolist())
    print(df.head(20))
    print()

# eligible
for name in ['validation_eligible_candidates.npz', 'test_eligible_candidates.npz']:
    z = np.load(data/name)
    print(f'=== {name} ===')
    print('keys', list(z.keys()))
    for k in z.keys():
        a = z[k]
        print(f'  {k}: shape={a.shape} dtype={a.dtype}')
    print('n_users', len(z['uids']), 'n_items', len(z['item_ids']))
    # eligibility density
    packed = z['packed_eligible']
    n_items = len(z['item_ids'])
    # unpack a few
    densities = []
    for r in range(min(20, packed.shape[0])):
        mask = np.unpackbits(packed[r], count=n_items, bitorder='little')
        densities.append(mask.sum())
    print('eligible counts sample (first 20 users)', densities)
    # all users
    counts = []
    for r in range(packed.shape[0]):
        mask = np.unpackbits(packed[r], count=n_items, bitorder='little')
        counts.append(int(mask.sum()))
    counts = np.array(counts)
    print('eligible stats: min', counts.min(), 'median', np.median(counts), 'mean', counts.mean(), 'max', counts.max())
    print()
PY

/opt/task-venv/bin/python - <<'PY'
import pandas as pd
import numpy as np
from pathlib import Path

data = Path('/task/data')

for fam in ['listens','likes','dislikes','unlikes','undislikes']:
    for split in ['train','val']:
        p = data/f'{split}_{fam}.parquet'
        df = pd.read_parquet(p)
        print(f'=== {split}_{fam} === shape={df.shape} cols={df.columns.tolist()}')
        print('ts min/max', int(df['timestamp'].min()), int(df['timestamp'].max()))
        print('n uid', df['uid'].nunique(), 'n item', df['item_id'].nunique())
        print('is_organic value counts:')
        print(df['is_organic'].value_counts())
        if 'played_ratio_pct' in df.columns:
            print('played_ratio_pct describe', df['played_ratio_pct'].describe().to_dict())
            print('>=50', (df.played_ratio_pct>=50).mean(), '>=80', (df.played_ratio_pct>=80).mean())
            print('organic & >=80', ((df.is_organic==1)&(df.played_ratio_pct>=80)).sum())
            print('organic & >=50', ((df.is_organic==1)&(df.played_ratio_pct>=50)).sum())
        print()
PY

/opt/task-venv/bin/python - <<'PY'
import pandas as pd
import numpy as np
from pathlib import Path

data = Path('/task/data')
cat = pd.read_parquet(data/'candidate_catalog.parquet')
catalog = set(cat.item_id.astype(np.int64))
tu = pd.read_parquet(data/'target_users.parquet')
targets = set(tu.uid.astype(np.int64))
print('n target', len(targets), 'n catalog', len(catalog))

# preference state overlap
pref = pd.read_parquet(data/'validation_preference_state.parquet')
print('pref states', pref.state.value_counts().to_dict())
print('pref unique users', pref.uid.nunique())
print('pref catalog overlap', pref.item_id.isin(catalog).mean())

# val labels vs history
val_likes = pd.read_parquet(data/'val_likes.parquet')
val_listens = pd.read_parquet(data/'val_listens.parquet')
ol = val_likes[(val_likes.uid.isin(targets)) & (val_likes.is_organic==1)]
print('val organic likes', len(ol), 'users', ol.uid.nunique(), 'catalog frac', ol.item_id.isin(catalog).mean())

sl = val_listens[(val_listens.uid.isin(targets)) & (val_listens.is_organic==1) & (val_listens.played_ratio_pct>=80)]
print('val organic strong listens', len(sl), 'users', sl.uid.nunique(), 'catalog frac', sl.item_id.isin(catalog).mean())

# How many val positives are in catalog
print('unique val organic like pairs catalog', ol[ol.item_id.isin(catalog)][['uid','item_id']].drop_duplicates().shape)
print('unique val strong listen pairs catalog', sl[sl.item_id.isin(catalog)][['uid','item_id']].drop_duplicates().shape)

# train listens sample for history stats - only target users, catalog items, positive
# memory: 43M rows is ok
print('loading train listens...')
cols = ['uid','timestamp','item_id','is_organic','played_ratio_pct']
# just dtypes and memory
import os
print('train_listens size MB', os.path.getsize(data/'train_listens.parquet')/1e6)
print('n cores')
import os
print(os.cpu_count())
import torch
print('torch', torch.__version__, 'threads', torch.get_num_threads())
PY

/opt/task-venv/bin/python - <<'PY'
import pandas as pd
import numpy as np
from pathlib import Path

data = Path('/task/data')
cat = pd.read_parquet(data/'candidate_catalog.parquet')
catalog = set(cat.item_id.astype(np.int64))
tu = pd.read_parquet(data/'target_users.parquet')
targets = set(tu.uid.astype(np.int64))

# load train listens with usecols
print('reading listens')
lst = pd.read_parquet(data/'train_listens.parquet', columns=['uid','item_id','is_organic','played_ratio_pct','timestamp'])
print('done', lst.shape, lst.memory_usage(deep=True).sum()/1e6)

# target user history
tmask = lst.uid.isin(targets)
print('target user listens', tmask.sum())
print('unique items in target listens', lst.loc[tmask,'item_id'].nunique())
pos = tmask & (lst.played_ratio_pct>=50)
print('target pos listens', pos.sum(), 'unique items', lst.loc[pos,'item_id'].nunique())
org80 = tmask & (lst.is_organic==1) & (lst.played_ratio_pct>=80)
print('target org80', org80.sum(), 'unique items', lst.loc[org80,'item_id'].nunique())
print('org80 catalog frac', lst.loc[org80,'item_id'].isin(catalog).mean())

# history length per user
h = lst.loc[org80].groupby('uid').size()
print('org80 hist len', h.describe().to_dict())

h2 = lst.loc[pos].groupby('uid').size()
print('pos hist len', h2.describe().to_dict())

# recency: last 90 days of train cutoff 24786800
cut = 24786800
day = 86400
recent = org80 & (lst.timestamp >= cut - 90*day)
print('recent 90d org80', recent.sum(), 'unique items', lst.loc[recent,'item_id'].nunique())
print('recent catalog frac', lst.loc[recent,'item_id'].isin(catalog).mean())

# how often val positives already in user train history
val_listens = pd.read_parquet(data/'val_listens.parquet', columns=['uid','item_id','is_organic','played_ratio_pct'])
sl = val_listens[(val_listens.uid.isin(targets)) & (val_listens.is_organic==1) & (val_listens.played_ratio_pct>=80)]
sl = sl[sl.item_id.isin(catalog)][['uid','item_id']].drop_duplicates()
hist_pairs = set(zip(lst.loc[pos,'uid'].to_numpy(), lst.loc[pos,'item_id'].to_numpy()))
# too big? pos is maybe 10M
print('pos pairs estimate', pos.sum())
PY

/opt/task-venv/bin/python - <<'PY'
import pandas as pd
import numpy as np
from pathlib import Path
from ranking_evaluation import CandidateSet, ndcg_at_k, score_rankings, TOP_N, QUALITY_K
import sys
sys.path.insert(0,'/task')
from ranking_evaluation import CandidateSet, ndcg_at_k, score_rankings

data = Path('/task/data')
cat = pd.read_parquet(data/'candidate_catalog.parquet')
catalog_items = cat.item_id.astype(np.int64).to_numpy()
pop_rank = cat.popularity_rank.to_numpy()  # 1 = most popular

tu = pd.read_parquet(data/'target_users.parquet')
targets = set(tu.uid.astype(np.int64))
eligibility = CandidateSet(data/'validation_eligible_candidates.npz')

# graded labels like evaluate_public
LISTEN_STRONG=80
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()
graded = {}
for uid, item in strong.itertuples(index=False, name=None):
    uid, item = int(uid), int(item)
    if eligibility.is_eligible(uid, item):
        graded.setdefault(uid, {})[item] = 1
for uid, item in organic_likes.itertuples(index=False, name=None):
    uid, item = int(uid), int(item)
    if eligibility.is_eligible(uid, item):
        graded.setdefault(uid, {})[item] = 3
print('scored users', len(graded))
nrel = [len(g) for g in graded.values()]
print('nrel per scored user', pd.Series(nrel).describe().to_dict())
grades = np.concatenate([np.array(list(g.values())) for g in graded.values()])
print('grade counts', {1: int((grades==1).sum()), 3: int((grades==3).sum())})

# popularity ranking among eligible
# precompute item order by popularity (catalog is already ordered by popularity)
# for each user, take first 100 eligible by popularity
item_ids = eligibility.item_ids
# catalog order should match popularity
print('item_ids vs catalog first 5', item_ids[:5], catalog_items[:5])
print('equal catalog order', np.array_equal(item_ids, catalog_items))

packed = eligibility.packed
n_items = len(item_ids)
ranked = {}
# vectorized unpack
for r, uid in enumerate(eligibility.uids):
    mask = np.unpackbits(packed[r], count=n_items, bitorder='little').astype(bool)
    # items already in popularity order if item_ids is catalog order
    elig = item_ids[mask]
    ranked[int(uid)] = elig[:100].tolist()

print(score_rankings(ranked, graded))

# overlap of labels with liked state
pref = pd.read_parquet(data/'validation_preference_state.parquet')
liked = pref[pref.state=='liked'][['uid','item_id']]
liked_set = set(zip(liked.uid.astype(int), liked.item_id.astype(int)))
n_like_hit=0; n_listen_hit=0; n_like=0; n_listen=0
for uid, g in graded.items():
    for item, gr in g.items():
        if gr==3:
            n_like += 1
            if (uid,item) in liked_set:
                n_like_hit += 1
        else:
            n_listen += 1
            if (uid,item) in liked_set:
                n_listen_hit += 1
print('grade3 already liked', n_like_hit, '/', n_like)
print('grade1 already liked', n_listen_hit, '/', n_listen)

# disliked overlap
dis = pref[pref.state=='disliked'][['uid','item_id']]
dis_set = set(zip(dis.uid.astype(int), dis.item_id.astype(int)))
n_dis=0
for uid,g in graded.items():
    for item in g:
        if (uid,item) in dis_set:
            n_dis += 1
print('relevant but disliked', n_dis)
PY

/opt/task-venv/bin/python - <<'PY'
import sys
sys.path.insert(0,'/task')
import pandas as pd
import numpy as np
from pathlib import Path
from ranking_evaluation import CandidateSet, ndcg_at_k, score_rankings

data = Path('/task/data')
cat = pd.read_parquet(data/'candidate_catalog.parquet')
catalog_items = cat.item_id.astype(np.int64).to_numpy()

tu = pd.read_parquet(data/'target_users.parquet')
targets = set(tu.uid.astype(np.int64))
eligibility = CandidateSet(data/'validation_eligible_candidates.npz')

LISTEN_STRONG=80
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()
graded = {}
for uid, item in strong.itertuples(index=False, name=None):
    uid, item = int(uid), int(item)
    if eligibility.is_eligible(uid, item):
        graded.setdefault(uid, {})[item] = 1
for uid, item in organic_likes.itertuples(index=False, name=None):
    uid, item = int(uid), int(item)
    if eligibility.is_eligible(uid, item):
        graded.setdefault(uid, {})[item] = 3
print('scored users', len(graded))
nrel = [len(g) for g in graded.values()]
print('nrel per scored user', pd.Series(nrel).describe().to_dict())
grades = np.concatenate([np.array(list(g.values())) for g in graded.values()])
print('grade counts', {1: int((grades==1).sum()), 3: int((grades==3).sum())})

item_ids = eligibility.item_ids
print('item_ids vs catalog first 5', item_ids[:5], catalog_items[:5])
print('equal catalog order', np.array_equal(item_ids, catalog_items))

packed = eligibility.packed
n_items = len(item_ids)
ranked = {}
for r, uid in enumerate(eligibility.uids):
    mask = np.unpackbits(packed[r], count=n_items, bitorder='little').astype(bool)
    elig = item_ids[mask]
    ranked[int(uid)] = elig[:100].tolist()

print('popularity', score_rankings(ranked, graded))

pref = pd.read_parquet(data/'validation_preference_state.parquet')
liked = pref[pref.state=='liked'][['uid','item_id']]
liked_set = set(zip(liked.uid.astype(int), liked.item_id.astype(int)))
n_like_hit=0; n_listen_hit=0; n_like=0; n_listen=0
for uid, g in graded.items():
    for item, gr in g.items():
        if gr==3:
            n_like += 1
            if (uid,item) in liked_set:
                n_like_hit += 1
        else:
            n_listen += 1
            if (uid,item) in liked_set:
                n_listen_hit += 1
print('grade3 already liked', n_like_hit, '/', n_like)
print('grade1 already liked', n_listen_hit, '/', n_listen)

dis = pref[pref.state=='disliked'][['uid','item_id']]
dis_set = set(zip(dis.uid.astype(int), dis.item_id.astype(int)))
n_dis=0
for uid,g in graded.items():
    for item in g:
        if (uid,item) in dis_set:
            n_dis += 1
print('relevant but disliked', n_dis)

# how many relevant items are ineligible
n_strong_all=0; n_strong_elig=0
for uid, item in strong.itertuples(index=False, name=None):
    uid,item=int(uid),int(item)
    n_strong_all += 1
    if eligibility.is_eligible(uid,item):
        n_strong_elig += 1
print('strong listen eligible among all unique pairs', n_strong_elig, 'wait this counts rows after drop dup but not catalog filter')
print('strong unique pairs', len(strong), 'likes unique', len(organic_likes))
PY

/opt/task-venv/bin/python - <<'PY'
import sys
sys.path.insert(0,'/task')
import pandas as pd
import numpy as np
from pathlib import Path
from ranking_evaluation import CandidateSet

data = Path('/task/data')
tu = pd.read_parquet(data/'target_users.parquet')
targets = set(tu.uid.astype(np.int64))
eligibility = CandidateSet(data/'validation_eligible_candidates.npz')
cat = pd.read_parquet(data/'candidate_catalog.parquet')
catalog = set(cat.item_id.astype(np.int64))

lst = pd.read_parquet(data/'train_listens.parquet', columns=['uid','item_id','played_ratio_pct'])
lst = lst[lst.uid.isin(targets)]
# catalog positives
pos = lst[(lst.played_ratio_pct>=50) & (lst.item_id.isin(catalog))]
# unique pairs
pos_pairs = set(zip(pos.uid.astype(int).to_numpy(), pos.item_id.astype(int).to_numpy()))
print('train pos catalog pairs', len(pos_pairs))

# any listen
any_pairs = set(zip(lst.uid.astype(int).to_numpy(), lst.item_id.astype(int).to_numpy()))
print('train any listen pairs', len(any_pairs))

val_listens = pd.read_parquet(data/'val_listens.parquet', columns=['uid','item_id','is_organic','played_ratio_pct'])
strong = val_listens[(val_listens.uid.isin(targets)) & (val_listens.is_organic==1) & (val_listens.played_ratio_pct>=80)][['uid','item_id']].drop_duplicates()

n_elig=n_in_pos=n_in_any=0
for uid,item in strong.itertuples(index=False,name=None):
    uid,item=int(uid),int(item)
    if eligibility.is_eligible(uid,item):
        n_elig += 1
        if (uid,item) in pos_pairs:
            n_in_pos += 1
        if (uid,item) in any_pairs:
            n_in_any += 1
print('eligible relevant listens', n_elig, 'in train pos', n_in_pos, 'in any train listen', n_in_any)

likes = pd.read_parquet(data/'val_likes.parquet')
ol = likes[(likes.uid.isin(targets))&(likes.is_organic==1)][['uid','item_id']].drop_duplicates()
n_elig=n_in_pos=n_in_any=0
for uid,item in ol.itertuples(index=False,name=None):
    uid,item=int(uid),int(item)
    if eligibility.is_eligible(uid,item):
        n_elig += 1
        if (uid,item) in pos_pairs:
            n_in_pos += 1
        if (uid,item) in any_pairs:
            n_in_any += 1
print('eligible relevant likes', n_elig, 'in train pos', n_in_pos, 'in any train listen', n_in_any)

# eligibility vs train history: fraction of eligible that were in train pos
# sample 50 users
rng = np.random.default_rng(0)
rows = rng.choice(len(eligibility.uids), size=50, replace=False)
item_ids = eligibility.item_ids
fracs=[]
for r in rows:
    uid=int(eligibility.uids[r])
    mask=np.unpackbits(eligibility.packed[r], count=len(item_ids), bitorder='little').astype(bool)
    elig_items=item_ids[mask]
    n_hist=sum((uid,int(i)) in pos_pairs for i in elig_items)
    fracs.append(n_hist/len(elig_items))
print('frac eligible that are train pos catalog (50 users)', float(np.mean(fracs)), 'max', float(np.max(fracs)))
PY

{
  "task_ids": [
    "call-ff5a5e5f-d78e-4a7c-b289-c88f0193bdbf-15"
  ],
  "timeout_ms": 120000
}

/opt/task-venv/bin/python - <<'PY'
import pandas as pd
import numpy as np
cat = pd.read_parquet('/task/data/candidate_catalog.parquet')
print(cat.describe())
print('min pos', cat.positive_listens.min(), 'max', cat.positive_listens.max())
print(cat.tail(3))
# track length from a sample of listens
import pyarrow.parquet as pq
# just check unique track length per item later
print('ok')
# how many train likes are catalog
likes = pd.read_parquet('/task/data/train_likes.parquet', columns=['item_id'])
print('train likes catalog frac', likes.item_id.isin(set(cat.item_id)).mean(), 'n', len(likes))
PY

{
  "todos": [
    {
      "id": "explore",
      "status": "completed"
    },
    {
      "id": "design",
      "status": "in_progress"
    }
  ],
  "merge": true
}

/opt/task-venv/bin/python -c "import torch; print(torch.__version__); print(torch.get_num_threads()); import os; print('cpu', os.cpu_count())"

/app/solution/model.py

"""Two-tower PyTorch ranker: user history encoder and independent item tower."""

from __future__ import annotations

import math

import torch
import torch.nn as nn
import torch.nn.functional as F


class AttnPool(nn.Module):
    """Learned query attention over a padded sequence, with recency bias."""

    def __init__(self, d_model: int):
        super().__init__()
        self.query = nn.Parameter(torch.zeros(d_model))
        nn.init.normal_(self.query, std=0.02)
        self.recency_scale = nn.Parameter(torch.tensor(0.4))
        self.scale = d_model**-0.5

    def forward(self, x: torch.Tensor, age_days: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        scores = (x * self.query).sum(-1) * self.scale
        recency = -F.softplus(self.recency_scale) * torch.log1p(age_days.clamp(min=0.0))
        scores = scores + recency
        scores = scores.masked_fill(~mask, torch.finfo(scores.dtype).min)
        empty = ~mask.any(dim=1)
        scores = scores.masked_fill(empty.unsqueeze(1), 0.0)
        weights = torch.softmax(scores, dim=1) * mask.float()
        denom = weights.sum(dim=1, keepdim=True).clamp(min=1e-6)
        weights = weights / denom
        return (weights.unsqueeze(-1) * x).sum(dim=1)


class ItemTower(nn.Module):
    """Maps a catalog item to an L2-normalized vector from ID + static features."""

    def __init__(self, embeddings: nn.Embedding, n_items: int, n_feat: int, d_model: int, dropout: float):
        super().__init__()
        self.embeddings = embeddings
        self.n_items = n_items
        self.feat_proj = nn.Linear(n_feat, d_model)
        self.net = nn.Sequential(
            nn.Linear(d_model * 2, d_model),
            nn.ReLU(inplace=True),
            nn.Dropout(dropout),
            nn.Linear(d_model, d_model),
        )

    def encode(self, catalog_idx: torch.Tensor, item_features: torch.Tensor) -> torch.Tensor:
        emb = self.embeddings(catalog_idx + 1)
        feat = self.feat_proj(item_features[catalog_idx])
        hidden = self.net(torch.cat([emb, feat], dim=-1))
        return F.normalize(hidden + emb, dim=-1)

    def all_vectors(self, item_features: torch.Tensor) -> torch.Tensor:
        idx = torch.arange(self.n_items, device=item_features.device)
        return self.encode(idx, item_features)


class UserTower(nn.Module):
    """Builds a user vector from interaction-history channels (no user-ID lookup)."""

    def __init__(self, embeddings: nn.Embedding, d_model: int, n_stats: int, dropout: float):
        super().__init__()
        self.embeddings = embeddings
        self.recent_pool = AttnPool(d_model)
        self.long_pool = AttnPool(d_model)
        self.like_pool = AttnPool(d_model)
        self.dislike_pool = AttnPool(d_model)
        self.oov_pool = AttnPool(d_model)
        in_dim = d_model * 5 + n_stats
        self.mlp = nn.Sequential(
            nn.Linear(in_dim, d_model * 2),
            nn.LayerNorm(d_model * 2),
            nn.ReLU(inplace=True),
            nn.Dropout(dropout),
            nn.Linear(d_model * 2, d_model),
            nn.LayerNorm(d_model),
            nn.ReLU(inplace=True),
            nn.Linear(d_model, d_model),
        )

    def _pool(self, pool: AttnPool, ids: torch.Tensor, age: torch.Tensor) -> torch.Tensor:
        mask = ids > 0
        x = self.embeddings(ids)
        return pool(x, age, mask)

    def forward(self, batch: dict[str, torch.Tensor]) -> torch.Tensor:
        parts = [
            self._pool(self.recent_pool, batch["recent_ids"], batch["recent_age"]),
            self._pool(self.long_pool, batch["long_ids"], batch["long_age"]),
            self._pool(self.like_pool, batch["like_ids"], batch["like_age"]),
            self._pool(self.dislike_pool, batch["dislike_ids"], batch["dislike_age"]),
            self._pool(self.oov_pool, batch["oov_ids"], batch["oov_age"]),
            batch["stats"],
        ]
        return F.normalize(self.mlp(torch.cat(parts, dim=-1)), dim=-1)


class TwoTower(nn.Module):
    """Dot-product two-tower model with a shared item-ID embedding table."""

    def __init__(
        self,
        vocab_size: int,
        n_items: int,
        n_feat: int,
        d_model: int = 64,
        n_stats: int = 8,
        dropout: float = 0.1,
    ):
        super().__init__()
        self.n_items = n_items
        self.d_model = d_model
        self.item_emb = nn.Embedding(vocab_size, d_model, padding_idx=0)
        nn.init.normal_(self.item_emb.weight, std=0.02)
        self.item_emb.weight.data[0].zero_()
        self.item_tower = ItemTower(self.item_emb, n_items, n_feat, d_model, dropout)
        self.user_tower = UserTower(self.item_emb, d_model, n_stats, dropout)
        self.logit_scale = nn.Parameter(torch.tensor(math.log(1.0 / 0.07)))
        self.register_buffer("item_features", torch.zeros(n_items, n_feat), persistent=True)

    def set_item_features(self, features: torch.Tensor) -> None:
        self.item_features.copy_(features)

    def user_vectors(self, batch: dict[str, torch.Tensor]) -> torch.Tensor:
        return self.user_tower(batch)

    def item_vectors(self, catalog_idx: torch.Tensor) -> torch.Tensor:
        return self.item_tower.encode(catalog_idx, self.item_features)

    def all_item_vectors(self) -> torch.Tensor:
        return self.item_tower.all_vectors(self.item_features)

    def score(self, user_vec: torch.Tensor, item_vec: torch.Tensor) -> torch.Tensor:
        return (user_vec * item_vec).sum(-1)

    def sampled_logits(
        self,
        user_vec: torch.Tensor,
        target_idx: torch.Tensor,
        extra_neg_idx: torch.Tensor,
        log_q: torch.Tensor | None = None,
    ) -> torch.Tensor:
        pos_vec = self.item_vectors(target_idx)
        neg_vec = self.item_vectors(extra_neg_idx)
        scale = self.logit_scale.exp().clamp(max=50.0)
        logits_pos = (user_vec * pos_vec).sum(-1, keepdim=True)
        logits_inbatch = user_vec @ pos_vec.t()
        same = target_idx.unsqueeze(1) == target_idx.unsqueeze(0)
        same = same.fill_diagonal_(False)
        logits_inbatch = logits_inbatch.masked_fill(same, torch.finfo(logits_inbatch.dtype).min)
        logits_inbatch = logits_inbatch.masked_fill(
            torch.eye(target_idx.size(0), device=target_idx.device, dtype=torch.bool),
            torch.finfo(logits_inbatch.dtype).min,
        )
        logits_neg = user_vec @ neg_vec.t()
        hit = extra_neg_idx.unsqueeze(0) == target_idx.unsqueeze(1)
        logits_neg = logits_neg.masked_fill(hit, torch.finfo(logits_neg.dtype).min)
        logits = torch.cat([logits_pos, logits_inbatch, logits_neg], dim=1) * scale
        if log_q is not None:
            cand = torch.cat(
                [target_idx.unsqueeze(1), target_idx.unsqueeze(0).expand(target_idx.size(0), -1), extra_neg_idx.unsqueeze(0).expand(target_idx.size(0), -1)],
                dim=1,
            )
            logits = logits - log_q[cand]
        return logits

/app/solution/train_two_tower.py

#!/usr/bin/env python3
"""Train a CPU PyTorch two-tower ranker and write validation + test rankings."""

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 torch
import torch.nn.functional as F
from torch.optim import AdamW
from torch.optim.lr_scheduler import CosineAnnealingLR

os.environ.setdefault("OMP_NUM_THREADS", "8")
os.environ.setdefault("MKL_NUM_THREADS", "8")
torch.set_num_threads(8)

from model import TwoTower

sys.path.insert(0, "/task")
from ranking_evaluation import CandidateSet, score_rankings

DATA = Path("/task/data")
SOLUTION = Path("/app/solution")
SOLUTION.mkdir(parents=True, exist_ok=True)

TRAIN_CUTOFF = 24_786_800
DAY = 86_400
WINDOW = 7 * DAY
N_HASH = 8192
N_STATS = 8
D_MODEL = 64
RECENT_DAYS = 21.0
L_RECENT = 40
L_LONG = 128
L_LIKE = 64
L_DISLIKE = 32
L_OOV = 40
MIN_HIST = 8
LISTEN_POS = 50
LISTEN_STRONG = 80
N_NEG = 512
BATCH_SIZE = 256
MAX_EXAMPLES_VAL = 360_000
MAX_EXAMPLES_REFIT = 480_000
SEED = 42

ITEM_LOOKUP_SIZE = 10_000_000


def log(msg: str) -> None:
    print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True)


@dataclass
class UserRec:
    all_item: np.ndarray
    all_ts: np.ndarray
    org_item: np.ndarray
    org_ts: np.ndarray
    org_like: np.ndarray
    pos_idx: np.ndarray
    pos_ts: np.ndarray
    pos_org: np.ndarray
    pos_ratio: np.ndarray
    strong_idx: np.ndarray
    strong_ts: np.ndarray
    oov_idx: np.ndarray
    oov_ts: np.ndarray
    like_idx: np.ndarray
    like_ts: np.ndarray
    dislike_idx: np.ndarray
    dislike_ts: np.ndarray


def _sort_uid_ts(df: pd.DataFrame) -> pd.DataFrame:
    return df.sort_values(["uid", "timestamp"], kind="mergesort")


def _partition(uids: np.ndarray, **cols: np.ndarray) -> dict[int, dict[str, np.ndarray]]:
    n = int(uids.shape[0])
    if n == 0:
        return {}
    breaks = np.flatnonzero(uids[1:] != uids[:-1]) + 1
    starts = np.concatenate(([0], breaks))
    ends = np.concatenate((breaks, [n]))
    out: dict[int, dict[str, np.ndarray]] = {}
    for a, b in zip(starts.tolist(), ends.tolist()):
        uid = int(uids[a])
        out[uid] = {k: v[a:b] for k, v in cols.items()}
    return out


def _empty_i64() -> np.ndarray:
    return np.zeros(0, dtype=np.int64)


def _empty_i8() -> np.ndarray:
    return np.zeros(0, dtype=np.int8)


def _pad_last(ids: np.ndarray, ages: np.ndarray, length: int) -> tuple[np.ndarray, np.ndarray]:
    out_ids = np.zeros(length, dtype=np.int64)
    out_age = np.zeros(length, dtype=np.float32)
    n = int(ids.shape[0])
    if n <= 0:
        return out_ids, out_age
    sl = min(n, length)
    out_ids[-sl:] = ids[-sl:]
    out_age[-sl:] = ages[-sl:].astype(np.float32, copy=False)
    return out_ids, out_age


def _age_days(ts: np.ndarray, cutoff: int) -> np.ndarray:
    return (cutoff - ts).astype(np.float64) / DAY


def _stats(rec: UserRec, cutoff: int) -> np.ndarray:
    i = int(np.searchsorted(rec.pos_ts, cutoff))
    j = int(np.searchsorted(rec.like_ts, cutoff))
    k = int(np.searchsorted(rec.dislike_ts, cutoff))
    pos_ts = rec.pos_ts[:i]
    n_pos = i
    n_like = j
    n_dis = k
    if n_pos:
        unique = np.unique(rec.pos_idx[:i]).size
        org_frac = float(rec.pos_org[:i].mean())
        mean_ratio = float(rec.pos_ratio[:i].mean()) / 100.0
        first_age = float(cutoff - int(pos_ts[0])) / DAY
        last_age = float(cutoff - int(pos_ts[-1])) / DAY
    else:
        unique = 0
        org_frac = 0.0
        mean_ratio = 0.0
        first_age = 0.0
        last_age = 0.0
    return np.array(
        [
            math.log1p(n_pos) / 10.0,
            math.log1p(n_like) / 5.0,
            math.log1p(n_dis) / 5.0,
            math.log1p(unique) / 10.0,
            org_frac,
            mean_ratio,
            math.log1p(max(first_age, 0.0)) / 6.0,
            math.log1p(max(last_age, 0.0)) / 4.0,
        ],
        dtype=np.float32,
    )


def load_catalog() -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    cat = pd.read_parquet(DATA / "candidate_catalog.parquet")
    item_ids = cat.item_id.to_numpy(dtype=np.int64)
    pos = cat.positive_listens.to_numpy(dtype=np.float32)
    org = cat.organic_positive_listens.to_numpy(dtype=np.float32)
    rank = cat.popularity_rank.to_numpy(dtype=np.float32)
    feat = np.stack(
        [
            np.log1p(pos) / 10.0,
            np.log1p(org) / 10.0,
            np.divide(org, pos, out=np.zeros_like(org), where=pos > 0),
            np.log(rank) / 10.0,
        ],
        axis=1,
    ).astype(np.float32)
    lookup = np.full(ITEM_LOOKUP_SIZE, -1, dtype=np.int32)
    valid = (item_ids >= 0) & (item_ids < ITEM_LOOKUP_SIZE)
    lookup[item_ids[valid]] = np.arange(len(item_ids), dtype=np.int32)[valid]
    return item_ids, feat, lookup


def _hash_oov(raw: np.ndarray, n_items: int) -> np.ndarray:
    return n_items + 1 + (raw.astype(np.int64) * np.int64(2654435761) % N_HASH)


def _map_catalog(raw: np.ndarray, lookup: np.ndarray) -> np.ndarray:
    mapped = np.full(raw.shape[0], -1, dtype=np.int32)
    in_range = (raw >= 0) & (raw < lookup.shape[0])
    mapped[in_range] = lookup[raw[in_range]]
    return mapped


def _read_events(splits: list[str], family: str, columns: list[str]) -> pd.DataFrame:
    frames = [pd.read_parquet(DATA / f"{split}_{family}.parquet", columns=columns) for split in splits]
    return pd.concat(frames, ignore_index=True) if len(frames) > 1 else frames[0]


def build_user_recs(splits: list[str], lookup: np.ndarray, n_items: int) -> dict[int, UserRec]:
    log(f"loading events for splits={splits}")
    listens = _read_events(
        splits, "listens", ["uid", "timestamp", "item_id", "is_organic", "played_ratio_pct"]
    )
    likes = _read_events(splits, "likes", ["uid", "timestamp", "item_id", "is_organic"])
    dislikes = _read_events(splits, "dislikes", ["uid", "timestamp", "item_id"])
    unlikes = _read_events(splits, "unlikes", ["uid", "timestamp", "item_id"])
    undislikes = _read_events(splits, "undislikes", ["uid", "timestamp", "item_id"])

    def arrays(df: pd.DataFrame) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
        return (
            df.uid.to_numpy(dtype=np.int64),
            df.timestamp.to_numpy(dtype=np.int64),
            df.item_id.to_numpy(dtype=np.int64),
        )

    uid_l, ts_l, item_l = arrays(listens)
    org_l = listens.is_organic.to_numpy(dtype=np.int8)
    ratio_l = listens.played_ratio_pct.to_numpy(dtype=np.int16)
    del listens

    cat_idx_l = _map_catalog(item_l, lookup)
    is_cat_l = cat_idx_l >= 0

    # All catalog interactions (listens of any completion + explicit actions).
    seen_uid = [uid_l[is_cat_l], ts_l[is_cat_l], cat_idx_l[is_cat_l].astype(np.int64)]
    for df in (likes, dislikes, unlikes, undislikes):
        u, t, it = arrays(df)
        mapped = _map_catalog(it, lookup)
        m = mapped >= 0
        seen_uid[0] = np.concatenate([seen_uid[0], u[m]])
        seen_uid[1] = np.concatenate([seen_uid[1], t[m]])
        seen_uid[2] = np.concatenate([seen_uid[2], mapped[m].astype(np.int64)])

    seen_order = np.lexsort((seen_uid[1], seen_uid[0]))
    seen_parts = _partition(
        seen_uid[0][seen_order],
        item=seen_uid[2][seen_order],
        ts=seen_uid[1][seen_order],
    )
    del seen_uid, seen_order

    pos_m = ratio_l >= LISTEN_POS
    pos_cat = pos_m & is_cat_l
    pos_oov = pos_m & ~is_cat_l
    pos_order = np.lexsort((ts_l[pos_cat], uid_l[pos_cat]))
    pos_parts = _partition(
        uid_l[pos_cat][pos_order],
        idx=(cat_idx_l[pos_cat].astype(np.int64)[pos_order] + 1),
        ts=ts_l[pos_cat][pos_order],
        org=org_l[pos_cat][pos_order],
        ratio=ratio_l[pos_cat][pos_order],
    )
    oov_order = np.lexsort((ts_l[pos_oov], uid_l[pos_oov]))
    oov_parts = _partition(
        uid_l[pos_oov][oov_order],
        idx=_hash_oov(item_l[pos_oov][oov_order], n_items),
        ts=ts_l[pos_oov][oov_order],
    )

    strong_m = is_cat_l & (org_l == 1) & (ratio_l >= LISTEN_STRONG)
    strong_order = np.lexsort((ts_l[strong_m], uid_l[strong_m]))
    strong_parts = _partition(
        uid_l[strong_m][strong_order],
        idx=(cat_idx_l[strong_m].astype(np.int64)[strong_order] + 1),
        ts=ts_l[strong_m][strong_order],
        item=cat_idx_l[strong_m].astype(np.int64)[strong_order],
    )

    like_u, like_ts, like_item = arrays(likes)
    like_org = likes.is_organic.to_numpy(dtype=np.int8)
    like_map = _map_catalog(like_item, lookup)
    like_hist_idx = np.where(like_map >= 0, like_map.astype(np.int64) + 1, _hash_oov(like_item, n_items))
    like_order = np.lexsort((like_ts, like_u))
    like_parts = _partition(
        like_u[like_order],
        idx=like_hist_idx[like_order],
        ts=like_ts[like_order],
        item=like_map[like_order].astype(np.int64),
        org=like_org[like_order],
    )

    org_like_m = (like_org == 1) & (like_map >= 0)
    org_like_order = np.lexsort((like_ts[org_like_m], like_u[org_like_m]))
    org_like_parts = _partition(
        like_u[org_like_m][org_like_order],
        item=like_map[org_like_m].astype(np.int64)[org_like_order],
        ts=like_ts[org_like_m][org_like_order],
    )

    dis_u, dis_ts, dis_item = arrays(dislikes)
    dis_map = _map_catalog(dis_item, lookup)
    dis_hist = np.where(dis_map >= 0, dis_map.astype(np.int64) + 1, _hash_oov(dis_item, n_items))
    dis_order = np.lexsort((dis_ts, dis_u))
    dis_parts = _partition(
        dis_u[dis_order],
        idx=dis_hist[dis_order],
        ts=dis_ts[dis_order],
    )

    del uid_l, ts_l, item_l, org_l, ratio_l, cat_idx_l, likes, dislikes, unlikes, undislikes

    uids = set(seen_parts) | set(pos_parts) | set(like_parts) | set(strong_parts)
    recs: dict[int, UserRec] = {}
    z64 = _empty_i64()
    z8 = _empty_i8()
    for uid in uids:
        seen = seen_parts.get(uid, {})
        pos = pos_parts.get(uid, {})
        oov = oov_parts.get(uid, {})
        strong = strong_parts.get(uid, {})
        like = like_parts.get(uid, {})
        dis = dis_parts.get(uid, {})
        ol = org_like_parts.get(uid, {})

        strong_item = strong.get("item", z64)
        strong_ts = strong.get("ts", z64)
        ol_item = ol.get("item", z64)
        ol_ts = ol.get("ts", z64)
        if strong_item.size and ol_item.size:
            org_item = np.concatenate([strong_item, ol_item])
            org_ts = np.concatenate([strong_ts, ol_ts])
            org_like_flag = np.concatenate(
                [np.zeros(strong_item.size, dtype=np.int8), np.ones(ol_item.size, dtype=np.int8)]
            )
            order = np.argsort(org_ts, kind="mergesort")
            org_item, org_ts, org_like_flag = org_item[order], org_ts[order], org_like_flag[order]
        elif ol_item.size:
            org_item, org_ts = ol_item, ol_ts
            org_like_flag = np.ones(ol_item.size, dtype=np.int8)
        else:
            org_item, org_ts = strong_item, strong_ts
            org_like_flag = np.zeros(strong_item.size, dtype=np.int8)

        recs[uid] = UserRec(
            all_item=seen.get("item", z64),
            all_ts=seen.get("ts", z64),
            org_item=org_item,
            org_ts=org_ts,
            org_like=org_like_flag,
            pos_idx=pos.get("idx", z64),
            pos_ts=pos.get("ts", z64),
            pos_org=pos.get("org", z8),
            pos_ratio=pos.get("ratio", z8),
            strong_idx=strong.get("idx", z64),
            strong_ts=strong_ts,
            oov_idx=oov.get("idx", z64),
            oov_ts=oov.get("ts", z64),
            like_idx=like.get("idx", z64),
            like_ts=like.get("ts", z64),
            dislike_idx=dis.get("idx", z64),
            dislike_ts=dis.get("ts", z64),
        )
    log(f"built user records for {len(recs)} users")
    return recs


def history_channels(
    rec: UserRec,
    cutoff: int,
    like_override: tuple[np.ndarray, np.ndarray] | None = None,
    dislike_override: tuple[np.ndarray, np.ndarray] | None = None,
) -> dict[str, np.ndarray]:
    def take(ids: np.ndarray, ts: np.ndarray, length: int, recent_only: bool = False) -> tuple[np.ndarray, np.ndarray]:
        end = int(np.searchsorted(ts, cutoff))
        ids, ts = ids[:end], ts[:end]
        if ids.size == 0:
            return _pad_last(ids, np.zeros(0, np.float32), length)
        age = _age_days(ts, cutoff)
        if recent_only:
            keep = age <= RECENT_DAYS
            ids, age = ids[keep], age[keep]
        return _pad_last(ids, age, length)

    recent_ids, recent_age = take(rec.strong_idx, rec.strong_ts, L_RECENT, recent_only=True)
    if int((recent_ids > 0).sum()) < 4:
        recent_ids, recent_age = take(rec.pos_idx, rec.pos_ts, L_RECENT, recent_only=True)
    long_ids, long_age = take(rec.pos_idx, rec.pos_ts, L_LONG, recent_only=False)
    oov_ids, oov_age = take(rec.oov_idx, rec.oov_ts, L_OOV, recent_only=False)

    if like_override is None:
        like_ids, like_age = take(rec.like_idx, rec.like_ts, L_LIKE)
    else:
        like_ids, like_age = _pad_last(like_override[0], like_override[1], L_LIKE)
    if dislike_override is None:
        dis_ids, dis_age = take(rec.dislike_idx, rec.dislike_ts, L_DISLIKE)
    else:
        dis_ids, dis_age = _pad_last(dislike_override[0], dislike_override[1], L_DISLIKE)

    return {
        "recent_ids": recent_ids,
        "recent_age": recent_age,
        "long_ids": long_ids,
        "long_age": long_age,
        "like_ids": like_ids,
        "like_age": like_age,
        "dislike_ids": dis_ids,
        "dislike_age": dis_age,
        "oov_ids": oov_ids,
        "oov_age": oov_age,
        "stats": _stats(rec, cutoff),
    }


def _seen_before(rec: UserRec, cutoff: int) -> np.ndarray:
    end = int(np.searchsorted(rec.all_ts, cutoff))
    if end == 0:
        return np.zeros(0, dtype=np.int64)
    return np.unique(rec.all_item[:end])


def collect_examples(
    recs: dict[int, UserRec],
    snapshots: list[int],
    max_examples: int,
    rng: np.random.Generator,
) -> dict[str, np.ndarray]:
    rows: list[dict[str, np.ndarray | float | int]] = []
    for snap in snapshots:
        win_end = snap + WINDOW
        for uid, rec in recs.items():
            if rec.pos_ts.size == 0:
                continue
            n_hist = int(np.searchsorted(rec.pos_ts, snap))
            if n_hist < MIN_HIST:
                continue
            seen = _seen_before(rec, snap)
            seen_set = set(seen.tolist()) if seen.size else set()
            lo = int(np.searchsorted(rec.org_ts, snap))
            hi = int(np.searchsorted(rec.org_ts, win_end))
            if hi <= lo:
                continue
            hist = history_channels(rec, snap)
            used: set[int] = set()
            for item, ts, is_like in zip(rec.org_item[lo:hi], rec.org_ts[lo:hi], rec.org_like[lo:hi]):
                item = int(item)
                if item in seen_set or item in used:
                    continue
                used.add(item)
                weight = 3.0 if int(is_like) == 1 else 1.0
                row = dict(hist)
                row["target"] = item
                row["weight"] = weight
                rows.append(row)
                seen_set.add(item)

    # Sequential next-discovery examples from later history.
    for uid, rec in recs.items():
        if rec.org_ts.size < MIN_HIST + 2:
            continue
        first_ts: dict[int, int] = {}
        for item, ts in zip(rec.all_item, rec.all_ts):
            item = int(item)
            if item not in first_ts:
                first_ts[item] = int(ts)
        discoveries: list[tuple[int, int, int]] = []
        seen_d: set[int] = set()
        for item, ts, is_like in zip(rec.org_item, rec.org_ts, rec.org_like):
            item = int(item)
            if item in seen_d:
                continue
            ft = first_ts.get(item, int(ts))
            if int(ts) - ft > WINDOW:
                continue
            seen_d.add(item)
            discoveries.append((ft, item, int(is_like)))
        if len(discoveries) < MIN_HIST + 2:
            continue
        start = max(MIN_HIST, int(0.55 * len(discoveries)))
        for i in range(start, len(discoveries)):
            if (i < len(discoveries) - 6) and (i % 4 != 0):
                continue
            snap, item, is_like = discoveries[i]
            n_hist = int(np.searchsorted(rec.pos_ts, snap))
            if n_hist < MIN_HIST:
                continue
            hist = history_channels(rec, snap)
            row = dict(hist)
            row["target"] = int(item)
            row["weight"] = 3.0 if is_like else 1.0
            rows.append(row)

    if not rows:
        raise RuntimeError("no training examples constructed")
    log(f"raw example count {len(rows)}")
    if len(rows) > max_examples:
        weights = np.array([float(r["weight"]) for r in rows], dtype=np.float64)
        p = weights / weights.sum()
        pick = rng.choice(len(rows), size=max_examples, replace=False, p=p)
        rows = [rows[i] for i in pick.tolist()]
        log(f"subsampled to {len(rows)} examples")

    keys = [
        "recent_ids",
        "recent_age",
        "long_ids",
        "long_age",
        "like_ids",
        "like_age",
        "dislike_ids",
        "dislike_age",
        "oov_ids",
        "oov_age",
        "stats",
        "target",
        "weight",
    ]
    out: dict[str, np.ndarray] = {}
    for key in keys:
        if key == "target":
            out[key] = np.asarray([int(r[key]) for r in rows], dtype=np.int64)
        elif key == "weight":
            out[key] = np.asarray([float(r[key]) for r in rows], dtype=np.float32)
        else:
            out[key] = np.stack([r[key] for r in rows], axis=0)
    log(
        f"example tensor shapes: "
        + ", ".join(f"{k}{tuple(v.shape)}" for k, v in out.items())
    )
    return out


def preference_overrides(
    path: Path,
    lookup: np.ndarray,
    n_items: int,
    recs: dict[int, UserRec],
    cutoff: int,
) -> dict[int, dict[str, tuple[np.ndarray, np.ndarray]]]:
    pref = pd.read_parquet(path)
    out: dict[int, dict[str, tuple[np.ndarray, np.ndarray]]] = {}
    for uid, grp in pref.groupby("uid"):
        uid = int(uid)
        rec = recs.get(uid)
        last_like: dict[int, int] = {}
        last_dis: dict[int, int] = {}
        if rec is not None:
            for idx, ts in zip(rec.like_idx, rec.like_ts):
                last_like[int(idx)] = int(ts)
            for idx, ts in zip(rec.dislike_idx, rec.dislike_ts):
                last_dis[int(idx)] = int(ts)
        likes_idx: list[int] = []
        likes_age: list[float] = []
        dis_idx: list[int] = []
        dis_age: list[float] = []
        for item, state in zip(grp.item_id.to_numpy(np.int64), grp.state.to_numpy()):
            mapped = int(lookup[item]) if 0 <= item < lookup.shape[0] else -1
            if mapped < 0:
                hist_idx = int(_hash_oov(np.array([item]), n_items)[0])
            else:
                hist_idx = mapped + 1
            if state == "liked":
                ts = last_like.get(hist_idx, cutoff - 14 * DAY)
                likes_idx.append(hist_idx)
                likes_age.append(max(0.0, (cutoff - ts) / DAY))
            elif state == "disliked":
                ts = last_dis.get(hist_idx, cutoff - 14 * DAY)
                dis_idx.append(hist_idx)
                dis_age.append(max(0.0, (cutoff - ts) / DAY))
        out[uid] = {
            "like": (np.asarray(likes_idx, dtype=np.int64), np.asarray(likes_age, dtype=np.float32)),
            "dislike": (np.asarray(dis_idx, dtype=np.int64), np.asarray(dis_age, dtype=np.float32)),
        }
    return out


def encode_users(
    model: TwoTower,
    recs: dict[int, UserRec],
    uids: np.ndarray,
    cutoff: int,
    overrides: dict[int, dict[str, tuple[np.ndarray, np.ndarray]]] | None,
    batch_size: int = 256,
) -> np.ndarray:
    model.eval()
    vectors = np.zeros((len(uids), model.d_model), dtype=np.float32)
    with torch.no_grad():
        for start in range(0, len(uids), batch_size):
            chunk = uids[start : start + batch_size]
            bag: dict[str, list[np.ndarray]] = {
                "recent_ids": [],
                "recent_age": [],
                "long_ids": [],
                "long_age": [],
                "like_ids": [],
                "like_age": [],
                "dislike_ids": [],
                "dislike_age": [],
                "oov_ids": [],
                "oov_age": [],
                "stats": [],
            }
            for uid in chunk:
                uid = int(uid)
                rec = recs.get(uid)
                if rec is None:
                    rec = UserRec(
                        all_item=_empty_i64(),
                        all_ts=_empty_i64(),
                        org_item=_empty_i64(),
                        org_ts=_empty_i64(),
                        org_like=_empty_i8(),
                        pos_idx=_empty_i64(),
                        pos_ts=_empty_i64(),
                        pos_org=_empty_i8(),
                        pos_ratio=_empty_i8(),
                        strong_idx=_empty_i64(),
                        strong_ts=_empty_i64(),
                        oov_idx=_empty_i64(),
                        oov_ts=_empty_i64(),
                        like_idx=_empty_i64(),
                        like_ts=_empty_i64(),
                        dislike_idx=_empty_i64(),
                        dislike_ts=_empty_i64(),
                    )
                like_ov = dis_ov = None
                if overrides is not None and uid in overrides:
                    like_ov = overrides[uid]["like"]
                    dis_ov = overrides[uid]["dislike"]
                hist = history_channels(rec, cutoff, like_ov, dis_ov)
                for k, v in hist.items():
                    bag[k].append(v)
            batch = {
                k: torch.from_numpy(np.stack(v, axis=0))
                for k, v in bag.items()
            }
            vec = model.user_vectors(batch)
            vectors[start : start + len(chunk)] = vec.cpu().numpy()
    return vectors


def rank_users(
    user_vecs: np.ndarray,
    item_vecs: np.ndarray,
    elig: np.ndarray,
    item_ids: np.ndarray,
    uids: np.ndarray,
    top_n: int = 100,
) -> pd.DataFrame:
    scores = user_vecs @ item_vecs.T
    scores = np.where(elig, scores, -1e9).astype(np.float32, copy=False)
    idx = np.argpartition(-scores, kth=top_n - 1, axis=1)[:, :top_n]
    part = np.take_along_axis(scores, idx, axis=1)
    order = np.argsort(-part, axis=1)
    top = np.take_along_axis(idx, order, axis=1)
    n_users = len(uids)
    out_uid = np.repeat(uids.astype(np.int64), top_n)
    out_item = item_ids[top.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 validation_graded(eligibility: CandidateSet, targets: set[int]) -> dict[int, dict[int, int]]:
    listens = pd.read_parquet(DATA / "val_listens.parquet")
    likes = pd.read_parquet(DATA / "val_likes.parquet")
    strong = listens[
        listens.uid.isin(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()
    graded: dict[int, dict[int, int]] = {}
    for uid, item in strong.itertuples(index=False, name=None):
        uid, item = int(uid), int(item)
        if eligibility.is_eligible(uid, item):
            graded.setdefault(uid, {})[item] = 1
    for uid, item in organic_likes.itertuples(index=False, name=None):
        uid, item = int(uid), int(item)
        if eligibility.is_eligible(uid, item):
            graded.setdefault(uid, {})[item] = 3
    return graded


def unpack_eligibility(path: Path, n_items: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    with np.load(path, allow_pickle=False) as art:
        uids = art["uids"].astype(np.int64, copy=True)
        item_ids = art["item_ids"].astype(np.int64, copy=True)
        packed = art["packed_eligible"]
    elig = np.unpackbits(packed, axis=1, bitorder="little")[:, :n_items].astype(bool)
    return uids, item_ids, elig


def tensors_from_examples(examples: dict[str, np.ndarray]) -> dict[str, torch.Tensor]:
    out = {}
    for k, v in examples.items():
        if k in {"recent_ids", "long_ids", "like_ids", "dislike_ids", "oov_ids", "target"}:
            out[k] = torch.from_numpy(np.ascontiguousarray(v, dtype=np.int64))
        else:
            out[k] = torch.from_numpy(np.ascontiguousarray(v, dtype=np.float32))
    return out


def iterate_batches(
    tensors: dict[str, torch.Tensor], batch_size: int, rng: torch.Generator
):
    n = tensors["target"].shape[0]
    perm = torch.randperm(n, generator=rng)
    for start in range(0, n, batch_size):
        idx = perm[start : start + batch_size]
        if idx.numel() < 8:
            continue
        yield {k: v[idx] for k, v in tensors.items()}


def train_model(
    model: TwoTower,
    examples: dict[str, np.ndarray],
    pop_probs: np.ndarray,
    log_q: torch.Tensor,
    epochs: int,
    lr: float,
    eval_fn,
    deadline: float,
    ckpt_path: Path,
) -> tuple[TwoTower, float, int]:
    tensors = tensors_from_examples(examples)
    n_items = model.n_items
    pop = torch.from_numpy(pop_probs.astype(np.float32))
    opt = AdamW(model.parameters(), lr=lr, weight_decay=1e-4)
    steps_per = max(1, examples["target"].shape[0] // BATCH_SIZE)
    sched = CosineAnnealingLR(opt, T_max=max(1, epochs * steps_per), eta_min=lr * 0.08)
    gen = torch.Generator()
    gen.manual_seed(SEED)
    best_ndcg = -1.0
    best_epoch = 0
    best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}

    for epoch in range(1, epochs + 1):
        if time.time() > deadline:
            log("hit time deadline during training")
            break
        model.train()
        losses = []
        t0 = time.time()
        n_steps = 0
        for batch in iterate_batches(tensors, BATCH_SIZE, gen):
            if time.time() > deadline:
                break
            user_vec = model.user_vectors(batch)
            extra = torch.multinomial(pop, N_NEG, replacement=True)
            logits = model.sampled_logits(user_vec, batch["target"], extra, log_q=log_q)
            labels = torch.zeros(logits.size(0), dtype=torch.long)
            loss = F.cross_entropy(logits, labels, reduction="none")
            loss = (loss * batch["weight"]).mean()
            opt.zero_grad(set_to_none=True)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
            sched.step()
            losses.append(float(loss.item()))
            n_steps += 1
        mean_loss = float(np.mean(losses)) if losses else float("nan")
        ndcg = eval_fn(model) if eval_fn is not None else 0.0
        log(
            f"epoch {epoch} steps={n_steps} loss={mean_loss:.4f} "
            f"ndcg@10={ndcg:.6f} time={time.time() - t0:.1f}s"
        )
        if ndcg >= best_ndcg:
            best_ndcg = ndcg
            best_epoch = epoch
            best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
            torch.save(best_state, ckpt_path)
            log(f"saved checkpoint ndcg={best_ndcg:.6f} -> {ckpt_path}")
        elif eval_fn is None:
            best_ndcg = mean_loss
            best_epoch = epoch
            best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
            torch.save(best_state, ckpt_path)
    model.load_state_dict(best_state)
    return model, float(best_ndcg), best_epoch


def snapshot_list(end_ts: int, n_windows: int) -> list[int]:
    snaps = []
    for w in range(n_windows, 0, -1):
        snap = end_ts - w * WINDOW
        if snap > DAY * 14:
            snaps.append(int(snap))
    return snaps


def write_ranking(df: pd.DataFrame, path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    df.to_parquet(path, index=False)
    log(f"wrote {path} rows={len(df)}")


def main() -> None:
    t_start = time.time()
    deadline = t_start + 54 * 60
    rng = np.random.default_rng(SEED)
    torch.manual_seed(SEED)
    np.random.seed(SEED)

    item_ids, item_feat, lookup = load_catalog()
    n_items = len(item_ids)
    vocab = 1 + n_items + N_HASH
    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.tolist()))

    val_elig_uids, val_elig_items, val_elig = unpack_eligibility(
        DATA / "validation_eligible_candidates.npz", n_items
    )
    test_elig_uids, test_elig_items, test_elig = unpack_eligibility(
        DATA / "test_eligible_candidates.npz", n_items
    )
    assert np.array_equal(val_elig_items, item_ids)
    assert np.array_equal(test_elig_items, item_ids)

    val_uid_row = {int(u): i for i, u in enumerate(val_elig_uids.tolist())}
    test_uid_row = {int(u): i for i, u in enumerate(test_elig_uids.tolist())}
    val_elig_aligned = np.stack([val_elig[val_uid_row[int(u)]] for u in target_uids], axis=0)
    test_elig_aligned = np.stack([test_elig[test_uid_row[int(u)]] for u in target_uids], axis=0)

    eligibility = CandidateSet(DATA / "validation_eligible_candidates.npz")
    graded = validation_graded(eligibility, target_set)
    log(f"validation scored users={len(graded)}")

    recs_train = build_user_recs(["train"], lookup, n_items)
    snaps = snapshot_list(TRAIN_CUTOFF, n_windows=8)
    log(f"val-stage snapshots {snaps}")
    examples = collect_examples(recs_train, snaps, MAX_EXAMPLES_VAL, rng)

    target_hist = np.bincount(examples["target"], minlength=n_items).astype(np.float64)
    pop = np.power(np.maximum(target_hist, 0.5), 0.75)
    pop_probs = pop / pop.sum()
    q = np.maximum(target_hist, 0.5)
    q = q / q.sum()
    log_q = torch.from_numpy(np.log(q).astype(np.float32))

    val_pref = preference_overrides(
        DATA / "validation_preference_state.parquet", lookup, n_items, recs_train, TRAIN_CUTOFF
    )

    def eval_val(model: TwoTower) -> float:
        model.eval()
        with torch.no_grad():
            item_vecs = model.all_item_vectors().cpu().numpy()
        user_vecs = encode_users(model, recs_train, target_uids, TRAIN_CUTOFF, val_pref)
        ranking = rank_users(user_vecs, item_vecs, val_elig_aligned, item_ids, target_uids)
        ranked = {
            int(uid): list(map(int, grp.sort_values("rank").item_id.tolist()))
            for uid, grp in ranking.groupby("uid")
        }
        scores = score_rankings(ranked, graded)
        return float(scores["quality_ndcg@10"])

    model = TwoTower(
        vocab_size=vocab,
        n_items=n_items,
        n_feat=item_feat.shape[1],
        d_model=D_MODEL,
        n_stats=N_STATS,
        dropout=0.1,
    )
    model.set_item_features(torch.from_numpy(item_feat))
    n_params = sum(p.numel() for p in model.parameters())
    log(f"model params={n_params} vocab={vocab} items={n_items}")

    val_deadline = min(deadline - 16 * 60, t_start + 28 * 60)
    ckpt = SOLUTION / "best_val.pt"
    model, best_ndcg, best_epoch = train_model(
        model,
        examples,
        pop_probs,
        log_q,
        epochs=6,
        lr=1.5e-3,
        eval_fn=eval_val,
        deadline=val_deadline,
        ckpt_path=ckpt,
    )
    log(f"best val ndcg@10={best_ndcg:.6f} epoch={best_epoch}")

    model.eval()
    with torch.no_grad():
        item_vecs = model.all_item_vectors().cpu().numpy()
    user_vecs = encode_users(model, recs_train, target_uids, TRAIN_CUTOFF, val_pref)
    val_ranking = rank_users(user_vecs, item_vecs, val_elig_aligned, item_ids, target_uids)
    write_ranking(val_ranking, SOLUTION / "validation.parquet")
    (SOLUTION / "val_score.json").write_text(
        json.dumps({"quality_ndcg@10": best_ndcg, "best_epoch": best_epoch}, indent=2)
    )

    recs_full = build_user_recs(["train", "val"], lookup, n_items)
    val_end = 25_393_400
    test_cutoff = val_end + 1800
    test_pref = preference_overrides(
        DATA / "test_preference_state.parquet", lookup, n_items, recs_full, test_cutoff
    )

    # Fallback submission from the val-selected model with test-time histories.
    user_vecs_test = encode_users(model, recs_full, target_uids, test_cutoff, test_pref)
    with torch.no_grad():
        item_vecs = model.all_item_vectors().cpu().numpy()
    fallback = rank_users(user_vecs_test, item_vecs, test_elig_aligned, item_ids, target_uids)
    write_ranking(fallback, SOLUTION / "submission.parquet")

    remaining = deadline - time.time()
    log(f"remaining seconds before refit {remaining:.0f}")
    if remaining < 10 * 60:
        log("skipping full refit; keeping fallback submission")
    else:
        snaps_full = snapshot_list(val_end, n_windows=9)
        log(f"refit snapshots {snaps_full}")
        examples_full = collect_examples(recs_full, snaps_full, MAX_EXAMPLES_REFIT, rng)
        target_hist = np.bincount(examples_full["target"], minlength=n_items).astype(np.float64)
        pop = np.power(np.maximum(target_hist, 0.5), 0.75)
        pop_probs = pop / pop.sum()
        q = np.maximum(target_hist, 0.5)
        q = q / q.sum()
        log_q = torch.from_numpy(np.log(q).astype(np.float32))
        refit = TwoTower(
            vocab_size=vocab,
            n_items=n_items,
            n_feat=item_feat.shape[1],
            d_model=D_MODEL,
            n_stats=N_STATS,
            dropout=0.1,
        )
        refit.set_item_features(torch.from_numpy(item_feat))
        if ckpt.exists():
            refit.load_state_dict(torch.load(ckpt, map_location="cpu"))
            log("warm-started refit from best val checkpoint")
        refit_epochs = max(2, min(4, best_epoch if best_epoch else 3))
        refit, _, _ = train_model(
            refit,
            examples_full,
            pop_probs,
            log_q,
            epochs=refit_epochs,
            lr=8e-4,
            eval_fn=None,
            deadline=deadline,
            ckpt_path=SOLUTION / "best_refit.pt",
        )
        refit.eval()
        with torch.no_grad():
            item_vecs = refit.all_item_vectors().cpu().numpy()
        user_vecs_test = encode_users(refit, recs_full, target_uids, test_cutoff, test_pref)
        submission = rank_users(user_vecs_test, item_vecs, test_elig_aligned, item_ids, target_uids)
        write_ranking(submission, SOLUTION / "submission.parquet")
        model = refit

    log(f"done in {time.time() - t_start:.1f}s val_ndcg={best_ndcg:.6f}")
    (SOLUTION / "run_meta.json").write_text(
        json.dumps(
            {
                "quality_ndcg@10": best_ndcg,
                "best_epoch": best_epoch,
                "seconds": time.time() - t_start,
                "d_model": D_MODEL,
            },
            indent=2,
        )
    )


if __name__ == "__main__":
    main()

/app/solution/train_two_tower.py

Before

from model import TwoTower

sys.path.insert(0, "/task")
from ranking_evaluation import CandidateSet, score_rankings

After

SOLUTION = Path("/app/solution")
sys.path.insert(0, str(SOLUTION))
from model import TwoTower

sys.path.insert(0, "/task")
from ranking_evaluation import CandidateSet, score_rankings

/app/solution/train_two_tower.py

Before

DATA = Path("/task/data")
SOLUTION = Path("/app/solution")
SOLUTION.mkdir(parents=True, exist_ok=True)

After

DATA = Path("/task/data")
SOLUTION.mkdir(parents=True, exist_ok=True)

/app/solution/train_two_tower.py

Before

            used: set[int] = set()
            for item, ts, is_like in zip(rec.org_item[lo:hi], rec.org_ts[lo:hi], rec.org_like[lo:hi]):
                item = int(item)
                if item in seen_set or item in used:
                    continue
                used.add(item)
                weight = 3.0 if int(is_like) == 1 else 1.0
                row = dict(hist)
                row["target"] = item
                row["weight"] = weight
                rows.append(row)
                seen_set.add(item)

After

            used: set[int] = set()
            n_added = 0
            for item, ts, is_like in zip(rec.org_item[lo:hi], rec.org_ts[lo:hi], rec.org_like[lo:hi]):
                item = int(item)
                if item in seen_set or item in used:
                    continue
                used.add(item)
                weight = 3.0 if int(is_like) == 1 else 1.0
                row = dict(hist)
                row["target"] = item
                row["weight"] = weight
                rows.append(row)
                seen_set.add(item)
                n_added += 1
                if n_added >= 10:
                    break

{
  "target_file": "/app/solution/train_two_tower.py",
  "offset": 1,
  "limit": 50
}

/app/solution/train_two_tower.py

Before

        start = max(MIN_HIST, int(0.55 * len(discoveries)))
        for i in range(start, len(discoveries)):
            if (i < len(discoveries) - 6) and (i % 4 != 0):
                continue
            snap, item, is_like = discoveries[i]
            n_hist = int(np.searchsorted(rec.pos_ts, snap))
            if n_hist < MIN_HIST:
                continue
            hist = history_channels(rec, snap)
            row = dict(hist)
            row["target"] = int(item)
            row["weight"] = 3.0 if is_like else 1.0
            rows.append(row)

After

        start = max(MIN_HIST, len(discoveries) - 16)
        for i in range(start, len(discoveries)):
            snap, item, is_like = discoveries[i]
            n_hist = int(np.searchsorted(rec.pos_ts, snap))
            if n_hist < MIN_HIST:
                continue
            hist = history_channels(rec, snap)
            row = dict(hist)
            row["target"] = int(item)
            row["weight"] = 3.0 if is_like else 1.0
            rows.append(row)

/opt/task-venv/bin/python - <<'PY'
import numpy as np
from pathlib import Path
p=Path('/task/data/validation_eligible_candidates.npz')
z=np.load(p)
packed=z['packed_eligible']
n_items=len(z['item_ids'])
a=np.unpackbits(packed, axis=1, bitorder='little')[:, :n_items]
r=3
b=np.unpackbits(packed[r], count=n_items, bitorder='little')
print('match row', r, np.array_equal(a[r], b), a.shape, a[r].sum(), b.sum())
# also test a few rows
ok=True
for r in [0,1,10,100,8455]:
    b=np.unpackbits(packed[r], count=n_items, bitorder='little')
    if not np.array_equal(a[r], b):
        ok=False
        print('mismatch', r)
print('all ok', ok)
PY

{
  "todos": [
    {
      "id": "design",
      "status": "completed"
    },
    {
      "id": "train_val",
      "status": "in_progress"
    }
  ],
  "merge": true
}

cd /app/solution && /opt/task-venv/bin/python -u /app/solution/train_two_tower.py 2>&1 | tee /app/solution/train.log

Options

{
  "block_until_ms": 0
}

{
  "task_ids": [
    "01a0edf9-0322-7160-8be7-7f0e0500b827"
  ],
  "timeout_ms": 60000
}

{
  "task_ids": [
    "01a0edf9-0322-7160-8be7-7f0e0500b827"
  ],
  "timeout_ms": 180000
}

{
  "target_file": "/app/solution/train.log"
}

{
  "target_directory": "/app/solution"
}

{
  "target_file": "/home/agent/.grok/sessions/%2Fapp/181a6127-40b9-40c0-a171-f664123ad260/terminal/call-3fb031cc-ccf1-4121-9d5b-dcb365184449-29.log",
  "offset": 1,
  "limit": 50
}

ls -la /app/solution/; echo '---'; tail -c 2000 /home/agent/.grok/sessions/%2Fapp/181a6127-40b9-40c0-a171-f664123ad260/terminal/call-3fb031cc-ccf1-4121-9d5b-dcb365184449-29.log; echo; echo EXIT; dmesg 2>/dev/null | tail -20 || true; echo MEM; free -h

cd /app/solution && /opt/task-venv/bin/python -u - <<'PY'
import traceback, sys, time
print('start debug', flush=True)
try:
    import train_two_tower as t
    print('imported', flush=True)
    import numpy as np, torch
    from pathlib import Path
    item_ids, item_feat, lookup = t.load_catalog()
    n_items = len(item_ids)
    print('catalog', n_items, flush=True)
    recs = t.build_user_recs(['train'], lookup, n_items)
    rng = np.random.default_rng(0)
    snaps = t.snapshot_list(t.TRAIN_CUTOFF, 2)
    print('snaps', snaps, flush=True)
    examples = t.collect_examples(recs, snaps, 20000, rng)
    print('examples', examples['target'].shape, 'max target', examples['target'].max(), 'min', examples['target'].min(), flush=True)
    print('id maxes', {k: int(examples[k].max()) for k in ['recent_ids','long_ids','like_ids','dislike_ids','oov_ids']}, flush=True)
    vocab = 1 + n_items + t.N_HASH
    model = t.TwoTower(vocab, n_items, item_feat.shape[1], t.D_MODEL, t.N_STATS, 0.1)
    model.set_item_features(torch.from_numpy(item_feat))
    tensors = t.tensors_from_examples(examples)
    batch = {k: v[:32] for k,v in tensors.items()}
    print('batch target', batch['target'][:10], flush=True)
    model.train()
    user_vec = model.user_vectors(batch)
    print('user', user_vec.shape, user_vec.norm(dim=1)[:3], flush=True)
    pop = torch.ones(n_items)/n_items
    extra = torch.multinomial(pop, 64, replacement=True)
    logits = model.sampled_logits(user_vec, batch['target'], extra, log_q=None)
    print('logits', logits.shape, logits[:,0].detach()[:5], flush=True)
    import torch.nn.functional as F
    labels = torch.zeros(logits.size(0), dtype=torch.long)
    loss = F.cross_entropy(logits, labels)
    print('loss', float(loss), flush=True)
    loss.backward()
    print('backward ok', flush=True)
except Exception:
    traceback.print_exc()
    sys.exit(1)
PY

{
  "task_ids": [
    "call-701bd4c6-89ac-45e0-86a0-44db30fc3bae-36"
  ],
  "timeout_ms": 120000
}

/app/solution/train_two_tower.py

Before

def rank_users(
    user_vecs: np.ndarray,
    item_vecs: np.ndarray,
    elig: np.ndarray,
    item_ids: np.ndarray,
    uids: np.ndarray,
    top_n: int = 100,
) -> pd.DataFrame:
    scores = user_vecs @ item_vecs.T
    scores = np.where(elig, scores, -1e9).astype(np.float32, copy=False)
    idx = np.argpartition(-scores, kth=top_n - 1, axis=1)[:, :top_n]
    part = np.take_along_axis(scores, idx, axis=1)
    order = np.argsort(-part, axis=1)
    top = np.take_along_axis(idx, order, axis=1)
    n_users = len(uids)
    out_uid = np.repeat(uids.astype(np.int64), top_n)
    out_item = item_ids[top.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})

After

def rank_users(
    user_vecs: np.ndarray,
    item_vecs: np.ndarray,
    elig: np.ndarray,
    item_ids: np.ndarray,
    uids: np.ndarray,
    top_n: int = 100,
) -> tuple[pd.DataFrame, dict[int, list[int]]]:
    scores = user_vecs @ item_vecs.T
    scores = np.where(elig, scores, np.float32(-1e9))
    idx = np.argpartition(-scores, kth=top_n - 1, axis=1)[:, :top_n]
    part = np.take_along_axis(scores, idx, axis=1)
    order = np.argsort(-part, axis=1)
    top = np.take_along_axis(idx, order, axis=1)
    ranked_ids = item_ids[top]
    n_users = len(uids)
    out_uid = np.repeat(uids.astype(np.int64), top_n)
    out_item = ranked_ids.reshape(-1).astype(np.int64)
    out_rank = np.tile(np.arange(1, top_n + 1, dtype=np.int64), n_users)
    frame = pd.DataFrame({"uid": out_uid, "item_id": out_item, "rank": out_rank})
    ranked = {int(uid): [int(x) for x in row] for uid, row in zip(uids.tolist(), ranked_ids)}
    return frame, ranked

/app/solution/train_two_tower.py

Before

def tensors_from_examples(examples: dict[str, np.ndarray]) -> dict[str, torch.Tensor]:
    out = {}
    for k, v in examples.items():
        if k in {"recent_ids", "long_ids", "like_ids", "dislike_ids", "oov_ids", "target"}:
            out[k] = torch.from_numpy(np.ascontiguousarray(v, dtype=np.int64))
        else:
            out[k] = torch.from_numpy(np.ascontiguousarray(v, dtype=np.float32))
    return out


def iterate_batches(
    tensors: dict[str, torch.Tensor], batch_size: int, rng: torch.Generator
):
    n = tensors["target"].shape[0]
    perm = torch.randperm(n, generator=rng)
    for start in range(0, n, batch_size):
        idx = perm[start : start + batch_size]
        if idx.numel() < 8:
            continue
        yield {k: v[idx] for k, v in tensors.items()}

After

INT_KEYS = {"recent_ids", "long_ids", "like_ids", "dislike_ids", "oov_ids", "target"}


def iterate_batches(examples: dict[str, np.ndarray], batch_size: int, rng: np.random.Generator):
    n = int(examples["target"].shape[0])
    perm = rng.permutation(n)
    for start in range(0, n, batch_size):
        idx = perm[start : start + batch_size]
        if idx.size < 8:
            continue
        batch = {}
        for k, v in examples.items():
            sl = np.ascontiguousarray(v[idx])
            if k in INT_KEYS:
                batch[k] = torch.from_numpy(sl.astype(np.int64, copy=False))
            else:
                batch[k] = torch.from_numpy(sl.astype(np.float32, copy=False))
        yield batch

/app/solution/train_two_tower.py

Before

    tensors = tensors_from_examples(examples)
    n_items = model.n_items
    pop = torch.from_numpy(pop_probs.astype(np.float32))
    opt = AdamW(model.parameters(), lr=lr, weight_decay=1e-4)
    steps_per = max(1, examples["target"].shape[0] // BATCH_SIZE)
    sched = CosineAnnealingLR(opt, T_max=max(1, epochs * steps_per), eta_min=lr * 0.08)
    gen = torch.Generator()
    gen.manual_seed(SEED)
    best_ndcg = -1.0
    best_epoch = 0
    best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}

    for epoch in range(1, epochs + 1):
        if time.time() > deadline:
            log("hit time deadline during training")
            break
        model.train()
        losses = []
        t0 = time.time()
        n_steps = 0
        for batch in iterate_batches(tensors, BATCH_SIZE, gen):
            if time.time() > deadline:
                break
            user_vec = model.user_vectors(batch)
            extra = torch.multinomial(pop, N_NEG, replacement=True)
            logits = model.sampled_logits(user_vec, batch["target"], extra, log_q=log_q)
            labels = torch.zeros(logits.size(0), dtype=torch.long)
            loss = F.cross_entropy(logits, labels, reduction="none")
            loss = (loss * batch["weight"]).mean()
            opt.zero_grad(set_to_none=True)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
            sched.step()
            losses.append(float(loss.item()))
            n_steps += 1
        mean_loss = float(np.mean(losses)) if losses else float("nan")
        ndcg = eval_fn(model) if eval_fn is not None else 0.0
        log(
            f"epoch {epoch} steps={n_steps} loss={mean_loss:.4f} "
            f"ndcg@10={ndcg:.6f} time={time.time() - t0:.1f}s"
        )

After

    n_items = model.n_items
    pop = torch.from_numpy(pop_probs.astype(np.float32))
    opt = AdamW(model.parameters(), lr=lr, weight_decay=1e-4)
    steps_per = max(1, examples["target"].shape[0] // BATCH_SIZE)
    sched = CosineAnnealingLR(opt, T_max=max(1, epochs * steps_per), eta_min=lr * 0.08)
    np_rng = np.random.default_rng(SEED + 7)
    best_ndcg = -1.0
    best_epoch = 0
    best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
    log(f"training examples={examples['target'].shape[0]} steps/epoch~{steps_per}")

    for epoch in range(1, epochs + 1):
        if time.time() > deadline:
            log("hit time deadline during training")
            break
        model.train()
        losses = []
        t0 = time.time()
        n_steps = 0
        for batch in iterate_batches(examples, BATCH_SIZE, np_rng):
            if time.time() > deadline:
                break
            user_vec = model.user_vectors(batch)
            extra = torch.multinomial(pop, N_NEG, replacement=True)
            logits = model.sampled_logits(user_vec, batch["target"], extra, log_q=log_q)
            labels = torch.zeros(logits.size(0), dtype=torch.long)
            loss = F.cross_entropy(logits, labels, reduction="none")
            loss = (loss * batch["weight"]).mean()
            opt.zero_grad(set_to_none=True)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
            sched.step()
            losses.append(float(loss.detach()))
            n_steps += 1
            if n_steps == 1 or n_steps % 100 == 0:
                log(f"epoch {epoch} step {n_steps} loss={losses[-1]:.4f}")
        mean_loss = float(np.mean(losses)) if losses else float("nan")
        ndcg = eval_fn(model) if eval_fn is not None else 0.0
        log(
            f"epoch {epoch} steps={n_steps} loss={mean_loss:.4f} "
            f"ndcg@10={ndcg:.6f} time={time.time() - t0:.1f}s"
        )

/app/solution/train_two_tower.py

Before

    def eval_val(model: TwoTower) -> float:
        model.eval()
        with torch.no_grad():
            item_vecs = model.all_item_vectors().cpu().numpy()
        user_vecs = encode_users(model, recs_train, target_uids, TRAIN_CUTOFF, val_pref)
        ranking = rank_users(user_vecs, item_vecs, val_elig_aligned, item_ids, target_uids)
        ranked = {
            int(uid): list(map(int, grp.sort_values("rank").item_id.tolist()))
            for uid, grp in ranking.groupby("uid")
        }
        scores = score_rankings(ranked, graded)
        return float(scores["quality_ndcg@10"])

After

    def eval_val(model: TwoTower) -> float:
        model.eval()
        with torch.no_grad():
            item_vecs = model.all_item_vectors().cpu().numpy()
        user_vecs = encode_users(model, recs_train, target_uids, TRAIN_CUTOFF, val_pref)
        _, ranked = rank_users(user_vecs, item_vecs, val_elig_aligned, item_ids, target_uids)
        scores = score_rankings(ranked, graded)
        log(f"eval ndcg@10={scores['quality_ndcg@10']:.6f} users={scores['quality_users']}")
        return float(scores["quality_ndcg@10"])

/app/solution/train_two_tower.py

Before

    user_vecs = encode_users(model, recs_train, target_uids, TRAIN_CUTOFF, val_pref)
    val_ranking = rank_users(user_vecs, item_vecs, val_elig_aligned, item_ids, target_uids)
    write_ranking(val_ranking, SOLUTION / "validation.parquet")

After

    user_vecs = encode_users(model, recs_train, target_uids, TRAIN_CUTOFF, val_pref)
    val_ranking, _ = rank_users(user_vecs, item_vecs, val_elig_aligned, item_ids, target_uids)
    write_ranking(val_ranking, SOLUTION / "validation.parquet")

/app/solution/train_two_tower.py

Before

    user_vecs_test = encode_users(model, recs_full, target_uids, test_cutoff, test_pref)
    with torch.no_grad():
        item_vecs = model.all_item_vectors().cpu().numpy()
    fallback = rank_users(user_vecs_test, item_vecs, test_elig_aligned, item_ids, target_uids)
    write_ranking(fallback, SOLUTION / "submission.parquet")

After

    user_vecs_test = encode_users(model, recs_full, target_uids, test_cutoff, test_pref)
    with torch.no_grad():
        item_vecs = model.all_item_vectors().cpu().numpy()
    fallback, _ = rank_users(user_vecs_test, item_vecs, test_elig_aligned, item_ids, target_uids)
    write_ranking(fallback, SOLUTION / "submission.parquet")

/app/solution/train_two_tower.py

Before

        user_vecs_test = encode_users(refit, recs_full, target_uids, test_cutoff, test_pref)
        submission = rank_users(user_vecs_test, item_vecs, test_elig_aligned, item_ids, target_uids)
        write_ranking(submission, SOLUTION / "submission.parquet")

After

        user_vecs_test = encode_users(refit, recs_full, target_uids, test_cutoff, test_pref)
        submission, _ = rank_users(user_vecs_test, item_vecs, test_elig_aligned, item_ids, target_uids)
        write_ranking(submission, SOLUTION / "submission.parquet")

/app/solution/train_two_tower.py

Before

def main() -> None:
    t_start = time.time()
    deadline = t_start + 54 * 60

After

def main() -> None:
    t_start = time.time()
    deadline = t_start + 52 * 60

/app/solution/train_two_tower.py

Before

if __name__ == "__main__":
    main()

After

if __name__ == "__main__":
    try:
        main()
    except Exception:
        import traceback

        tb = traceback.format_exc()
        Path("/app/solution/error.log").write_text(tb)
        log(tb)
        raise

set -o pipefail
cd /app/solution && /opt/task-venv/bin/python -u /app/solution/train_two_tower.py 2>&1 | tee /app/solution/train.log
echo PIPE_EXIT:$?

Options

{
  "block_until_ms": 0
}

{
  "task_ids": [
    "01a0edfd-94fa-7700-978c-3231cc7e1eab"
  ],
  "timeout_ms": 120000
}

/app/solution/train_two_tower.py

Before

import json
import math
import os
import sys
import time
from dataclasses import dataclass
from pathlib import Path

After

import gc
import json
import math
import os
import sys
import time
from dataclasses import dataclass
from pathlib import Path

/app/solution/train_two_tower.py

Before

def log(msg: str) -> None:
    print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True)

After

def rss_mb() -> float:
    try:
        with open("/proc/self/status") as f:
            for line in f:
                if line.startswith("VmRSS:"):
                    return int(line.split()[1]) / 1024.0
    except OSError:
        return float("nan")
    return float("nan")


def log(msg: str) -> None:
    print(f"[{time.strftime('%H:%M:%S')}] {msg} rss={rss_mb():.0f}MB", flush=True)

/app/solution/train_two_tower.py

Before

            out[key] = np.stack([r[key] for r in rows], axis=0)
    log(
        f"example tensor shapes: "
        + ", ".join(f"{k}{tuple(v.shape)}" for k, v in out.items())
    )
    return out

After

            stacked = np.stack([r[key] for r in rows], axis=0)
            if key.endswith("_ids"):
                stacked = stacked.astype(np.int32, copy=False)
            out[key] = stacked
    del rows
    gc.collect()
    log(
        "example tensor shapes: "
        + ", ".join(f"{k}{tuple(v.shape)}" for k, v in out.items())
    )
    return out

/app/solution/train_two_tower.py

Before

def rank_users(
    user_vecs: np.ndarray,
    item_vecs: np.ndarray,
    elig: np.ndarray,
    item_ids: np.ndarray,
    uids: np.ndarray,
    top_n: int = 100,
) -> tuple[pd.DataFrame, dict[int, list[int]]]:
    scores = user_vecs @ item_vecs.T
    scores = np.where(elig, scores, np.float32(-1e9))
    idx = np.argpartition(-scores, kth=top_n - 1, axis=1)[:, :top_n]
    part = np.take_along_axis(scores, idx, axis=1)
    order = np.argsort(-part, axis=1)
    top = np.take_along_axis(idx, order, axis=1)
    ranked_ids = item_ids[top]
    n_users = len(uids)
    out_uid = np.repeat(uids.astype(np.int64), top_n)
    out_item = ranked_ids.reshape(-1).astype(np.int64)
    out_rank = np.tile(np.arange(1, top_n + 1, dtype=np.int64), n_users)
    frame = pd.DataFrame({"uid": out_uid, "item_id": out_item, "rank": out_rank})
    ranked = {int(uid): [int(x) for x in row] for uid, row in zip(uids.tolist(), ranked_ids)}
    return frame, ranked

After

def rank_users(
    user_vecs: np.ndarray,
    item_vecs: np.ndarray,
    elig: np.ndarray,
    item_ids: np.ndarray,
    uids: np.ndarray,
    top_n: int = 100,
    chunk: int = 256,
) -> tuple[pd.DataFrame, dict[int, list[int]]]:
    n_users = len(uids)
    top = np.empty((n_users, top_n), dtype=np.int32)
    for start in range(0, n_users, chunk):
        end = min(start + chunk, n_users)
        scores = user_vecs[start:end] @ item_vecs.T
        scores[~elig[start:end]] = np.float32(-1e9)
        idx = np.argpartition(-scores, kth=top_n - 1, axis=1)[:, :top_n]
        part = np.take_along_axis(scores, idx, axis=1)
        order = np.argsort(-part, axis=1)
        top[start:end] = np.take_along_axis(idx, order, axis=1)
        del scores, idx, part, order
    ranked_ids = item_ids[top]
    out_uid = np.repeat(uids.astype(np.int64), top_n)
    out_item = ranked_ids.reshape(-1).astype(np.int64)
    out_rank = np.tile(np.arange(1, top_n + 1, dtype=np.int64), n_users)
    frame = pd.DataFrame({"uid": out_uid, "item_id": out_item, "rank": out_rank})
    ranked = {int(uid): [int(x) for x in row] for uid, row in zip(uids.tolist(), ranked_ids)}
    return frame, ranked

/app/solution/train_two_tower.py

Before

            extra = torch.multinomial(pop, N_NEG, replacement=True)
            logits = model.sampled_logits(user_vec, batch["target"], extra, log_q=log_q)
            labels = torch.zeros(logits.size(0), dtype=torch.long)
            loss = F.cross_entropy(logits, labels, reduction="none")
            loss = (loss * batch["weight"]).mean()
            opt.zero_grad(set_to_none=True)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
            sched.step()
            losses.append(float(loss.detach()))
            n_steps += 1
            if n_steps == 1 or n_steps % 100 == 0:
                log(f"epoch {epoch} step {n_steps} loss={losses[-1]:.4f}")
        mean_loss = float(np.mean(losses)) if losses else float("nan")
        ndcg = eval_fn(model) if eval_fn is not None else 0.0

After

            extra = torch.multinomial(pop, N_NEG, replacement=True)
            logits = model.sampled_logits(user_vec, batch["target"], extra, log_q=log_q)
            labels = torch.zeros(logits.size(0), dtype=torch.long)
            per = F.cross_entropy(logits, labels, reduction="none")
            loss = (per * batch["weight"]).mean()
            opt.zero_grad(set_to_none=True)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
            sched.step()
            losses.append(float(loss.detach()))
            n_steps += 1
            if n_steps == 1 or n_steps % 100 == 0:
                log(f"epoch {epoch} step {n_steps} loss={losses[-1]:.4f}")
            del batch, user_vec, extra, logits, labels, per, loss
        mean_loss = float(np.mean(losses)) if losses else float("nan")
        gc.collect()
        log(f"epoch {epoch} train done, evaluating")
        ndcg = eval_fn(model) if eval_fn is not None else 0.0

/app/solution/train_two_tower.py

Before

N_NEG = 512
BATCH_SIZE = 256

After

N_NEG = 256
BATCH_SIZE = 192

/app/solution/train_two_tower.py

Before

def preference_overrides(

After

def strip_example_only_fields(recs: dict[int, UserRec]) -> None:
    z64 = _empty_i64()
    z8 = _empty_i8()
    for rec in recs.values():
        rec.all_item = z64
        rec.all_ts = z64
        rec.org_item = z64
        rec.org_ts = z64
        rec.org_like = z8


def preference_overrides(

/app/solution/train_two_tower.py

Before

    examples = collect_examples(recs_train, snaps, MAX_EXAMPLES_VAL, rng)

    target_hist = np.bincount(examples["target"], minlength=n_items).astype(np.float64)

After

    examples = collect_examples(recs_train, snaps, MAX_EXAMPLES_VAL, rng)
    strip_example_only_fields(recs_train)
    gc.collect()

    target_hist = np.bincount(examples["target"], minlength=n_items).astype(np.float64)

/app/solution/train_two_tower.py

Before

        examples_full = collect_examples(recs_full, snaps_full, MAX_EXAMPLES_REFIT, rng)
        target_hist = np.bincount(examples_full["target"], minlength=n_items).astype(np.float64)

After

        examples_full = collect_examples(recs_full, snaps_full, MAX_EXAMPLES_REFIT, rng)
        strip_example_only_fields(recs_full)
        gc.collect()
        target_hist = np.bincount(examples_full["target"], minlength=n_items).astype(np.float64)

/app/solution/train_two_tower.py

Before

            del batch, user_vec, extra, logits, labels, per, loss
        mean_loss = float(np.mean(losses)) if losses else float("nan")
        gc.collect()
        log(f"epoch {epoch} train done, evaluating")
        ndcg = eval_fn(model) if eval_fn is not None else 0.0

After

            del batch, user_vec, extra, logits, labels, per, loss
        mean_loss = float(np.mean(losses)) if losses else float("nan")
        epoch_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
        torch.save(epoch_state, ckpt_path.with_name(ckpt_path.stem + f"_epoch{epoch}.pt"))
        gc.collect()
        log(f"epoch {epoch} train done, evaluating")
        ndcg = eval_fn(model) if eval_fn is not None else 0.0

/app/solution/train_two_tower.py

Before

        user_vecs = encode_users(model, recs_train, target_uids, TRAIN_CUTOFF, val_pref)
        _, ranked = rank_users(user_vecs, item_vecs, val_elig_aligned, item_ids, target_uids)
        scores = score_rankings(ranked, graded)
        log(f"eval ndcg@10={scores['quality_ndcg@10']:.6f} users={scores['quality_users']}")
        return float(scores["quality_ndcg@10"])

After

        user_vecs = encode_users(model, recs_train, target_uids, TRAIN_CUTOFF, val_pref)
        frame, ranked = rank_users(user_vecs, item_vecs, val_elig_aligned, item_ids, target_uids)
        scores = score_rankings(ranked, graded)
        log(f"eval ndcg@10={scores['quality_ndcg@10']:.6f} users={scores['quality_users']}")
        write_ranking(frame, SOLUTION / "validation.parquet")
        return float(scores["quality_ndcg@10"])

set -o pipefail
cd /app/solution && /opt/task-venv/bin/python -u /app/solution/train_two_tower.py 2>&1 | tee /app/solution/train.log
echo PIPE_EXIT:$?

Options

{
  "block_until_ms": 0
}

{
  "task_ids": [
    "01a0ee00-92fa-7812-9f5a-54627875f222"
  ],
  "timeout_ms": 180000
}

{
  "task_ids": [
    "01a0ee00-92fa-7812-9f5a-54627875f222"
  ],
  "timeout_ms": 300000
}

{
  "task_ids": [
    "01a0ee00-92fa-7812-9f5a-54627875f222"
  ],
  "timeout_ms": 180000
}

cp -f /app/solution/submission.parquet /app/solution/submission_fallback.parquet
cp -f /app/solution/validation.parquet /app/solution/validation_v1.parquet
cp -f /app/solution/best_val.pt /app/solution/best_val_v1.pt
/opt/task-venv/bin/python /task/evaluate_public.py /app/solution/validation.parquet

/app/solution/model.py

"""Two-tower PyTorch ranker: user history encoder and independent item tower."""

from __future__ import annotations

import math

import torch
import torch.nn as nn
import torch.nn.functional as F


class RecencyPool(nn.Module):
    """Softmax attention with a learned recency bias over padded token vectors."""

    def __init__(self, d_model: int):
        super().__init__()
        self.query = nn.Parameter(torch.zeros(d_model))
        nn.init.normal_(self.query, std=0.02)
        self.recency_scale = nn.Parameter(torch.tensor(0.6))
        self.scale = d_model**-0.5

    def forward(self, x: torch.Tensor, age_days: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        scores = (x * self.query).sum(-1) * self.scale
        recency = -F.softplus(self.recency_scale) * torch.log1p(age_days.clamp(min=0.0))
        scores = scores + recency
        scores = scores.masked_fill(~mask, torch.finfo(scores.dtype).min)
        empty = ~mask.any(dim=1)
        scores = scores.masked_fill(empty.unsqueeze(1), 0.0)
        weights = torch.softmax(scores, dim=1) * mask.float()
        denom = weights.sum(dim=1, keepdim=True).clamp(min=1e-6)
        weights = weights / denom
        return (weights.unsqueeze(-1) * x).sum(dim=1)


class ItemTower(nn.Module):
    """Maps a catalog item to an L2-normalized vector from ID + static features."""

    def __init__(self, embeddings: nn.Embedding, n_items: int, n_feat: int, d_model: int, dropout: float):
        super().__init__()
        self.embeddings = embeddings
        self.n_items = n_items
        self.feat_proj = nn.Linear(n_feat, d_model)
        self.net = nn.Sequential(
            nn.Linear(d_model * 2, d_model),
            nn.ReLU(inplace=True),
            nn.Dropout(dropout),
            nn.Linear(d_model, d_model),
        )

    def encode(self, catalog_idx: torch.Tensor, item_features: torch.Tensor) -> torch.Tensor:
        emb = self.embeddings(catalog_idx + 1)
        feat = self.feat_proj(item_features[catalog_idx])
        hidden = self.net(torch.cat([emb, feat], dim=-1))
        return F.normalize(hidden + emb, dim=-1)

    def all_vectors(self, item_features: torch.Tensor) -> torch.Tensor:
        idx = torch.arange(self.n_items, device=item_features.device)
        return self.encode(idx, item_features)


class UserTower(nn.Module):
    """Builds a user vector from history item vectors (no user-ID lookup)."""

    def __init__(self, embeddings: nn.Embedding, n_items: int, d_model: int, n_stats: int, dropout: float):
        super().__init__()
        self.embeddings = embeddings
        self.n_items = n_items
        self.recent_pool = RecencyPool(d_model)
        self.long_pool = RecencyPool(d_model)
        self.like_pool = RecencyPool(d_model)
        self.dislike_pool = RecencyPool(d_model)
        self.oov_pool = RecencyPool(d_model)
        self.half_life = nn.Parameter(torch.tensor(24.0))
        in_dim = d_model * 5 + n_stats
        self.mlp = nn.Sequential(
            nn.Linear(in_dim, d_model * 2),
            nn.LayerNorm(d_model * 2),
            nn.ReLU(inplace=True),
            nn.Dropout(dropout),
            nn.Linear(d_model * 2, d_model),
        )
        self.channel_w = nn.Parameter(torch.tensor([1.0, 0.55, 0.40, -0.25, 0.12]))

    def _gather(self, ids: torch.Tensor, item_mat: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        mask = ids > 0
        cat_mask = (ids >= 1) & (ids <= self.n_items)
        safe = ids.clamp(min=1, max=self.n_items) - 1
        catalog_vecs = item_mat[safe]
        oov_vecs = F.normalize(self.embeddings(ids.clamp(min=0)), dim=-1)
        vecs = torch.where(cat_mask.unsqueeze(-1), catalog_vecs, oov_vecs)
        vecs = vecs * mask.unsqueeze(-1).float()
        return vecs, mask

    def _mean(self, x: torch.Tensor, age: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        hl = F.softplus(self.half_life) + 1.0
        weights = torch.exp(-math.log(2.0) * age.clamp(min=0.0) / hl) * mask.float()
        denom = weights.sum(dim=1, keepdim=True).clamp(min=1e-6)
        return (x * (weights / denom).unsqueeze(-1)).sum(dim=1)

    def forward(self, batch: dict[str, torch.Tensor], item_mat: torch.Tensor) -> torch.Tensor:
        channels = []
        means = []
        specs = [
            ("recent_ids", "recent_age", self.recent_pool),
            ("long_ids", "long_age", self.long_pool),
            ("like_ids", "like_age", self.like_pool),
            ("dislike_ids", "dislike_age", self.dislike_pool),
            ("oov_ids", "oov_age", self.oov_pool),
        ]
        for id_key, age_key, pool in specs:
            vecs, mask = self._gather(batch[id_key], item_mat)
            channels.append(pool(vecs, batch[age_key], mask))
            means.append(self._mean(vecs, batch[age_key], mask))
        stacked_mean = torch.stack(means, dim=1)
        base = (stacked_mean * self.channel_w.view(1, -1, 1)).sum(dim=1)
        mix = self.mlp(torch.cat(channels + [batch["stats"]], dim=-1))
        return F.normalize(base + mix, dim=-1)


class TwoTower(nn.Module):
    """Dot-product two-tower model with a shared item-ID embedding table."""

    def __init__(
        self,
        vocab_size: int,
        n_items: int,
        n_feat: int,
        d_model: int = 64,
        n_stats: int = 8,
        dropout: float = 0.1,
    ):
        super().__init__()
        self.n_items = n_items
        self.d_model = d_model
        self.item_emb = nn.Embedding(vocab_size, d_model, padding_idx=0)
        nn.init.normal_(self.item_emb.weight, std=0.02)
        self.item_emb.weight.data[0].zero_()
        self.item_tower = ItemTower(self.item_emb, n_items, n_feat, d_model, dropout)
        self.user_tower = UserTower(self.item_emb, n_items, d_model, n_stats, dropout)
        self.logit_scale = nn.Parameter(torch.tensor(math.log(1.0 / 0.07)))
        self.register_buffer("item_features", torch.zeros(n_items, n_feat), persistent=True)

    def set_item_features(self, features: torch.Tensor) -> None:
        self.item_features.copy_(features)

    def all_item_vectors(self) -> torch.Tensor:
        return self.item_tower.all_vectors(self.item_features)

    def user_vectors(self, batch: dict[str, torch.Tensor], item_mat: torch.Tensor | None = None) -> torch.Tensor:
        if item_mat is None:
            item_mat = self.all_item_vectors()
        return self.user_tower(batch, item_mat)

    def item_vectors(self, catalog_idx: torch.Tensor) -> torch.Tensor:
        return self.item_tower.encode(catalog_idx, self.item_features)

    def score(self, user_vec: torch.Tensor, item_vec: torch.Tensor) -> torch.Tensor:
        return (user_vec * item_vec).sum(-1)

    def full_softmax_loss(self, batch: dict[str, torch.Tensor]) -> torch.Tensor:
        item_mat = self.all_item_vectors()
        user_vec = self.user_vectors(batch, item_mat)
        scale = self.logit_scale.exp().clamp(max=50.0)
        logits = user_vec @ item_mat.t() * scale
        per = F.cross_entropy(logits, batch["target"], reduction="none", label_smoothing=0.03)
        return (per * batch["weight"]).mean()

/app/solution/train_two_tower.py

Before

def _pad_last(ids: np.ndarray, ages: np.ndarray, length: int) -> tuple[np.ndarray, np.ndarray]:
    out_ids = np.zeros(length, dtype=np.int64)
    out_age = np.zeros(length, dtype=np.float32)
    n = int(ids.shape[0])
    if n <= 0:
        return out_ids, out_age
    sl = min(n, length)
    out_ids[-sl:] = ids[-sl:]
    out_age[-sl:] = ages[-sl:].astype(np.float32, copy=False)
    return out_ids, out_age

After

def _pad_last(ids: np.ndarray, ages: np.ndarray, length: int) -> tuple[np.ndarray, np.ndarray]:
    out_ids = np.zeros(length, dtype=np.int64)
    out_age = np.zeros(length, dtype=np.float32)
    n = int(ids.shape[0])
    if n <= 0:
        return out_ids, out_age
    sl = min(n, length)
    out_ids[-sl:] = ids[-sl:]
    out_age[-sl:] = ages[-sl:].astype(np.float32, copy=False)
    return out_ids, out_age


def _last_unique(ids: np.ndarray, ages: np.ndarray, length: int) -> tuple[np.ndarray, np.ndarray]:
    n = int(ids.shape[0])
    if n <= 0:
        return _pad_last(ids, ages, length)
    rev = ids[::-1]
    _, first = np.unique(rev, return_index=True)
    last_pos = np.sort(n - 1 - first)
    if last_pos.size > length:
        last_pos = last_pos[-length:]
    return _pad_last(ids[last_pos], ages[last_pos], length)

/app/solution/train_two_tower.py

Before

    def take(ids: np.ndarray, ts: np.ndarray, length: int, recent_only: bool = False) -> tuple[np.ndarray, np.ndarray]:
        end = int(np.searchsorted(ts, cutoff))
        ids, ts = ids[:end], ts[:end]
        if ids.size == 0:
            return _pad_last(ids, np.zeros(0, np.float32), length)
        age = _age_days(ts, cutoff)
        if recent_only:
            keep = age <= RECENT_DAYS
            ids, age = ids[keep], age[keep]
        return _pad_last(ids, age, length)

After

    def take(ids: np.ndarray, ts: np.ndarray, length: int, recent_only: bool = False) -> tuple[np.ndarray, np.ndarray]:
        end = int(np.searchsorted(ts, cutoff))
        ids, ts = ids[:end], ts[:end]
        if ids.size == 0:
            return _pad_last(ids, np.zeros(0, np.float32), length)
        age = _age_days(ts, cutoff)
        if recent_only:
            keep = age <= RECENT_DAYS
            ids, age = ids[keep], age[keep]
        return _last_unique(ids, age, length)

/app/solution/train_two_tower.py

Before

            vec = model.user_vectors(batch)
            vectors[start : start + len(chunk)] = vec.cpu().numpy()
    return vectors

After

            vec = model.user_vectors(batch, item_mat)
            vectors[start : start + len(chunk)] = vec.cpu().numpy()
    return vectors

/app/solution/train_two_tower.py

Before

    model.eval()
    vectors = np.zeros((len(uids), model.d_model), dtype=np.float32)
    with torch.no_grad():
        for start in range(0, len(uids), batch_size):

After

    model.eval()
    vectors = np.zeros((len(uids), model.d_model), dtype=np.float32)
    with torch.no_grad():
        item_mat = model.all_item_vectors()
        for start in range(0, len(uids), batch_size):

/app/solution/train_two_tower.py

Before

def train_model(
    model: TwoTower,
    examples: dict[str, np.ndarray],
    pop_probs: np.ndarray,
    log_q: torch.Tensor,
    epochs: int,
    lr: float,
    eval_fn,
    deadline: float,
    ckpt_path: Path,
) -> tuple[TwoTower, float, int]:
    n_items = model.n_items
    pop = torch.from_numpy(pop_probs.astype(np.float32))
    opt = AdamW(model.parameters(), lr=lr, weight_decay=1e-4)
    steps_per = max(1, examples["target"].shape[0] // BATCH_SIZE)
    sched = CosineAnnealingLR(opt, T_max=max(1, epochs * steps_per), eta_min=lr * 0.08)
    np_rng = np.random.default_rng(SEED + 7)
    best_ndcg = -1.0
    best_epoch = 0
    best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
    log(f"training examples={examples['target'].shape[0]} steps/epoch~{steps_per}")

    for epoch in range(1, epochs + 1):
        if time.time() > deadline:
            log("hit time deadline during training")
            break
        model.train()
        losses = []
        t0 = time.time()
        n_steps = 0
        for batch in iterate_batches(examples, BATCH_SIZE, np_rng):
            if time.time() > deadline:
                break
            user_vec = model.user_vectors(batch)
            extra = torch.multinomial(pop, N_NEG, replacement=True)
            logits = model.sampled_logits(user_vec, batch["target"], extra, log_q=log_q)
            labels = torch.zeros(logits.size(0), dtype=torch.long)
            per = F.cross_entropy(logits, labels, reduction="none")
            loss = (per * batch["weight"]).mean()
            opt.zero_grad(set_to_none=True)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
            sched.step()
            losses.append(float(loss.detach()))
            n_steps += 1
            if n_steps == 1 or n_steps % 100 == 0:
                log(f"epoch {epoch} step {n_steps} loss={losses[-1]:.4f}")
            del batch, user_vec, extra, logits, labels, per, loss
        mean_loss = float(np.mean(losses)) if losses else float("nan")
        epoch_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
        torch.save(epoch_state, ckpt_path.with_name(ckpt_path.stem + f"_epoch{epoch}.pt"))
        gc.collect()
        log(f"epoch {epoch} train done, evaluating")
        ndcg = eval_fn(model) if eval_fn is not None else 0.0
        log(
            f"epoch {epoch} steps={n_steps} loss={mean_loss:.4f} "
            f"ndcg@10={ndcg:.6f} time={time.time() - t0:.1f}s"
        )
        if ndcg >= best_ndcg:
            best_ndcg = ndcg
            best_epoch = epoch
            best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
            torch.save(best_state, ckpt_path)
            log(f"saved checkpoint ndcg={best_ndcg:.6f} -> {ckpt_path}")
        elif eval_fn is None:
            best_ndcg = mean_loss
            best_epoch = epoch
            best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
            torch.save(best_state, ckpt_path)
    model.load_state_dict(best_state)
    return model, float(best_ndcg), best_epoch

After

def train_model(
    model: TwoTower,
    examples: dict[str, np.ndarray],
    pop_probs: np.ndarray,
    log_q: torch.Tensor,
    epochs: int,
    lr: float,
    eval_fn,
    deadline: float,
    ckpt_path: Path,
) -> tuple[TwoTower, float, int]:
    opt = AdamW(model.parameters(), lr=lr, weight_decay=1e-4)
    steps_per = max(1, examples["target"].shape[0] // BATCH_SIZE)
    sched = CosineAnnealingLR(opt, T_max=max(1, epochs * steps_per), eta_min=lr * 0.1)
    np_rng = np.random.default_rng(SEED + 7)
    best_ndcg = -1.0
    best_epoch = 0
    patience = 0
    best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
    log(f"training examples={examples['target'].shape[0]} steps/epoch~{steps_per}")

    for epoch in range(1, epochs + 1):
        if time.time() > deadline:
            log("hit time deadline during training")
            break
        model.train()
        losses = []
        t0 = time.time()
        n_steps = 0
        for batch in iterate_batches(examples, BATCH_SIZE, np_rng):
            if time.time() > deadline:
                break
            loss = model.full_softmax_loss(batch)
            opt.zero_grad(set_to_none=True)
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            opt.step()
            sched.step()
            losses.append(float(loss.detach()))
            n_steps += 1
            if n_steps == 1 or n_steps % 50 == 0:
                log(f"epoch {epoch} step {n_steps} loss={losses[-1]:.4f}")
            del batch, loss
        mean_loss = float(np.mean(losses)) if losses else float("nan")
        gc.collect()
        log(f"epoch {epoch} train done, evaluating")
        ndcg = eval_fn(model) if eval_fn is not None else 0.0
        log(
            f"epoch {epoch} steps={n_steps} loss={mean_loss:.4f} "
            f"ndcg@10={ndcg:.6f} time={time.time() - t0:.1f}s"
        )
        if eval_fn is None:
            best_ndcg = mean_loss
            best_epoch = epoch
            best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
            torch.save(best_state, ckpt_path)
        elif ndcg > best_ndcg + 1e-6:
            best_ndcg = ndcg
            best_epoch = epoch
            patience = 0
            best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
            torch.save(best_state, ckpt_path)
            log(f"saved checkpoint ndcg={best_ndcg:.6f} -> {ckpt_path}")
        else:
            patience += 1
            if patience >= 1 and epoch >= 2:
                log("early stopping")
                break
    model.load_state_dict(best_state)
    return model, float(best_ndcg), best_epoch

/app/solution/train_two_tower.py

Before

N_NEG = 256
BATCH_SIZE = 192
MAX_EXAMPLES_VAL = 360_000
MAX_EXAMPLES_REFIT = 480_000

After

N_NEG = 256
BATCH_SIZE = 128
MAX_EXAMPLES_VAL = 360_000
MAX_EXAMPLES_REFIT = 420_000

/app/solution/train_two_tower.py

Before

        frame, ranked = rank_users(user_vecs, item_vecs, val_elig_aligned, item_ids, target_uids)
        scores = score_rankings(ranked, graded)
        log(f"eval ndcg@10={scores['quality_ndcg@10']:.6f} users={scores['quality_users']}")
        write_ranking(frame, SOLUTION / "validation.parquet")
        return float(scores["quality_ndcg@10"])

After

        frame, ranked = rank_users(user_vecs, item_vecs, val_elig_aligned, item_ids, target_uids)
        scores = score_rankings(ranked, graded)
        ndcg = float(scores["quality_ndcg@10"])
        log(f"eval ndcg@10={ndcg:.6f} users={scores['quality_users']}")
        if ndcg >= getattr(eval_val, "best", -1.0):
            eval_val.best = ndcg
            write_ranking(frame, SOLUTION / "validation.parquet")
            write_ranking(frame, SOLUTION / "validation_best.parquet")
        return ndcg

/app/solution/train_two_tower.py

Before

    val_deadline = min(deadline - 16 * 60, t_start + 28 * 60)
    ckpt = SOLUTION / "best_val.pt"
    model, best_ndcg, best_epoch = train_model(
        model,
        examples,
        pop_probs,
        log_q,
        epochs=6,
        lr=1.5e-3,
        eval_fn=eval_val,
        deadline=val_deadline,
        ckpt_path=ckpt,
    )

After

    val_deadline = min(deadline - 14 * 60, t_start + 26 * 60)
    ckpt = SOLUTION / "best_val.pt"
    model, best_ndcg, best_epoch = train_model(
        model,
        examples,
        pop_probs,
        log_q,
        epochs=4,
        lr=1.2e-3,
        eval_fn=eval_val,
        deadline=val_deadline,
        ckpt_path=ckpt,
    )

/app/solution/train_two_tower.py

Before

        refit_epochs = max(2, min(4, best_epoch if best_epoch else 3))
        refit, _, _ = train_model(
            refit,
            examples_full,
            pop_probs,
            log_q,
            epochs=refit_epochs,
            lr=8e-4,
            eval_fn=None,
            deadline=deadline,
            ckpt_path=SOLUTION / "best_refit.pt",
        )

After

        refit_epochs = max(1, min(2, best_epoch if best_epoch else 1))
        refit, _, _ = train_model(
            refit,
            examples_full,
            pop_probs,
            log_q,
            epochs=refit_epochs,
            lr=6e-4,
            eval_fn=None,
            deadline=deadline,
            ckpt_path=SOLUTION / "best_refit.pt",
        )

/app/solution/model_v1.py

"""Checkpoint-compatible v1 two-tower (raw-embedding user tower)."""

from __future__ import annotations

import math

import torch
import torch.nn as nn
import torch.nn.functional as F


class AttnPool(nn.Module):
    def __init__(self, d_model: int):
        super().__init__()
        self.query = nn.Parameter(torch.zeros(d_model))
        nn.init.normal_(self.query, std=0.02)
        self.recency_scale = nn.Parameter(torch.tensor(0.4))
        self.scale = d_model**-0.5

    def forward(self, x: torch.Tensor, age_days: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
        scores = (x * self.query).sum(-1) * self.scale
        recency = -F.softplus(self.recency_scale) * torch.log1p(age_days.clamp(min=0.0))
        scores = scores + recency
        scores = scores.masked_fill(~mask, torch.finfo(scores.dtype).min)
        empty = ~mask.any(dim=1)
        scores = scores.masked_fill(empty.unsqueeze(1), 0.0)
        weights = torch.softmax(scores, dim=1) * mask.float()
        denom = weights.sum(dim=1, keepdim=True).clamp(min=1e-6)
        weights = weights / denom
        return (weights.unsqueeze(-1) * x).sum(dim=1)


class ItemTower(nn.Module):
    def __init__(self, embeddings: nn.Embedding, n_items: int, n_feat: int, d_model: int, dropout: float):
        super().__init__()
        self.embeddings = embeddings
        self.n_items = n_items
        self.feat_proj = nn.Linear(n_feat, d_model)
        self.net = nn.Sequential(
            nn.Linear(d_model * 2, d_model),
            nn.ReLU(inplace=True),
            nn.Dropout(dropout),
            nn.Linear(d_model, d_model),
        )

    def encode(self, catalog_idx: torch.Tensor, item_features: torch.Tensor) -> torch.Tensor:
        emb = self.embeddings(catalog_idx + 1)
        feat = self.feat_proj(item_features[catalog_idx])
        hidden = self.net(torch.cat([emb, feat], dim=-1))
        return F.normalize(hidden + emb, dim=-1)

    def all_vectors(self, item_features: torch.Tensor) -> torch.Tensor:
        idx = torch.arange(self.n_items, device=item_features.device)
        return self.encode(idx, item_features)


class UserTower(nn.Module):
    def __init__(self, embeddings: nn.Embedding, d_model: int, n_stats: int, dropout: float):
        super().__init__()
        self.embeddings = embeddings
        self.recent_pool = AttnPool(d_model)
        self.long_pool = AttnPool(d_model)
        self.like_pool = AttnPool(d_model)
        self.dislike_pool = AttnPool(d_model)
        self.oov_pool = AttnPool(d_model)
        in_dim = d_model * 5 + n_stats
        self.mlp = nn.Sequential(
            nn.Linear(in_dim, d_model * 2),
            nn.LayerNorm(d_model * 2),
            nn.ReLU(inplace=True),
            nn.Dropout(dropout),
            nn.Linear(d_model * 2, d_model),
            nn.LayerNorm(d_model),
            nn.ReLU(inplace=True),
            nn.Linear(d_model, d_model),
        )

    def _pool(self, pool: AttnPool, ids: torch.Tensor, age: torch.Tensor) -> torch.Tensor:
        mask = ids > 0
        x = self.embeddings(ids)
        return pool(x, age, mask)

    def forward(self, batch: dict[str, torch.Tensor]) -> torch.Tensor:
        parts = [
            self._pool(self.recent_pool, batch["recent_ids"], batch["recent_age"]),
            self._pool(self.long_pool, batch["long_ids"], batch["long_age"]),
            self._pool(self.like_pool, batch["like_ids"], batch["like_age"]),
            self._pool(self.dislike_pool, batch["dislike_ids"], batch["dislike_age"]),
            self._pool(self.oov_pool, batch["oov_ids"], batch["oov_age"]),
            batch["stats"],
        ]
        return F.normalize(self.mlp(torch.cat(parts, dim=-1)), dim=-1)


class TwoTowerV1(nn.Module):
    def __init__(self, vocab_size: int, n_items: int, n_feat: int, d_model: int = 64, n_stats: int = 8, dropout: float = 0.1):
        super().__init__()
        self.n_items = n_items
        self.d_model = d_model
        self.item_emb = nn.Embedding(vocab_size, d_model, padding_idx=0)
        nn.init.normal_(self.item_emb.weight, std=0.02)
        self.item_emb.weight.data[0].zero_()
        self.item_tower = ItemTower(self.item_emb, n_items, n_feat, d_model, dropout)
        self.user_tower = UserTower(self.item_emb, d_model, n_stats, dropout)
        self.logit_scale = nn.Parameter(torch.tensor(math.log(1.0 / 0.07)))
        self.register_buffer("item_features", torch.zeros(n_items, n_feat), persistent=True)

    def set_item_features(self, features: torch.Tensor) -> None:
        self.item_features.copy_(features)

    def user_vectors(self, batch: dict[str, torch.Tensor]) -> torch.Tensor:
        return self.user_tower(batch)

    def all_item_vectors(self) -> torch.Tensor:
        return self.item_tower.all_vectors(self.item_features)

/app/solution/train_two_tower.py

Before

    log(f"best val ndcg@10={best_ndcg:.6f} epoch={best_epoch}")

    model.eval()
    with torch.no_grad():
        item_vecs = model.all_item_vectors().cpu().numpy()
    user_vecs = encode_users(model, recs_train, target_uids, TRAIN_CUTOFF, val_pref)
    val_ranking, _ = rank_users(user_vecs, item_vecs, val_elig_aligned, item_ids, target_uids)
    write_ranking(val_ranking, SOLUTION / "validation.parquet")
    (SOLUTION / "val_score.json").write_text(
        json.dumps({"quality_ndcg@10": best_ndcg, "best_epoch": best_epoch}, indent=2)
    )

After

    log(f"best val ndcg@10={best_ndcg:.6f} epoch={best_epoch}")

    model.eval()
    with torch.no_grad():
        item_vecs = model.all_item_vectors().cpu().numpy()
    user_vecs = encode_users(model, recs_train, target_uids, TRAIN_CUTOFF, val_pref)
    val_ranking, _ = rank_users(user_vecs, item_vecs, val_elig_aligned, item_ids, target_uids)
    write_ranking(val_ranking, SOLUTION / "validation.parquet")
    (SOLUTION / "val_score.json").write_text(
        json.dumps({"quality_ndcg@10": best_ndcg, "best_epoch": best_epoch}, indent=2)
    )
    prev_best = 0.026616
    use_v2 = best_ndcg + 1e-6 >= prev_best
    log(f"use_v2={use_v2} new={best_ndcg:.6f} prev={prev_best:.6f}")

/app/solution/train_two_tower.py

Before

    remaining = deadline - time.time()
    log(f"remaining seconds before refit {remaining:.0f}")
    if remaining < 10 * 60:
        log("skipping full refit; keeping fallback submission")
    else:

After

    remaining = deadline - time.time()
    log(f"remaining seconds before refit {remaining:.0f}")
    if not use_v2:
        log("v2 underperformed public v1; restoring v1 validation ranking")
        v1_val = SOLUTION / "validation_v1.parquet"
        if v1_val.exists():
            write_ranking(pd.read_parquet(v1_val), SOLUTION / "validation.parquet")
        log("skipping v2 refit")
    elif remaining < 10 * 60:
        log("skipping full refit; keeping fallback submission")
    else:

/app/solution/train_two_tower.py

Before

    if not use_v2:
        log("v2 underperformed public v1; restoring v1 validation ranking")
        v1_val = SOLUTION / "validation_v1.parquet"
        if v1_val.exists():
            write_ranking(pd.read_parquet(v1_val), SOLUTION / "validation.parquet")
        log("skipping v2 refit")
    elif remaining < 10 * 60:

After

    if not use_v2:
        log("v2 underperformed public v1; restoring v1 validation ranking and v1 test ranker")
        v1_val = SOLUTION / "validation_v1.parquet"
        if v1_val.exists():
            write_ranking(pd.read_parquet(v1_val), SOLUTION / "validation.parquet")
        from model_v1 import TwoTowerV1

        v1 = TwoTowerV1(
            vocab_size=vocab,
            n_items=n_items,
            n_feat=item_feat.shape[1],
            d_model=D_MODEL,
            n_stats=N_STATS,
            dropout=0.1,
        )
        v1.set_item_features(torch.from_numpy(item_feat))
        v1.load_state_dict(torch.load(SOLUTION / "best_val_v1.pt", map_location="cpu"))
        v1.eval()
        with torch.no_grad():
            item_vecs = v1.all_item_vectors().cpu().numpy()
            vectors = np.zeros((len(target_uids), v1.d_model), dtype=np.float32)
            for start in range(0, len(target_uids), 256):
                chunk = target_uids[start : start + 256]
                bag: dict[str, list[np.ndarray]] = {k: [] for k in [
                    "recent_ids", "recent_age", "long_ids", "long_age",
                    "like_ids", "like_age", "dislike_ids", "dislike_age",
                    "oov_ids", "oov_age", "stats",
                ]}
                for uid in chunk:
                    uid = int(uid)
                    rec = recs_full.get(uid)
                    like_ov = dis_ov = None
                    if uid in test_pref:
                        like_ov = test_pref[uid]["like"]
                        dis_ov = test_pref[uid]["dislike"]
                    hist = history_channels(rec, test_cutoff, like_ov, dis_ov)
                    for k, v in hist.items():
                        bag[k].append(v)
                batch = {k: torch.from_numpy(np.stack(v, axis=0)) for k, v in bag.items()}
                vectors[start : start + len(chunk)] = v1.user_vectors(batch).cpu().numpy()
        frame, _ = rank_users(vectors, item_vecs, test_elig_aligned, item_ids, target_uids)
        write_ranking(frame, SOLUTION / "submission.parquet")
        log("skipping v2 refit")
    elif remaining < 10 * 60:

cp -f /app/solution/validation_v1.parquet /app/solution/validation.parquet
set -o pipefail
cd /app/solution && /opt/task-venv/bin/python -u /app/solution/train_two_tower.py 2>&1 | tee /app/solution/train_v2.log
echo PIPE_EXIT:$?

Options

{
  "block_until_ms": 0
}

{
  "task_ids": [
    "01a0ee0c-66b5-7de3-ad8f-ecfab391fe09"
  ],
  "timeout_ms": 180000
}

{
  "task_ids": [
    "01a0ee0c-66b5-7de3-ad8f-ecfab391fe09"
  ],
  "timeout_ms": 240000
}

{
  "task_ids": [
    "01a0ee0c-66b5-7de3-ad8f-ecfab391fe09"
  ],
  "timeout_ms": 300000
}

{
  "task_ids": [
    "01a0ee0c-66b5-7de3-ad8f-ecfab391fe09"
  ],
  "timeout_ms": 300000
}

/app/solution/model.py

Before

        self.mlp = nn.Sequential(
            nn.Linear(in_dim, d_model * 2),
            nn.LayerNorm(d_model * 2),
            nn.ReLU(inplace=True),
            nn.Dropout(dropout),
            nn.Linear(d_model * 2, d_model),
        )
        self.channel_w = nn.Parameter(torch.tensor([1.0, 0.55, 0.40, -0.25, 0.12]))

After

        self.mlp = nn.Sequential(
            nn.Linear(in_dim, d_model * 2),
            nn.LayerNorm(d_model * 2),
            nn.ReLU(inplace=True),
            nn.Dropout(dropout),
            nn.Linear(d_model * 2, d_model),
        )
        nn.init.zeros_(self.mlp[-1].weight)
        nn.init.zeros_(self.mlp[-1].bias)
        self.channel_w = nn.Parameter(torch.tensor([1.0, 0.55, 0.40, -0.25, 0.12]))

/app/solution/apply_and_refit.py

#!/usr/bin/env python3
"""Apply the val-selected checkpoint, then refit one epoch on train+val."""

from __future__ import annotations

import gc
import time

import numpy as np
import pandas as pd
import torch

import train_two_tower as t
from model import TwoTower

SEED = 42


def make_model(n_items: int, n_feat: int, vocab: int) -> TwoTower:
    model = TwoTower(
        vocab_size=vocab,
        n_items=n_items,
        n_feat=n_feat,
        d_model=t.D_MODEL,
        n_stats=t.N_STATS,
        dropout=0.1,
    )
    return model


def apply_model(model, recs, uids, cutoff, pref, elig, item_ids) -> pd.DataFrame:
    model.eval()
    with torch.no_grad():
        item_vecs = model.all_item_vectors().cpu().numpy()
    user_vecs = t.encode_users(model, recs, uids, cutoff, pref)
    frame, _ = t.rank_users(user_vecs, item_vecs, elig, item_ids, uids)
    return frame


def main() -> None:
    rng = np.random.default_rng(SEED)
    torch.manual_seed(SEED)
    item_ids, item_feat, lookup = t.load_catalog()
    n_items = len(item_ids)
    vocab = 1 + n_items + t.N_HASH
    target_uids = pd.read_parquet(t.DATA / "target_users.parquet").uid.to_numpy(np.int64)
    test_uids, test_items, test_elig = t.unpack_eligibility(
        t.DATA / "test_eligible_candidates.npz", n_items
    )
    row = {int(u): i for i, u in enumerate(test_uids.tolist())}
    test_elig_aligned = np.stack([test_elig[row[int(u)]] for u in target_uids], axis=0)
    val_end = 25_393_400
    test_cutoff = val_end + 1800

    t.log("loading train+val history")
    recs_full = t.build_user_recs(["train", "val"], lookup, n_items)
    test_pref = t.preference_overrides(
        t.DATA / "test_preference_state.parquet", lookup, n_items, recs_full, test_cutoff
    )

    ckpt = t.SOLUTION / "best_val.pt"
    model = make_model(n_items, item_feat.shape[1], vocab)
    model.set_item_features(torch.from_numpy(item_feat))
    model.load_state_dict(torch.load(ckpt, map_location="cpu"))
    t.log("applying val-selected checkpoint to test histories")
    frame = apply_model(model, recs_full, target_uids, test_cutoff, test_pref, test_elig_aligned, item_ids)
    t.write_ranking(frame, t.SOLUTION / "submission.parquet")
    t.write_ranking(frame, t.SOLUTION / "submission_valckpt.parquet")

    t.log("refitting 1 epoch from scratch on train+val")
    snaps = t.snapshot_list(val_end, n_windows=9)
    examples = t.collect_examples(recs_full, snaps, t.MAX_EXAMPLES_REFIT, rng)
    t.strip_example_only_fields(recs_full)
    gc.collect()
    target_hist = np.bincount(examples["target"], minlength=n_items).astype(np.float64)
    pop = np.power(np.maximum(target_hist, 0.5), 0.75)
    pop_probs = pop / pop.sum()
    q = np.maximum(target_hist, 0.5)
    q = q / q.sum()
    log_q = torch.from_numpy(np.log(q).astype(np.float32))
    refit = make_model(n_items, item_feat.shape[1], vocab)
    refit.set_item_features(torch.from_numpy(item_feat))
    deadline = time.time() + 14 * 60
    refit, _, _ = t.train_model(
        refit,
        examples,
        pop_probs,
        log_q,
        epochs=1,
        lr=1.2e-3,
        eval_fn=None,
        deadline=deadline,
        ckpt_path=t.SOLUTION / "best_refit.pt",
    )
    frame = apply_model(refit, recs_full, target_uids, test_cutoff, test_pref, test_elig_aligned, item_ids)
    t.write_ranking(frame, t.SOLUTION / "submission.parquet")
    t.log("refit submission written")


if __name__ == "__main__":
    main()

cp -f /app/solution/validation_best.parquet /app/solution/validation.parquet
set -o pipefail
cd /app/solution && /opt/task-venv/bin/python -u /app/solution/apply_and_refit.py 2>&1 | tee /app/solution/apply_refit.log
echo PIPE_EXIT:$?

Options

{
  "block_until_ms": 0
}

{
  "task_ids": [
    "01a0ee1d-4bc3-74b2-ab91-5f5465725bca"
  ],
  "timeout_ms": 120000
}

/opt/task-venv/bin/python /task/evaluate_public.py /app/solution/validation.parquet
/opt/task-venv/bin/python - <<'PY'
import pandas as pd
s=pd.read_parquet('/app/solution/submission.parquet')
print(s.shape, s.columns.tolist(), s.dtypes.to_dict())
print('users', s.uid.nunique(), 'rows/user', s.groupby('uid').size().min(), s.groupby('uid').size().max())
print('rank set ok', set(s.groupby('uid')['rank'].nunique()) )
print(s.head(3))
PY