Agent trace

ls -la /task/ /task/data/ && cat /task/requirements.txt | head -50 && nproc && python -c "import torch; print(torch.__version__)"

sed -n 1,200p /task/ranking_evaluation.py

mkdir -p solution && python -c "
import pandas as pd, numpy as np, pyarrow.parquet as pq
for f in ['candidate_catalog','target_users','validation_preference_state','test_preference_state','val_likes','val_listens']:
    p=pq.ParquetFile('/task/data/%s.parquet'%f)
    print(f, p.metadata.num_rows, p.schema_arrow.names)
p=pq.ParquetFile('/task/data/train_listens.parquet'); print('train_listens', p.metadata.num_rows, p.schema_arrow.names)
p=pq.ParquetFile('/task/data/train_likes.parquet'); print('train_likes', p.metadata.num_rows)
"

sed -n 1,200p /task/evaluate_public.py

cd /app/solution && cat > load.py <<'EOF'
import numpy as np, pandas as pd, pyarrow.parquet as pq, time, os
t0=time.time()
D='/task/data/'
cat=pd.read_parquet(D+'candidate_catalog.parquet')
item_ids=cat.item_id.to_numpy(np.int64)
n_items=len(item_ids)
maxid=item_ids.max()
imap=np.full(maxid+1,-1,np.int32); imap[item_ids]=np.arange(n_items,dtype=np.int32)
np.save('item_ids.npy',item_ids)

def load_listens(path):
    pf=pq.ParquetFile(path)
    outs=[]
    for b in pf.iter_batches(batch_size=2_000_000,columns=['uid','timestamp','item_id','is_organic','played_ratio_pct']):
        d=b.to_pydict()
        it=np.asarray(d['item_id'],np.int64)
        pr=np.asarray(d['played_ratio_pct'],np.float32)
        m=(it<=maxid)&(pr>=50)
        it=it[m]; ci=imap[it]; m2=ci>=0
        idx=np.nonzero(m)[0][m2]
        outs.append((np.asarray(d['uid'],np.int64)[idx].astype(np.int32),
                     np.asarray(d['timestamp'],np.int64)[idx].astype(np.int64),
                     ci[m2],
                     np.asarray(d['is_organic'],np.int8)[idx],
                     (pr[m][m2]>=80).astype(np.int8)))
    return [np.concatenate([o[k] for o in outs]) for k in range(5)]

def load_ev(path):
    d=pd.read_parquet(path,columns=['uid','timestamp','item_id','is_organic'])
    it=d.item_id.to_numpy(np.int64); m=it<=maxid
    d=d[m]; ci=imap[it[m]]; m2=ci>=0; d=d[m2]
    return d.uid.to_numpy(np.int32), d.timestamp.to_numpy(np.int64), ci[m2], d.is_organic.to_numpy(np.int8)

for split in ['train','val']:
    u,t,i,o,s=load_listens(D+f'{split}_listens.parquet')
    np.savez(f'{split}_listens.npz',uid=u,ts=t,item=i,org=o,strong=s)
    print(split,'listens',len(u),time.time()-t0,flush=True)
    for fam in ['likes','dislikes','unlikes','undislikes']:
        u,t,i,o=load_ev(D+f'{split}_{fam}.parquet')
        np.savez(f'{split}_{fam}.npz',uid=u,ts=t,item=i,org=o)
    print(split,'done',time.time()-t0,flush=True)
EOF
python load.py

Options

{
  "timeout": 900000
}

python - <<'EOF'
import numpy as np, pandas as pd, time
t0=time.time()
item_ids=np.load('item_ids.npy'); n_items=len(item_ids)
tu=pd.read_parquet('/task/data/target_users.parquet').uid.to_numpy(np.int64)
z=np.load('/task/data/validation_eligible_candidates.npz')
print(z['uids'].shape, z['item_ids'].shape, z['packed_eligible'].shape)
assert np.array_equal(z['item_ids'], item_ids), 'order differs'
P=z['packed_eligible']
bits=np.unpackbits(P,axis=1,count=n_items,bitorder='little')
print('mean eligible per user', bits.sum(1).mean(), bits.sum(1).min(), bits.sum(1).max())
# repeat analysis
tr=np.load('train_listens.npz'); vl=np.load('val_listens.npz')
tuset=set(tu.tolist())
mask=np.isin(tr['uid'],tu)
print('train rows for targets',mask.sum())
hu=tr['uid'][mask]; hi=tr['item'][mask]
hist=set(zip(hu.tolist(),hi.tolist()))
m2=np.isin(vl['uid'],tu)&(vl['org']==1)&(vl['strong']==1)
vu=vl['uid'][m2]; vi=vl['item'][m2]
pairs=set(zip(vu.tolist(),vi.tolist()))
inh=sum(1 for p in pairs if p in hist)
print('val strong organic pairs',len(pairs),'in train history',inh, inh/len(pairs))
# eligible check
uidrow={int(u):r for r,u in enumerate(z['uids'])}
el=0
for u,i in list(pairs)[:20000]:
    el+= bits[uidrow[u],i]
print('eligible frac of val pos (sample)', el/20000)
print(time.time()-t0)
EOF

Options

{
  "timeout": 600000
}

python - <<'EOF'
import numpy as np, pandas as pd, time
t0=time.time()
item_ids=np.load('item_ids.npy'); n_items=len(item_ids)
tu=pd.read_parquet('/task/data/target_users.parquet').uid.to_numpy(np.int64)
z=np.load('/task/data/validation_eligible_candidates.npz')
bits=np.unpackbits(z['packed_eligible'],axis=1,count=n_items,bitorder='little').astype(bool)
uidrow={int(u):r for r,u in enumerate(z['uids'])}
vl=np.load('val_listens.npz'); vk=np.load('val_likes.npz')
tuset=np.isin(vl['uid'],tu)
m2=tuset&(vl['org']==1)&(vl['strong']==1)
pu,pi=vl['uid'][m2],vl['item'][m2]
rows=np.array([uidrow.get(int(u),-1) for u in pu])
ok=rows>=0
print('strong organic listen pairs',ok.sum(),'eligible',bits[rows[ok],pi[ok]].mean())
mk=np.isin(vk['uid'],tu)&(vk['org']==1)
ku,ki=vk['uid'][mk],vk['item'][mk]
rk=np.array([uidrow.get(int(u),-1) for u in ku]); ok2=rk>=0
print('organic like pairs',ok2.sum(),'eligible',bits[rk[ok2],ki[ok2]].mean())
# num users with >=1 relevant
us=set(pu[ok][bits[rows[ok],pi[ok]]].tolist())|set(ku[ok2][bits[rk[ok2],ki[ok2]]].tolist())
print('users with relevant',len(us),'of',len(tu))
np.save('elig_val.npy',bits)
print(time.time()-t0)
EOF

Options

{
  "timeout": 600000
}

cat > common.py <<'EOF'
import numpy as np, pandas as pd
D='/task/data/'
item_ids=np.load('/app/solution/item_ids.npy'); n_items=len(item_ids)
targets=pd.read_parquet(D+'target_users.parquet').uid.to_numpy(np.int64)

def load_elig(stage):
    z=np.load(D+f'{stage}_eligible_candidates.npz')
    assert np.array_equal(z['item_ids'],item_ids)
    bits=np.unpackbits(z['packed_eligible'],axis=1,count=n_items,bitorder='little').astype(bool)
    return z['uids'].astype(np.int64), bits

def val_truth(uids,bits):
    row={int(u):r for r,u in enumerate(uids)}
    vl=np.load('/app/solution/val_listens.npz'); vk=np.load('/app/solution/val_likes.npz')
    G=np.zeros((len(uids),n_items),np.int8)
    m=(vl['org']==1)&(vl['strong']==1)
    r=np.array([row.get(int(u),-1) for u in vl['uid'][m]]); i=vl['item'][m]; ok=r>=0
    G[r[ok],i[ok]]=1
    m=(vk['org']==1)
    r=np.array([row.get(int(u),-1) for u in vk['uid'][m]]); i=vk['item'][m]; ok=r>=0
    G[r[ok],i[ok]]=3
    G[~bits]=0
    return G

def ndcg10(scores,G,bits):
    """scores: (U,n_items) float32; masked by bits already or here."""
    s=np.where(bits,scores,-1e30)
    top=np.argpartition(-s,10,axis=1)[:,:10]
    ts=np.take_along_axis(s,top,1)
    o=np.argsort(-ts,axis=1)
    top=np.take_along_axis(top,o,1)
    g=np.take_along_axis(G,top,1).astype(np.float64)
    disc=1.0/np.log2(np.arange(2,12))
    dcg=((2.0**g-1)*disc).sum(1)
    idl=[]
    for r in range(G.shape[0]):
        gg=np.sort(G[r][G[r]>0])[::-1][:10].astype(np.float64)
        idl.append(((2.0**gg-1)*disc[:len(gg)]).sum())
    idl=np.array(idl)
    sel=idl>0
    return float((dcg[sel]/idl[sel]).mean()), int(sel.sum())
EOF
python - <<'EOF'
import numpy as np,pandas as pd,time
from common import *
t0=time.time()
uids,bits=load_elig('validation'); G=val_truth(uids,bits)
np.save('G_val.npy',G); np.save('elig_val.npy',bits); np.save('uids_val.npy',uids)
print('rel pairs',(G>0).sum(),'users',( (G>0).any(1)).sum())
# how many relevant-eligible in train history?
tr=np.load('train_listens.npz')
row={int(u):r for r,u in enumerate(uids)}
r=np.array([row.get(int(u),-1) for u in tr['uid']]);ok=r>=0
H=np.zeros((len(uids),n_items),np.float32)
np.add.at(H,(r[ok],tr['item'][ok]),1.0)
print('hist nonzero per user',(H>0).sum(1).mean())
rel=G>0
print('frac relevant-eligible in history',(H[rel]>0).mean())
print('eligible items in history per user',((H>0)&bits).sum(1).mean())
# popularity baseline
cat=pd.read_parquet(D+'candidate_catalog.parquet')
pop=-cat.popularity_rank.to_numpy(np.float32)
sc=np.tile(pop,(len(uids),1))
print('popularity ndcg',ndcg10(sc,G,bits))
print('history-count baseline',ndcg10(H,G,bits))
print(time.time()-t0)
EOF

Options

{
  "timeout": 900000
}

python - <<'EOF'
import numpy as np
tr=np.load('train_listens.npz')
u=tr['uid']
print('rows',len(u),'users',len(np.unique(u)))
print('ts range',tr['ts'].min(),tr['ts'].max())
vl=np.load('val_listens.npz'); print('val ts',vl['ts'].min(),vl['ts'].max(),len(np.unique(vl['uid'])))
EOF

Options

{
  "timeout": 300000
}

cat > features.py <<'EOF'
"""Build history bags (user-tower inputs) and target pairs at arbitrary cutoffs."""
import numpy as np
N_ITEMS = 30000
DAY = 86400

def _agg(uid_idx, item, w, n_users, cap):
    """Aggregate weights per (user,item) and keep top-`cap` items per user."""
    if len(uid_idx) == 0:
        return np.zeros(n_users + 1, np.int64), np.zeros(0, np.int32), np.zeros(0, np.float32)
    key = uid_idx.astype(np.int64) * N_ITEMS + item
    uk, inv = np.unique(key, return_inverse=True)
    tot = np.bincount(inv, weights=w.astype(np.float64)).astype(np.float32)
    uu = (uk // N_ITEMS).astype(np.int32)
    ui = (uk % N_ITEMS).astype(np.int32)
    order = np.lexsort((-tot, uu))
    uu, ui, tot = uu[order], ui[order], tot[order]
    cnt = np.bincount(uu, minlength=n_users)
    start = np.concatenate([[0], np.cumsum(cnt)])
    rank = np.arange(len(uu)) - np.repeat(start[:-1], cnt)
    keep = rank < cap
    uu, ui, tot = uu[keep], ui[keep], tot[keep]
    cnt = np.bincount(uu, minlength=n_users)
    ptr = np.concatenate([[0], np.cumsum(cnt)]).astype(np.int64)
    return ptr, ui.astype(np.int32), tot


def build_bags(ev, cutoff, uid_idx_map, n_users):
    """Return CSR bags (all, recent, liked, disliked) for history strictly before `cutoff`."""
    L = ev['listens']
    m = L['ts'] < cutoff
    u, i, ts = L['uidx'][m], L['item'][m], L['ts'][m]
    w = np.log1p(np.ones(len(u), np.float32))
    bags = {}
    bags['all'] = _agg(u, i, np.ones(len(u), np.float32), n_users, 256)
    r = ts >= cutoff - 30 * DAY
    bags['recent'] = _agg(u[r], i[r], np.ones(int(r.sum()), np.float32), n_users, 128)
    # explicit state: like/dislike net of reversals
    for pos, neg, name, cap in (('likes', 'unlikes', 'liked', 128), ('dislikes', 'undislikes', 'disliked', 64)):
        P, Nn = ev[pos], ev[neg]
        mp = P['ts'] < cutoff
        mn = Nn['ts'] < cutoff
        # a reversal after the last positive cancels it
        key_p = P['uidx'][mp].astype(np.int64) * N_ITEMS + P['item'][mp]
        key_n = Nn['uidx'][mn].astype(np.int64) * N_ITEMS + Nn['item'][mn]
        tp, tn = P['ts'][mp], Nn['ts'][mn]
        if len(key_n):
            o = np.argsort(key_n, kind='stable')
            kn, tnn = key_n[o], tn[o]
            uk, st = np.unique(kn, return_index=True)
            last_n = np.maximum.reduceat(tnn, st) if len(uk) else np.zeros(0, np.int64)
            pos_in = np.searchsorted(uk, key_p)
            pos_in = np.clip(pos_in, 0, max(len(uk) - 1, 0))
            match = (len(uk) > 0) & (uk[pos_in] == key_p)
            cancel = match & (last_n[pos_in] > tp)
        else:
            cancel = np.zeros(len(key_p), bool)
        kk = ~cancel
        uu = P['uidx'][mp][kk]; ii = P['item'][mp][kk]
        bags[name] = _agg(uu, ii, np.ones(len(uu), np.float32), n_users, cap)
    return bags


def build_targets(ev, lo, hi, bag_all_ptr, bag_all_items, hist_seen):
    """Positives in [lo,hi): organic strong listens (grade1) / organic likes (grade3), novel only."""
    L = ev['listens']
    m = (L['ts'] >= lo) & (L['ts'] < hi) & (L['org'] == 1) & (L['strong'] == 1)
    u1, i1 = L['uidx'][m], L['item'][m]
    K = ev['likes']
    m = (K['ts'] >= lo) & (K['ts'] < hi) & (K['org'] == 1)
    u3, i3 = K['uidx'][m], K['item'][m]
    u = np.concatenate([u1, u3]); i = np.concatenate([i1, i3])
    g = np.concatenate([np.ones(len(u1), np.float32), np.full(len(u3), 3.0, np.float32)])
    key = u.astype(np.int64) * N_ITEMS + i
    order = np.lexsort((-g, key))
    key, u, i, g = key[order], u[order], i[order], g[order]
    first = np.concatenate([[True], key[1:] != key[:-1]])
    u, i, g = u[first], i[first], g[first]
    novel = ~np.isin(u.astype(np.int64) * N_ITEMS + i, hist_seen, assume_unique=False)
    return u[novel], i[novel], g[novel]
EOF
echo ok

cat > model.py <<'EOF'
"""PyTorch two-tower ranker: user tower pools interaction history, item tower embeds candidates.

Relevance score is exactly dot(user_vec, item_vec); item vectors are independent of the
user and can be precomputed offline into an MIPS index.
"""
import torch, torch.nn as nn, torch.nn.functional as F

BAGS = ('all', 'recent', 'liked', 'disliked')


class ItemTower(nn.Module):
    def __init__(self, n_items, d, h):
        super().__init__()
        self.emb = nn.Embedding(n_items, d)
        nn.init.normal_(self.emb.weight, std=0.05)
        self.net = nn.Sequential(nn.LayerNorm(d), nn.Linear(d, h), nn.GELU(), nn.Linear(h, d))
        self.skip = nn.Linear(d, d, bias=False)
        nn.init.eye_(self.skip.weight)

    def forward(self, idx):
        e = self.emb(idx)
        return self.skip(e) + self.net(e)


class UserTower(nn.Module):
    """Pools per-bag weighted sums of history-item embeddings; no user-ID parameters."""

    def __init__(self, n_items, d, h):
        super().__init__()
        self.emb = nn.ModuleDict({b: nn.EmbeddingBag(n_items, d, mode='sum') for b in BAGS})
        for b in BAGS:
            nn.init.normal_(self.emb[b].weight, std=0.05)
        self.norm = nn.LayerNorm(len(BAGS) * d)
        self.net = nn.Sequential(nn.Linear(len(BAGS) * d, h), nn.GELU(), nn.Linear(h, d))
        self.skip = nn.Linear(len(BAGS) * d, d, bias=False)
        self.scale = nn.Parameter(torch.tensor(1.0))

    def forward(self, bag_inputs):
        parts = []
        for b in BAGS:
            idx, off, w = bag_inputs[b]
            parts.append(self.emb[b](idx, off, per_sample_weights=w))
        x = self.norm(torch.cat(parts, dim=1))
        return (self.skip(x) + self.net(x)) * self.scale


class TwoTower(nn.Module):
    def __init__(self, n_items, d=128, h=256):
        super().__init__()
        self.user = UserTower(n_items, d, h)
        self.item = ItemTower(n_items, d, h)

    def score(self, bag_inputs, item_idx):
        return self.user(bag_inputs) @ self.item(item_idx).T
EOF
echo ok

cat > pipeline.py <<'EOF'
import numpy as np, pandas as pd, torch, time, sys, os
import features as FT
from features import N_ITEMS, DAY
from model import TwoTower, BAGS

torch.set_num_threads(8)
S = '/app/solution/'
D = '/task/data/'
T_TRAIN = 24786800
VAL_LO, VAL_HI = 24786800, 25393600


def load_events():
    ev_tr, ev_va = {}, {}
    z = np.load(S + 'train_listens.npz'); ev_tr['listens'] = dict(z)
    z = np.load(S + 'val_listens.npz'); ev_va['listens'] = dict(z)
    for f in ('likes', 'dislikes', 'unlikes', 'undislikes'):
        ev_tr[f] = dict(np.load(S + f'train_{f}.npz'))
        ev_va[f] = dict(np.load(S + f'val_{f}.npz'))
    return ev_tr, ev_va


def merge(ev_tr, ev_va, uid_map):
    ev = {}
    for f in ev_tr:
        d = {}
        for k in ('uid', 'ts', 'item'):
            d[k] = np.concatenate([ev_tr[f][k], ev_va[f][k]])
        for k in ('org', 'strong'):
            if k in ev_tr[f]:
                d[k] = np.concatenate([ev_tr[f][k], ev_va[f][k]])
        d['uidx'] = uid_map[d['uid']]
        ev[f] = d
    return ev


def hist_keys(ev, cutoff):
    ks = []
    for f, d in ev.items():
        m = d['ts'] < cutoff
        ks.append(d['uidx'][m].astype(np.int64) * N_ITEMS + d['item'][m])
    return np.unique(np.concatenate(ks))


def make_slice(ev, cutoff, lo, hi, n_users):
    bags = FT.build_bags(ev, cutoff, None, n_users)
    hk = hist_keys(ev, cutoff)
    u, i, g = FT.build_targets(ev, lo, hi, None, None, hk)
    return bags, (u, i, g)


class Ctx:
    """Concatenated per-slice user bags addressed by a flat context id."""

    def __init__(self):
        self.ptr = {b: [np.zeros(1, np.int64)] for b in BAGS}
        self.items = {b: [] for b in BAGS}
        self.w = {b: [] for b in BAGS}
        self.n = 0

    def add(self, bags, n_users):
        for b in BAGS:
            ptr, it, w = bags[b]
            self.ptr[b].append(ptr[1:] + self.ptr[b][-1][-1])
            self.items[b].append(it)
            self.w[b].append(w)
        base = self.n
        self.n += n_users
        return base

    def finish(self):
        for b in BAGS:
            self.ptr[b] = np.concatenate(self.ptr[b])
            self.items[b] = np.concatenate(self.items[b]) if self.items[b] else np.zeros(0, np.int32)
            wv = np.concatenate(self.w[b]) if self.w[b] else np.zeros(0, np.float32)
            self.w[b] = np.log1p(wv).astype(np.float32)

    def batch(self, ctx_ids):
        out = {}
        for b in BAGS:
            ptr = self.ptr[b]
            s, e = ptr[ctx_ids], ptr[ctx_ids + 1]
            ln = (e - s).astype(np.int64)
            off = np.concatenate([[0], np.cumsum(ln)[:-1]])
            tot = int(ln.sum())
            if tot == 0:
                idx = np.zeros(len(ctx_ids), np.int64); wt = np.zeros(len(ctx_ids), np.float32)
                off = np.arange(len(ctx_ids), dtype=np.int64)
            else:
                gather = np.repeat(s, ln) + (np.arange(tot) - np.repeat(off, ln))
                idx = self.items[b][gather].astype(np.int64)
                wt = self.w[b][gather]
                # L2-ish normalisation per bag so long histories do not dominate
                nrm = np.zeros(len(ctx_ids), np.float32)
                np.add.at(nrm, np.repeat(np.arange(len(ctx_ids)), ln), wt)
                nrm = np.maximum(nrm, 1e-6)
                wt = wt / np.repeat(np.sqrt(nrm), ln)
                empty = ln == 0
                if empty.any():
                    # EmbeddingBag needs valid offsets; empty bags contribute zero weight
                    idx = np.concatenate([idx, np.zeros(int(empty.sum()), np.int64)])
                    wt = np.concatenate([wt, np.zeros(int(empty.sum()), np.float32)])
                    ln2 = ln.copy(); ln2[empty] = 1
                    off = np.concatenate([[0], np.cumsum(ln2)[:-1]])
                    order = np.argsort(np.concatenate([np.repeat(np.arange(len(ctx_ids)), ln),
                                                       np.nonzero(empty)[0]]), kind='stable')
                    idx, wt = idx[order], wt[order]
            out[b] = (torch.from_numpy(np.ascontiguousarray(idx)),
                      torch.from_numpy(np.ascontiguousarray(off)),
                      torch.from_numpy(np.ascontiguousarray(wt)))
        return out
EOF
echo ok

cat > model.py <<'EOF'
"""PyTorch two-tower ranker.

The user tower pools a user's interaction history (no user-ID parameters); the item
tower maps an item to an independently precomputable vector. The relevance score is
exactly the dot product of the two vectors, so item vectors can be built offline into
a MIPS index and request time costs one user-tower forward pass plus one ANN lookup.
The two towers share a single item-embedding table but apply separate MLP heads.
"""
import torch, torch.nn as nn, torch.nn.functional as F

BAGS = ('all', 'recent', 'liked', 'disliked')


class TwoTower(nn.Module):
    def __init__(self, n_items, d=128, h=256):
        super().__init__()
        self.d = d
        self.emb = nn.Parameter(torch.randn(n_items, d) * 0.05)
        nb = len(BAGS)
        self.unorm = nn.LayerNorm(nb * d)
        self.umlp = nn.Sequential(nn.Linear(nb * d, h), nn.GELU(), nn.Linear(h, d))
        self.uskip = nn.Linear(nb * d, d, bias=False)
        self.inorm = nn.LayerNorm(d)
        self.imlp = nn.Sequential(nn.Linear(d, h), nn.GELU(), nn.Linear(h, d))
        self.iskip = nn.Linear(d, d, bias=False)
        nn.init.eye_(self.iskip.weight)
        self.logit_scale = nn.Parameter(torch.tensor(2.0))

    def item_vec(self, idx):
        e = self.emb[idx]
        e = self.inorm(e)
        return self.iskip(e) + self.imlp(e)

    def user_vec(self, bag_inputs):
        parts = [F.embedding_bag(idx, self.emb, off, mode='sum', per_sample_weights=w)
                 for idx, off, w in (bag_inputs[b] for b in BAGS)]
        x = self.unorm(torch.cat(parts, dim=1))
        u = self.uskip(x) + self.umlp(x)
        return u * self.logit_scale
EOF
echo ok

cat > run.py <<'EOF'
import numpy as np, pandas as pd, torch, time, sys, os
import features as FT
from features import N_ITEMS, DAY
from model import TwoTower, BAGS
from pipeline import Ctx, load_events, merge, hist_keys, T_TRAIN, VAL_LO, VAL_HI, S, D
import common

torch.set_num_threads(8)
t0 = time.time()
STAGE = sys.argv[1] if len(sys.argv) > 1 else 'val'
STEPS = int(sys.argv[2]) if len(sys.argv) > 2 else 1500
SEED = int(sys.argv[3]) if len(sys.argv) > 3 else 0

ev_tr, ev_va = load_events()
all_uid = np.unique(np.concatenate([ev_tr['listens']['uid'], ev_va['listens']['uid'],
                                    common.targets.astype(np.int32)]))
uid_map = np.full(int(all_uid.max()) + 1, -1, np.int32)
uid_map[all_uid] = np.arange(len(all_uid), dtype=np.int32)
n_users = len(all_uid)
ev = merge(ev_tr, ev_va, uid_map)
print('users', n_users, 'load', time.time() - t0, flush=True)

WEEK = 7 * DAY
slices = [(T_TRAIN - k * WEEK, T_TRAIN - k * WEEK, T_TRAIN - (k - 1) * WEEK) for k in range(1, 7)]
if STAGE == 'final':
    slices = [(T_TRAIN, VAL_LO, VAL_HI)] + slices

ctx = Ctx()
tr_ctx, tr_item, tr_g = [], [], []
for cutoff, lo, hi in slices:
    bags = FT.build_bags(ev, cutoff, None, n_users)
    base = ctx.add(bags, n_users)
    hk = hist_keys(ev, cutoff)
    u, i, g = FT.build_targets(ev, lo, hi, None, None, hk)
    tr_ctx.append(u.astype(np.int64) + base); tr_item.append(i.astype(np.int64)); tr_g.append(g)
    print('slice', cutoff, len(u), round(time.time() - t0, 1), flush=True)

INF_CUTOFF = VAL_HI if STAGE == 'final' else T_TRAIN
bags = FT.build_bags(ev, INF_CUTOFF, None, n_users)
inf_base = ctx.add(bags, n_users)
ctx.finish()
tr_ctx = np.concatenate(tr_ctx); tr_item = np.concatenate(tr_item); tr_g = np.concatenate(tr_g)
print('pairs', len(tr_ctx), round(time.time() - t0, 1), flush=True)

# negative sampling distribution
cnt = np.bincount(ev['listens']['item'][ev['listens']['ts'] < T_TRAIN], minlength=N_ITEMS).astype(np.float64)
pq = (cnt + 1.0) ** 0.75; pq /= pq.sum()

torch.manual_seed(SEED); np.random.seed(SEED)
model = TwoTower(N_ITEMS, d=128, h=256)
opt = torch.optim.AdamW(model.parameters(), lr=2.5e-3, weight_decay=1e-6)
sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=2.5e-3, total_steps=STEPS, pct_start=0.15)
B, NEG = 512, 3072
w_g = np.where(tr_g >= 3, 2.0, 1.0).astype(np.float32)
rng = np.random.default_rng(SEED)
model.train()
for step in range(STEPS):
    sel = rng.integers(0, len(tr_ctx), B)
    c, pos, wt = tr_ctx[sel], tr_item[sel], w_g[sel]
    neg = np.concatenate([rng.integers(0, N_ITEMS, NEG // 2), rng.choice(N_ITEMS, NEG // 2, p=pq)])
    cand = np.concatenate([pos, neg])
    bi = ctx.batch(c)
    u = model.user_vec(bi)
    v = model.item_vec(torch.from_numpy(cand))
    logits = u @ v.T
    # mask other in-batch positives so they are not treated as negatives for this row
    same = torch.from_numpy((pos[:, None] == cand[None, :]))
    tgt = torch.arange(B)
    same[tgt, tgt] = False
    logits = logits.masked_fill(same, -1e9)
    loss = (torch.nn.functional.cross_entropy(logits, tgt, reduction='none')
            * torch.from_numpy(wt)).mean()
    opt.zero_grad(); loss.backward(); opt.step(); sched.step()
    if step % 100 == 0:
        print(step, round(float(loss), 4), round(time.time() - t0, 1), flush=True)

# ---- inference ----
model.eval()
uids_e, bits = common.load_elig('validation' if STAGE == 'val' else 'test')
rows_ctx = uid_map[uids_e.astype(np.int64)].astype(np.int64) + inf_base
with torch.no_grad():
    V = model.item_vec(torch.arange(N_ITEMS)).numpy()
    US = []
    for s in range(0, len(rows_ctx), 512):
        US.append(model.user_vec(ctx.batch(rows_ctx[s:s + 512])).numpy())
    U = np.concatenate(US)
scores = U @ V.T
np.save(S + f'scores_{STAGE}.npy', scores.astype(np.float32))
torch.save(model.state_dict(), S + f'model_{STAGE}.pt')

if STAGE == 'val':
    G = np.load(S + 'G_val.npy')
    print('NDCG@10', common.ndcg10(scores, G, bits), flush=True)

# write ranking
s = np.where(bits, scores, -1e30)
top = np.argpartition(-s, 100, axis=1)[:, :100]
ts = np.take_along_axis(s, top, 1)
top = np.take_along_axis(top, np.argsort(-ts, axis=1), 1)
out = pd.DataFrame({'uid': np.repeat(uids_e, 100),
                    'item_id': common.item_ids[top.ravel()],
                    'rank': np.tile(np.arange(1, 101), len(uids_e))})
name = 'validation.parquet' if STAGE == 'val' else 'submission.parquet'
out.to_parquet(S + name, index=False)
print('wrote', name, round(time.time() - t0, 1), flush=True)
EOF
timeout 1200 python run.py val 900 2>&1 | tail -40

Options

{
  "timeout": 1300000
}

python - <<'EOF'
import re
p='run.py'; s=open(p).read()
s=s.replace("""all_uid = np.unique(np.concatenate([ev_tr['listens']['uid'], ev_va['listens']['uid'],
                                    common.targets.astype(np.int32)]))""",
"""_us = [common.targets.astype(np.int32)]
for _e in (ev_tr, ev_va):
    for _f in _e:
        _us.append(_e[_f]['uid'].astype(np.int32))
all_uid = np.unique(np.concatenate(_us))""")
open(p,'w').write(s)
EOF
timeout 1500 python run.py val 900 2>&1 | tail -40

Options

{
  "timeout": 1600000
}

cp validation.parquet validation_v1.parquet && python /task/evaluate_public.py validation_v1.parquet

Options

{
  "timeout": 600000
}

sed -i 's/d=128, h=256/d=int(os.environ.get("DIM",128)), h=int(os.environ.get("HID",256))/' run.py && grep -n 'TwoTower(N_ITEMS' run.py

python - <<'EOF'
p='model.py'; s=open(p).read()
s=s.replace("""        self.logit_scale = nn.Parameter(torch.tensor(2.0))""",
"""        self.item_bias = nn.Parameter(torch.zeros(n_items))
        self.logit_scale = nn.Parameter(torch.tensor(2.0))""")
s=s.replace("""        e = self.inorm(e)
        return self.iskip(e) + self.imlp(e)""",
"""        e = self.inorm(e)
        v = self.iskip(e) + self.imlp(e)
        # last coordinate carries a per-item popularity bias; the user tower emits a
        # constant 1 there, so the score stays an exact dot product of equal-dim vectors
        return torch.cat([v, self.item_bias[idx, None]], dim=1)""")
s=s.replace("""        u = self.uskip(x) + self.umlp(x)
        return u * self.logit_scale""",
"""        u = (self.uskip(x) + self.umlp(x)) * self.logit_scale
        return torch.cat([u, torch.ones(u.shape[0], 1)], dim=1)""")
open(p,'w').write(s)
EOF
DIM=160 HID=384 nohup python run.py val 4000 > log_v2.txt 2>&1 &
sleep 60; tail -5 log_v2.txt

Options

{
  "timeout": 300000
}

python - <<'EOF'
import numpy as np, pandas as pd, common
uids,bits=common.load_elig('test')
cat=pd.read_parquet('/task/data/candidate_catalog.parquet')
pop=-cat.popularity_rank.to_numpy(np.float32)
s=np.where(bits,pop[None,:],-1e30)
top=np.argpartition(-s,100,axis=1)[:,:100]
ts=np.take_along_axis(s,top,1); top=np.take_along_axis(top,np.argsort(-ts,axis=1),1)
pd.DataFrame({'uid':np.repeat(uids,100),'item_id':common.item_ids[top.ravel()],'rank':np.tile(np.arange(1,101),len(uids))}).to_parquet('/app/solution/submission.parquet',index=False)
print('fallback written',len(uids))
EOF
sleep 120; tail -4 log_v2.txt

Options

{
  "timeout": 400000
}

sleep 150; tail -6 log_v2.txt

Options

{
  "timeout": 400000
}

python - <<'EOF'
p='run.py'; s=open(p).read()
s=s.replace("tr_ctx, tr_item, tr_g = [], [], []","tr_ctx, tr_item, tr_g, tr_s = [], [], [], []")
s=s.replace("""for cutoff, lo, hi in slices:""","""for si, (cutoff, lo, hi) in enumerate(slices):""")
s=s.replace("""    tr_ctx.append(u.astype(np.int64) + base); tr_item.append(i.astype(np.int64)); tr_g.append(g)""",
"""    tr_ctx.append(u.astype(np.int64) + base); tr_item.append(i.astype(np.int64)); tr_g.append(g)
    tr_s.append(np.full(len(u), si, np.int32))""")
s=s.replace("""tr_ctx = np.concatenate(tr_ctx); tr_item = np.concatenate(tr_item); tr_g = np.concatenate(tr_g)""",
"""tr_ctx = np.concatenate(tr_ctx); tr_item = np.concatenate(tr_item); tr_g = np.concatenate(tr_g)
tr_s = np.concatenate(tr_s)
age = tr_s - (1 if STAGE == 'final' else 0)
slice_w = np.power(0.85, np.maximum(age, 0)).astype(np.float32)""")
s=s.replace("""w_g = np.where(tr_g >= 3, 2.0, 1.0).astype(np.float32)""",
"""w_g = (np.where(tr_g >= 3, 2.0, 1.0) * slice_w).astype(np.float32)""")

# evaluation harness inside training loop
s=s.replace("""model.train()
for step in range(STEPS):""",
"""uids_e, bits = common.load_elig('validation' if STAGE == 'val' else 'test')
G = np.load(S + 'G_val.npy') if STAGE == 'val' else None


def infer():
    model.eval()
    with torch.no_grad():
        V = model.item_vec(torch.arange(N_ITEMS)).numpy()
        US = [model.user_vec(ctx.batch(rows_ctx[s_:s_ + 1024])).numpy()
              for s_ in range(0, len(rows_ctx), 1024)]
    model.train()
    return np.concatenate(US) @ V.T


rows_ctx = uid_map[uids_e.astype(np.int64)].astype(np.int64) + inf_base
best, best_scores, best_step = -1.0, None, 0
EVERY = int(os.environ.get('EVERY', 400))
model.train()
for step in range(STEPS):""")
s=s.replace("""    if step % 100 == 0:
        print(step, round(float(loss), 4), round(time.time() - t0, 1), flush=True)""",
"""    if step % 100 == 0:
        print(step, round(float(loss.detach()), 4), round(time.time() - t0, 1), flush=True)
    if STAGE == 'val' and (step + 1) % EVERY == 0:
        sc = infer()
        nd = common.ndcg10(sc, G, bits)[0]
        print('  eval', step + 1, round(nd, 5), round(time.time() - t0, 1), flush=True)
        if nd > best:
            best, best_scores, best_step = nd, sc, step + 1""")
s=s.replace("""# ---- inference ----
model.eval()
uids_e, bits = common.load_elig('validation' if STAGE == 'val' else 'test')
rows_ctx = uid_map[uids_e.astype(np.int64)].astype(np.int64) + inf_base
with torch.no_grad():
    V = model.item_vec(torch.arange(N_ITEMS)).numpy()
    US = []
    for s in range(0, len(rows_ctx), 512):
        US.append(model.user_vec(ctx.batch(rows_ctx[s:s + 512])).numpy())
    U = np.concatenate(US)
scores = U @ V.T
np.save(S + f'scores_{STAGE}.npy', scores.astype(np.float32))
torch.save(model.state_dict(), S + f'model_{STAGE}.pt')

if STAGE == 'val':
    G = np.load(S + 'G_val.npy')
    print('NDCG@10', common.ndcg10(scores, G, bits), flush=True)
""",
"""# ---- inference ----
scores = best_scores if (STAGE == 'val' and best_scores is not None) else infer()
if STAGE == 'val':
    print('BEST NDCG@10', round(best, 6), 'at step', best_step, flush=True)
np.save(S + f'scores_{STAGE}.npy', scores.astype(np.float32))
torch.save(model.state_dict(), S + f'model_{STAGE}.pt')
""")
s=s.replace("""s = np.where(bits, scores, -1e30)""","""sq = np.where(bits, scores, -1e30)""")
s=s.replace("""top = np.argpartition(-s, 100, axis=1)[:, :100]
ts = np.take_along_axis(s, top, 1)""","""top = np.argpartition(-sq, 100, axis=1)[:, :100]
ts = np.take_along_axis(sq, top, 1)""")
open(p,'w').write(s)
print('patched')
EOF
DIM=128 HID=256 nohup python run.py val 4000 > log_v3.txt 2>&1 &
sleep 90; tail -8 log_v3.txt

Options

{
  "timeout": 300000
}

grep eval log_v3.txt; sleep 100; grep eval log_v3.txt

Options

{
  "timeout": 300000
}

tail -3 log_v3.txt; ps aux|grep -c "[r]un.py"

Options

{
  "timeout": 120000
}

python - <<'EOF'
p='common.py'; s=open(p).read()
s=s.replace("""    s=np.where(bits,scores,-1e30)
    top=np.argpartition(-s,10,axis=1)[:,:10]
    ts=np.take_along_axis(s,top,1)
    o=np.argsort(-ts,axis=1)
    top=np.take_along_axis(top,o,1)
    g=np.take_along_axis(G,top,1).astype(np.float64)
    disc=1.0/np.log2(np.arange(2,12))
    dcg=((2.0**g-1)*disc).sum(1)""",
"""    disc=1.0/np.log2(np.arange(2,12))
    dcg=np.zeros(len(scores))
    for a in range(0,len(scores),512):
        b=min(a+512,len(scores))
        s=np.where(bits[a:b],scores[a:b],np.float32(-1e30))
        top=np.argpartition(-s,10,axis=1)[:,:10]
        ts=np.take_along_axis(s,top,1)
        top=np.take_along_axis(top,np.argsort(-ts,axis=1),1)
        g=np.take_along_axis(G[a:b],top,1).astype(np.float64)
        dcg[a:b]=((2.0**g-1)*disc).sum(1)
        del s""")
open(p,'w').write(s)
p='run.py'; s=open(p).read()
s=s.replace("slice_w = np.power(0.85,","slice_w = np.power(float(os.environ.get('SDECAY',1.0)),")
s=s.replace("            best, best_scores, best_step = nd, sc, step + 1",
            "            best, best_scores, best_step = nd, sc.astype(np.float16), step + 1\n        del sc")
s=s.replace("scores = best_scores if (STAGE == 'val' and best_scores is not None) else infer()",
            "scores = best_scores.astype(np.float32) if (STAGE == 'val' and best_scores is not None) else infer()")
open(p,'w').write(s)
print('ok')
EOF
DIM=160 HID=384 EVERY=500 SDECAY=1.0 nohup python run.py val 3000 > log_v4.txt 2>&1 &
sleep 115; grep -E "eval|BEST" log_v4.txt; tail -2 log_v4.txt

Options

{
  "timeout": 300000
}

pkill -f "run.py" ; python - <<'EOF'
p='model.py'; s=open(p).read()
s=s.replace("class TwoTower(nn.Module):\n    def __init__(self, n_items, d=128, h=256):\n        super().__init__()",
 "import os\n\n\nclass TwoTower(nn.Module):\n    def __init__(self, n_items, d=128, h=256):\n        super().__init__()\n        self.use_bias = os.environ.get('BIAS', '1') == '1'")
s=s.replace("        return torch.cat([v, self.item_bias[idx, None]], dim=1)",
 "        if not self.use_bias:\n            return v\n        return torch.cat([v, self.item_bias[idx, None]], dim=1)")
s=s.replace("        return torch.cat([u, torch.ones(u.shape[0], 1)], dim=1)",
 "        if not self.use_bias:\n            return u\n        return torch.cat([u, torch.ones(u.shape[0], 1)], dim=1)")
open(p,'w').write(s)
EOF
DIM=128 HID=256 EVERY=250 SDECAY=1.0 BIAS=0 nohup python run.py val 2000 > log_v5.txt 2>&1 &
sleep 105; grep -E "eval|BEST" log_v5.txt; tail -2 log_v5.txt

Options

{
  "timeout": 300000
}

grep -E "eval|BEST|Error|error" log_v5.txt | tail; tail -3 log_v5.txt; ps aux|grep -c "[r]un.py"

Options

{
  "timeout": 120000
}

DIM=128 HID=256 EVERY=250 SDECAY=1.0 BIAS=0 setsid nohup python run.py val 2000 > log_v5.txt 2>&1 < /dev/null & 
sleep 110; grep -E "eval|BEST" /app/solution/log_v5.txt; tail -2 /app/solution/log_v5.txt

Options

{
  "timeout": 300000
}

DIM=128 HID=256 SDECAY=1.0 BIAS=0 setsid nohup python run.py final 900 > log_final.txt 2>&1 < /dev/null &
sleep 100; echo "--v5--"; grep -E "eval|BEST" log_v5.txt|tail -4; echo "--final--"; tail -3 log_final.txt

Options

{
  "timeout": 300000
}

cat > README.md <<'EOF'
# Two-tower discovery ranker (CPU PyTorch)

## Architecture (`model.py`)
A single `TwoTower` module with two heads over one shared 30,000 x 128 item-embedding
table (the retrieval catalog is the only ID vocabulary; there is **no user-ID table**).

* **Item tower** — `emb[i] -> LayerNorm -> (identity-initialised linear skip + 2-layer
  GELU MLP) -> R^128`. Depends on the item alone, so all 30,000 item vectors are
  precomputed once and loaded into a MIPS/ANN index.
* **User tower** — four `embedding_bag` pools over the same table, built purely from the
  user's interaction history at the scoring cutoff:
  `all` positive listens (top 256), `recent` positive listens in the last 30 days
  (top 128), currently `liked` items and currently `disliked` items (explicit state, net
  of unlikes/undislikes). Bag weights are `log1p(count)`, L2-style normalised so long
  histories do not dominate. The four pooled vectors are concatenated,
  `LayerNorm`-ed, and passed through a linear skip + 2-layer GELU MLP to `R^128`.
* **Score** — exactly `dot(user_vec, item_vec)`. Nothing else enters the score, so
  request time is one user-tower forward pass plus one inner-product search.

Training never materialises or factorises a global user-item matrix: gradients only
touch the items in a sampled mini-batch, so the method scales to millions of
users/items.

## Training objective (`run.py`)
Sampled-softmax (InfoNCE) over novel positives, using six weekly temporal slices of the
training period. For each slice the user tower sees only history strictly before the
slice cutoff, and the positives are the *novel* organic outcomes inside the following
seven days — the same event definition the metric uses (organic like -> grade 3, weight
2; organic listen >= 80% -> grade 1, weight 1). Items already in the user's history are
excluded from the targets, matching the platform's candidate-eligibility rule (verified:
eligibility removes 100% of previously-interacted items, so this is a pure
new-item-discovery problem).

Each step draws 512 (history, positive) pairs and 3,072 shared negatives (half uniform,
half popularity^0.75); other in-batch positives are masked out of the row's negatives.
AdamW + OneCycle, 900 steps.

## Data pipeline
* `load.py` — streams the 43M-row listen log once, keeps positive listens (>= 50%
  completion) on catalog items, and caches compact `.npz` arrays.
* `features.py` — builds the capped, weighted history bags and the novel-positive
  targets at an arbitrary cutoff.
* `pipeline.py` — concatenates per-slice bags into a flat CSR "context" store and
  assembles `embedding_bag` batches.
* `common.py` — eligibility unpacking, validation ground truth, NDCG@10.

## Results
Public validation NDCG@10 (`evaluate_public.py`): **0.014107** over 4,263 scored users,
against **0.010710** for the popularity ordering of the same eligible candidates
(+32%).

## Final refit and submission
The selected configuration is refit on `train_* + val_*`: a seventh slice is added whose
history is everything up to the training cutoff and whose targets are the validation
window's novel organic outcomes, and inference uses history through the end of the
validation window (train + val as observed history) with `test_eligible_candidates.npz`.

```sh
python load.py                 # one-time event caching
python run.py val 900          # validation ranker -> solution/validation.parquet
python run.py final 900        # hidden-test ranker -> solution/submission.parquet
```

Environment knobs: `DIM`, `HID`, `BIAS`, `SDECAY`, `EVERY` (defaults used above:
`DIM=128 HID=256 BIAS=0 SDECAY=1.0`).
EOF
sleep 65; tail -3 log_final.txt

Options

{
  "timeout": 300000
}

sleep 90; tail -3 log_final.txt; ps aux|grep "[r]un.py"|wc -l

Options

{
  "timeout": 300000
}

ps aux|grep "[r]un.py"|awk '{print $2,$3,$4,$12,$13,$14}'; tail -2 log_v5.txt; free -g|head -2; tail -2 log_final.txt

Options

{
  "timeout": 120000
}

kill -9 8327; sleep 75; tail -3 /app/solution/log_final.txt

Options

{
  "timeout": 200000
}

python - <<'EOF'
import pandas as pd, numpy as np, common
from ranking_evaluation import CandidateSet, validate_ranking, read_ranking
from pathlib import Path
f=read_ranking(Path('submission.parquet'))
n=validate_ranking(f,set(common.targets.astype(int)))
cs=CandidateSet(Path('/task/data/test_eligible_candidates.npz'))
uids,bits=common.load_elig('test')
row={int(u):r for r,u in enumerate(uids)}
im={int(i):c for c,i in enumerate(common.item_ids)}
r=np.array([row[int(u)] for u in n.uid]); c=np.array([im[int(i)] for i in n.item_id])
print('rows',len(n),'all eligible',bool(bits[r,c].all()))
print(n.head(3))
EOF
cp validation_v1.parquet validation.parquet; ls -la /app/solution/*.parquet /app/solution/*.md

Options

{
  "timeout": 300000
}

PYTHONPATH=/task python - <<'EOF'
import pandas as pd, numpy as np, common
from ranking_evaluation import validate_ranking, read_ranking
from pathlib import Path
n=validate_ranking(read_ranking(Path('submission.parquet')),set(common.targets.astype(int)))
uids,bits=common.load_elig('test')
row={int(u):r for r,u in enumerate(uids)}
im={int(i):c for c,i in enumerate(common.item_ids)}
r=np.array([row[int(u)] for u in n.uid]); c=np.array([im[int(i)] for i in n.item_id])
print('rows',len(n),'users',n.uid.nunique(),'all eligible',bool(bits[r,c].all()))
EOF
python /task/evaluate_public.py /app/solution/validation.parquet

Options

{
  "timeout": 300000
}

rm -f /app/solution/scores_val.npy /app/solution/scores_final.npy /app/solution/G_val.npy /app/solution/elig_val.npy /app/solution/uids_val.npy; ls /app/solution/

Options

{
  "timeout": 120000
}