MIT LICENSE, README license fix, make_dataset.py, 019 re-evaluation, technical report and CITATION.cff
86415de verified | # CUERDAS-019: re-evaluation of the saved 018 K=2 checkpoints (no training). | |
| # 1) per-example predictions (text vs K=2) on every test set, 2) paired bootstrap CIs | |
| # on the differences, 3) random-weight receiver control on K=2. | |
| import os, sys, json, random, platform | |
| import torch, transformers | |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) | |
| import run018 as R # loads frozen receiver, data, splits and test sets exactly as in 018 | |
| from transformers import AutoModelForCausalLM, AutoConfig | |
| INP = os.environ.get('INP', '.') | |
| OD = os.environ.get('OD', '.') | |
| SEEDS = [int(x) for x in os.environ.get('SEEDS', '17,29,43').split(',')] | |
| NBOOT = int(os.environ.get('NBOOT', 10000)) | |
| PRED = open(f'{OD}/predictions019.jsonl', 'w') | |
| SUMM = open(f'{OD}/results019.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__, | |
| 'cuda': torch.version.cuda, 'device': torch.cuda.get_device_name() if R.DEV == 'cuda' else 'cpu'}}) | |
| def per_example(route, msgs, tpl=R.SUF_MAIN, cod=None, bs=24): | |
| """Same scoring as R.acc, but returns per-message correctness and predicted first tokens.""" | |
| out = [] | |
| for b in range(0, len(msgs), bs): | |
| mb = msgs[b:b+bs] | |
| lg, mt = (R.text_logits(mb, tpl) if route == 'text' else R.codec_logits(cod, mb, tpl)) | |
| pred = [[None]*4 for _ in mb]; good = [True]*len(mb) | |
| for r, (i, j) in enumerate(mt): | |
| k = int(lg[r][R.CAND[j]].argmax()) | |
| pred[i][j] = R.tok.decode([int(R.CAND[j][k])]).strip() | |
| good[i] &= (k == R.gold_idx(j, mb[i][j])) | |
| out += list(zip(good, pred)) | |
| return out | |
| SUB = R.FRESH['id'][:48] + R.FRESH['ood'][:48] | |
| SETS = {'fresh_id': (R.FRESH['id'], R.SUF_MAIN), 'fresh_ood': (R.FRESH['ood'], R.SUF_MAIN)} | |
| for k, v in R.TRANSFER.items(): SETS['transfer_' + k] = (v, R.SUF_MAIN) | |
| for k, t in R.SUF_HELD.items(): SETS['query_' + k] = (SUB, t) | |
| def boot(a, b, n=NBOOT, seed=0): | |
| """Paired bootstrap over test items of mean(a)-mean(b); returns point, 95% percentile CI.""" | |
| 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), (n, 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) | |
| # text baseline (receiver frozen -> identical for every seed) | |
| TEXT = {} | |
| for name, (msgs, tpl) in SETS.items(): | |
| TEXT[name] = per_example('text', msgs, tpl) | |
| for i, (ok, p) in enumerate(TEXT[name]): | |
| PRED.write(json.dumps({'route': 'text', 'set': name, 'i': i, 'msg': msgs[i], 'pred': p, 'correct': bool(ok)}) + '\n') | |
| log({'route': 'text', 'set': name, 'n': len(msgs), 'acc': sum(o for o, _ in TEXT[name]) / len(msgs)}) | |
| CODEC = {}; cods = {} | |
| 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(); cods[sd] = cod | |
| for name, (msgs, tpl) in SETS.items(): | |
| res = per_example('codec', msgs, tpl, cod); CODEC[(sd, name)] = res | |
| for i, (ok, p) in enumerate(res): | |
| PRED.write(json.dumps({'route': f'k2_s{sd}', 'set': name, 'i': i, 'msg': msgs[i], 'pred': p, 'correct': bool(ok)}) + '\n') | |
| a = [float(o) for o, _ in res]; t = [float(o) for o, _ in TEXT[name]] | |
| d, lo, hi = boot(a, t, seed=sd) | |
| log({'route': f'k2_s{sd}', 'set': name, 'n': len(msgs), 'acc': sum(a) / len(a), 'diff_vs_text_pp': d, 'boot95': [lo, hi]}) | |
| log({'route': f'k2_s{sd}', 'noLLM_fresh_ood200': R.nolm(cod, R.FRESH['ood'][:200]), 'noLLM_transfer_4new': R.nolm(cod, R.TRANSFER['4_new']), **R.probe(cod)}) | |
| # seed-averaged difference (per item: mean over seeds of codec correctness minus text), paired bootstrap over items | |
| for name in SETS: | |
| a = [sum(float(CODEC[(sd, name)][i][0]) for sd in SEEDS) / len(SEEDS) for i in range(len(TEXT[name]))] | |
| t = [float(o) for o, _ in TEXT[name]] | |
| d, lo, hi = boot(a, t, seed=1) | |
| log({'route': 'k2_mean3seeds', 'set': name, 'diff_vs_text_pp': d, 'boot95': [lo, hi]}) | |
| # random-weight receiver control: same frozen lexicon (input embeddings, tied LM head), random transformer weights | |
| 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) | |
| log({'random_receiver_tied_head': bool(rnd.lm_head.weight.data_ptr() == rnd.model.embed_tokens.weight.data_ptr())}) | |
| R.lm = rnd # R.text_logits / R.codec_logits read the module-level receiver | |
| for name in ('fresh_ood', 'transfer_4_new'): | |
| msgs, tpl = SETS[name] | |
| r = {'control': 'random_receiver', 'set': name, 'n': len(msgs), | |
| 'text': sum(o for o, _ in per_example('text', msgs, tpl)) / len(msgs)} | |
| for sd in SEEDS: r[f'k2_s{sd}'] = sum(o for o, _ in per_example('codec', msgs, tpl, cods[sd])) / len(msgs) | |
| log(r) | |
| PRED.close(); SUMM.close() | |