File size: 5,182 Bytes
86415de | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | # 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'}})
@torch.no_grad()
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()
|