Revision 5: negative-result model card, report rev. 5 (PDF/MD/figure), CITATION.cff, iterations 020-021 code, logs, sealed tests and checkpoints
7420c83 verified | # CUERDAS-020: sealed test (evaluated once), compact text baselines, stronger no-LLM decoders. | |
| import os, sys, json, hashlib, random, math, platform | |
| import torch, torch.nn as nn, torch.nn.functional as F, transformers | |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) | |
| import run018 as R | |
| from transformers import AutoModelForCausalLM, AutoConfig | |
| INP = os.environ['INP']; OD = os.environ['OD']; NBOOT = 10000; SEEDS = [17, 29, 43] | |
| raw = open(f'{INP}/sealed020.json').read() | |
| H = hashlib.sha256(raw.encode()).hexdigest(); assert H == '4c7737f66364a44786b4d9acb4686a6f31829d78cb5914f19eca1df74fe73393', H | |
| S = json.loads(raw); NEW = S['new_words'] | |
| SUMM = open(f'{OD}/results020.jsonl', 'w'); PRED = open(f'{OD}/predictions020.jsonl', 'w') | |
| def log(d): SUMM.write(json.dumps(d) + '\n'); SUMM.flush(); print(json.dumps(d), flush=True) | |
| log({'env': {'python': platform.python_version(), 'torch': torch.__version__, 'transformers': transformers.__version__, 'device': torch.cuda.get_device_name() if R.DEV == 'cuda' else 'cpu'}, 'sealed_sha256': H}) | |
| for j in range(4): | |
| for w in NEW[j]: assert len(R.t1(w)) == 1, w | |
| CAND = [torch.tensor(sorted({R.t1(w)[0] for w in R.V[j] + NEW[j]})) for j in range(4)] | |
| assert all(len(CAND[j]) == len(R.V[j]) + 8 for j in range(4)) | |
| def gidx(j, w): return int((CAND[j] == R.t1(w)[0]).nonzero()[0]) | |
| SETS = ['sealed_id', 'sealed_ood', 'sealed_1_new', 'sealed_2_new', 'sealed_4_new'] | |
| FMT = {'template': R.TPL, 'kv': "action: {action}; target: {target}; status: {status}; window: {window}.", | |
| 'abbrev': "act={action} tgt={target} st={status} win={window}", 'csv': "action,target,status,window\n{action},{target},{status},{window}", | |
| 'values_csv': "{action}, {target}, {status}, {window}", 'values': "{action} {target} {status} {window}"} | |
| def text_logits(msgs, fmt): | |
| embs, mt = [], [] | |
| for i, m in enumerate(msgs): | |
| for j in range(4): | |
| embs.append(R.Emb[R.tok.encode(fmt.format(**dict(zip(R.FIELDS, m))) + R.SUF_MAIN.format(field=R.FIELDS[j]), add_special_tokens=False)]); mt.append((i, j)) | |
| return R.run_rows(embs, None), mt | |
| def per_ex(msgs, fmt=None, cod=None, bs=24): | |
| out = [] | |
| for b in range(0, len(msgs), bs): | |
| mb = msgs[b:b+bs]; lg, mt = text_logits(mb, fmt) if cod is None else R.codec_logits(cod, mb, R.SUF_MAIN) | |
| good = [True] * len(mb) | |
| for r, (i, j) in enumerate(mt): good[i] &= int(lg[r][CAND[j]].argmax()) == gidx(j, mb[i][j]) | |
| out += good | |
| return [float(g) for g in out] | |
| def wilson(k, n, z=1.96): | |
| p = k / n; d = 1 + z*z/n; c = (p + z*z/(2*n)) / d; h = z * math.sqrt(p*(1-p)/n + z*z/(4*n*n)) / d | |
| return [round(100*(c-h), 1), round(100*(c+h), 1)] | |
| def boot(a, b, seed=0): | |
| g = torch.Generator().manual_seed(seed); a = torch.tensor(a, dtype=torch.float64, device='cpu'); b = torch.tensor(b, dtype=torch.float64, device='cpu') | |
| idx = torch.randint(0, len(a), (NBOOT, len(a)), generator=g, device='cpu'); d = (a[idx].mean(1) - b[idx].mean(1)) * 100 | |
| return round(float((a.mean()-b.mean())*100), 2), [round(float(d.quantile(.025)), 2), round(float(d.quantile(.975)), 2)] | |
| def positions(msgs, fmt): return sum(len(R.tok.encode(fmt.format(**dict(zip(R.FIELDS, m))), add_special_tokens=False)) for m in msgs) / len(msgs) + 15 | |
| # ---- text formats ---- | |
| TXT = {} | |
| for fn, fmt in FMT.items(): | |
| for s in SETS: | |
| r = per_ex(S[s], fmt); TXT[(fn, s)] = r; k = int(sum(r)) | |
| for i, c in enumerate(r): PRED.write(json.dumps({'route': 'text_' + fn, 'set': s, 'i': i, 'correct': bool(c)}) + '\n') | |
| e = {'route': 'text_' + fn, 'set': s, 'n': len(r), 'acc': k/len(r), 'wilson95': wilson(k, len(r)), 'positions_per_event': round(positions(S[s], fmt), 2)} | |
| if fn != 'template': e['diff_vs_template'], e['boot95'] = boot(r, TXT[('template', s)], seed=7) | |
| log(e) | |
| # ---- K=2, sealed, evaluated once ---- | |
| cods = {}; K2 = {} | |
| for sd in SEEDS: | |
| cod = R.PackK(2, sd); cod.load_state_dict(torch.load(f'{INP}/k2long_s{sd}.pt', map_location=R.DEV)); cod.eval().requires_grad_(False); cods[sd] = cod | |
| for s in SETS: | |
| r = per_ex(S[s], cod=cod); K2[(sd, s)] = r; k = int(sum(r)) | |
| for i, c in enumerate(r): PRED.write(json.dumps({'route': f'k2_s{sd}', 'set': s, 'i': i, 'correct': bool(c)}) + '\n') | |
| d, ci = boot(r, TXT[('template', s)], seed=sd) | |
| log({'route': f'k2_s{sd}', 'set': s, 'n': len(r), 'acc': k/len(r), 'wilson95': wilson(k, len(r)), 'positions_per_event': 17, 'diff_vs_template': d, 'boot95': ci}) | |
| for s in SETS: | |
| a = [sum(K2[(sd, s)][i] for sd in SEEDS)/3 for i in range(len(S[s]))] | |
| for fn in FMT: | |
| d, ci = boot(a, TXT[(fn, s)], seed=1); log({'route': 'k2_mean3', 'set': s, 'acc': sum(a)/len(a), 'vs': 'text_' + fn, 'diff': d, 'boot95': ci}) | |
| # ---- no-LLM decoders on K=2 slots ---- | |
| def slots(cod, ms): | |
| with torch.no_grad(): return torch.cat([cod(torch.stack([torch.stack([R.lex(w) for w in m]) for m in ms[i:i+512]])).flatten(1) for i in range(0, len(ms), 512)]) | |
| g = random.Random(2020); AUG = [] | |
| for _ in range(20000): | |
| m = list(g.choice(R.FIT)) | |
| for j in range(4): | |
| if g.random() < .4: m[j] = g.choice(R.POOL) | |
| AUG.append(m) | |
| YA = torch.stack([torch.stack([R.lex(w) for w in m]) for m in AUG]) # [N,4,D] | |
| CE = [F.normalize(R.Emb[CAND[j]], dim=-1) for j in range(4)] | |
| def dec_emb(P, ms): # P [n,4,D] predicted value embeddings -> exact-match | |
| ok = torch.ones(len(ms), dtype=torch.bool, device='cpu') | |
| for j in range(4): ok &= ((F.normalize(P[:, j], dim=-1) @ CE[j].T).argmax(1).cpu() == torch.tensor([gidx(j, m[j]) for m in ms], device='cpu')) | |
| return ok.float().tolist() | |
| def train_mlp(X, Y, shuffle=False, steps=4000, seed=0): | |
| torch.manual_seed(seed); Y = Y[torch.randperm(len(Y))] if shuffle else Y | |
| mu, sdv = X.mean(0), X.std(0) + 1e-6 | |
| net = nn.Sequential(nn.Linear(X.shape[1], 1024), nn.GELU(), nn.Linear(1024, 1024), nn.GELU(), nn.Linear(1024, 4*R.D)); opt = torch.optim.Adam(net.parameters(), 1e-3) | |
| for _ in range(steps): | |
| b = torch.randint(0, len(X), (256,)); p = net((X[b]-mu)/sdv).view(-1, 4, R.D) | |
| loss = (1 - F.cosine_similarity(p, Y[b], -1)).mean(); opt.zero_grad(); loss.backward(); opt.step() | |
| net.eval(); return lambda Z: net((Z-mu)/sdv).view(-1, 4, R.D).detach() | |
| PROBE = {} | |
| for sd in SEEDS: | |
| cod = cods[sd]; XA = slots(cod, AUG); XF = slots(cod, R.FIT); res = {'route': f'probes_k2_s{sd}'} | |
| A = torch.linalg.solve(XA.T@XA + 10*torch.eye(XA.shape[1]), XA.T@YA.flatten(1)); ridge = lambda Z: (Z@A).view(-1, 4, R.D) | |
| mlp = train_mlp(XA, YA, seed=sd); ctrl = train_mlp(XA, YA, shuffle=True, seed=sd) | |
| YF = torch.tensor([[R.V[j].index(m[j]) for j in range(4)] for m in R.FIT]) | |
| torch.manual_seed(sd); cls = nn.Sequential(nn.Linear(XF.shape[1], 1024), nn.GELU(), nn.Linear(1024, sum(len(v) for v in R.V))); opt = torch.optim.Adam(cls.parameters(), 1e-3) | |
| off = [0, 12, 24, 40, 50]; muF, sdF = XF.mean(0), XF.std(0)+1e-6 | |
| for _ in range(2000): | |
| b = torch.randint(0, len(XF), (128,)); o = cls((XF[b]-muF)/sdF); loss = sum(F.cross_entropy(o[:, off[j]:off[j+1]], YF[b, j]) for j in range(4)); opt.zero_grad(); loss.backward(); opt.step() | |
| XFn = F.normalize(XF, dim=-1) | |
| for s in ('sealed_ood', 'sealed_4_new', 'sealed_1_new'): | |
| ms = S[s]; Z = slots(cod, ms) | |
| for name, fn in (('ridge_aug', ridge), ('mlp_aug', mlp), ('mlp_ctrl_shuffled', ctrl)): | |
| r = dec_emb(fn(Z), ms); PROBE[(sd, name, s)] = r; res[f'{name}:{s}'] = sum(r)/len(r) | |
| with torch.no_grad(): o = cls((Z-muF)/sdF) | |
| okc = []; okk1 = []; okk5 = []; sim = F.normalize(Z, dim=-1) @ XFn.T; top = sim.topk(5, dim=1).indices.cpu() | |
| for i, m in enumerate(ms): | |
| okc.append(float(all(m[j] in R.V[j] and int(o[i, off[j]:off[j+1]].argmax()) == R.V[j].index(m[j]) for j in range(4)))) | |
| n1 = R.FIT[int(top[i, 0])]; okk1.append(float(all(n1[j] == m[j] for j in range(4)))) | |
| vote = [max(set(R.FIT[int(t)][j] for t in top[i]), key=lambda v: sum(R.FIT[int(t)][j] == v for t in top[i])) for j in range(4)] | |
| okk5.append(float(all(vote[j] == m[j] for j in range(4)))) | |
| for name, r in (('mlp_cls_fit', okc), ('knn1', okk1), ('knn5', okk5)): PROBE[(sd, name, s)] = r; res[f'{name}:{s}'] = sum(r)/len(r) | |
| log(res) | |
| for s in ('sealed_ood', 'sealed_4_new', 'sealed_1_new'): | |
| for name in ('ridge_aug', 'mlp_aug', 'mlp_ctrl_shuffled', 'mlp_cls_fit', 'knn1', 'knn5'): | |
| a = [sum(PROBE[(sd, name, s)][i] for sd in SEEDS)/3 for i in range(len(S[s]))]; l = [sum(K2[(sd, s)][i] for sd in SEEDS)/3 for i in range(len(S[s]))] | |
| d, ci = boot(a, l, seed=3); log({'probe_mean3': name, 'set': s, 'acc': sum(a)/len(a), 'llm_k2_acc': sum(l)/len(l), 'diff_vs_llm': d, 'boot95': ci}) | |
| # ---- random-weight receiver ---- | |
| torch.manual_seed(0) | |
| rnd = AutoModelForCausalLM.from_config(AutoConfig.from_pretrained(R.MP)).float().eval().requires_grad_(False).to(R.DEV) | |
| rnd.model.embed_tokens.weight.data.copy_(R.Emb.data); R.lm = rnd | |
| for s in ('sealed_ood', 'sealed_4_new'): | |
| e = {'control': 'random_receiver', 'set': s, 'text_template': sum(per_ex(S[s], R.TPL))/len(S[s])} | |
| for sd in SEEDS: e[f'k2_s{sd}'] = sum(per_ex(S[s], cod=cods[sd]))/len(S[s]) | |
| log(e) | |
| SUMM.close(); PRED.close() | |