#!/usr/bin/env python3 """ Slice 20: fine-tune cointegrated/rubert-tiny (11.9M, 3-layer BERT) end-to-end for the binary executable-intent boundary (action vs not_action) on the frozen v2 dev corpus, following the slice 20 brief. Rules honoured: - full end-to-end fine-tuning, CLS-pooled native classification head - tokenizer used unchanged (audit in slice20_audit.py) - max length from corpus stats (p99+margin, cap 128): 25 here - narrow search: LR in {1e-5, 2e-5, 5e-5}, <= 6 epochs, early stop on a development (within-fold) split, best checkpoint restored - >= 3 seeds (42/17/7) for every config - grouped 5-fold CV reuse; cap-Q leave-generator-out as primary stress case - two input regimes: A = natural text (orig), B = punctuation-stripped (strip) Artifacts under /tmp/mvn-s20/: pre/{regime}_ids.npy, _attn.npy tokenized corpus (all three views) oof/{regime}_{lr}_{seed}_probs.npz OOF probs per view (var_orig/nofinal/strip) oof/{regime}_{lr}_{seed}_metrics.json lfo/{regime}_{lr}_{seed}.json capability-Q LOFO (held-out family) results/summary.json models/{regime}_{lr}_{seed}_fold{i}.pt, lfo_{seed}.pt CLI: slice20_pretrained.py {pre, grouped, lfo, metrics, runtime, onnx} """ import argparse import importlib.util import json import os import sys import time import numpy as np import torch from torch import nn from transformers import AutoConfig, AutoTokenizer from transformers import BertForSequenceClassification sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) import slice19_main as s19 RESULTS_DIR = "/tmp/mvn-s20" MODEL_DIR = ("/tmp/mvn-s20/hf-cache/models--cointegrated--rubert-tiny/" "snapshots/5441c5ea8026d4f6d7505ec004845409f1259fb1") MODEL_NAME = "cointegrated/rubert-tiny" MODEL_SHA = "5441c5ea8026d4f6d7505ec004845409f1259fb1" REGIMES = ["A", "B"] LRS = [1e-5, 2e-5, 5e-5] SEEDS = [42, 17, 7] VIEWS = ["orig", "nofinal", "strip"] MAX_LEN = 25 BATCH = 32 MAX_EPOCHS = 4 EARLY_STOP = 1 # patience in epochs on val PR-AUC VAL_FRACTION = 0.12 WEIGHT_DECAY = 0.01 torch.set_num_threads(4) # ─── tokenizer / input preparation ────────────────────────────────────────── def _load_tokenizer(): return AutoTokenizer.from_pretrained(MODEL_DIR) def tokenize(texts, tok): e = tok(list(texts), add_special_tokens=True, padding="max_length", truncation=True, max_length=MAX_LEN) return np.array(e["input_ids"], np.int64), np.array(e["attention_mask"], np.int64) def cmd_pre(args): os.makedirs(os.path.join(RESULTS_DIR, "pre"), exist_ok=True) tok = _load_tokenizer() rows = s19.load_dev() # distributed over all variants, all rows, both regimes for regime in REGIMES: train_view = "orig" if regime == "A" else "strip" tsrc = [r[f"text_{train_view}"] for r in rows] ids, attn = tokenize(tsrc, tok) np.save(os.path.join(RESULTS_DIR, "pre", f"{regime}_ids.npy"), ids) np.save(os.path.join(RESULTS_DIR, "pre", f"{regime}_attn.npy"), attn) # eval stress views tokenized under the same regime's vocab/format for v in VIEWS: ids_v, attn_v = tokenize([r[f"text_{v}"] for r in rows], tok) np.save(os.path.join(RESULTS_DIR, "pre", f"ev_{regime}_{v}_ids.npy"), ids_v) np.save(os.path.join(RESULTS_DIR, "pre", f"ev_{regime}_{v}_attn.npy"), attn_v) print(f"regime {regime} done, train view={train_view}") def _make_model(): cfg = AutoConfig.from_pretrained(MODEL_DIR) cfg.num_labels = 1 # sine logit, BCEWithLogits — matches slice 19 head model = BertForSequenceClassification.from_pretrained(MODEL_DIR, config=cfg) return model # ─── training ─────────────────────────────────────────────────────────────── def _val_split(rows_idx, y, seed): rng = np.random.RandomState(seed) idx = rng.permutation(rows_idx) # stratified: keep the y-ratio in both parts pos = idx[y[idx] == 1] neg = idx[y[idx] == 0] nv_pos = max(1, int(round(len(pos) * VAL_FRACTION))) nv_neg = max(1, int(round(len(neg) * VAL_FRACTION))) v = np.concatenate([pos[:nv_pos], neg[:nv_neg]]) t = np.concatenate([pos[nv_pos:], neg[nv_neg:]]) return t, v def train_model(id_arr, attn, y, tr_idx, val_idx, lr, seed, builder=None): """Fine-tune the full encoder; early-stop on val PR-AUC; return best state.""" net = (builder or _make_model)() opt = torch.optim.AdamW([p for p in net.parameters()], lr=lr, weight_decay=WEIGHT_DECAY) lossf = nn.BCEWithLogitsLoss() from sklearn.metrics import average_precision_score tr = torch.from_numpy(np.ascontiguousarray(id_arr[tr_idx])) ta = torch.from_numpy(np.ascontiguousarray(attn[tr_idx])) ty = torch.from_numpy(y[tr_idx].astype(np.float32)) va = torch.from_numpy(np.ascontiguousarray(id_arr[val_idx])) vaa = torch.from_numpy(np.ascontiguousarray(attn[val_idx])) vy = y[val_idx] best_pr = -1.0 best_state = None best_epoch = 0 patience = 0 n = len(tr_idx) rng = np.random.RandomState(seed * 97 % 2**31) for epoch in range(MAX_EPOCHS): net.train() perm = rng.permutation(n) running = 0.0 nb = 0 for st in range(0, n, BATCH): bidx = torch.from_numpy(perm[st:st + BATCH]) logits = net(input_ids=tr[bidx], attention_mask=ta[bidx]).logits.squeeze(-1) loss = lossf(logits, ty[bidx]) opt.zero_grad() loss.backward() opt.step() running += float(loss) nb += 1 net.eval() with torch.no_grad(): pval = torch.sigmoid(net(input_ids=va, attention_mask=vaa).logits.squeeze(-1)).numpy() if len(np.unique(vy)) > 1: pr = average_precision_score(vy, pval) else: pr = 0.0 if pr > best_pr: best_pr = pr best_state = {k: v.detach().clone() for k, v in net.state_dict().items()} best_epoch = epoch + 1 patience = 0 else: patience += 1 if patience >= EARLY_STOP: break net.load_state_dict(best_state) return net, best_epoch, best_pr, running / max(nb, 1) def predict_proba(net, id_arr, attn, idx=None): net.eval() idx = np.arange(len(id_arr)) if idx is None else idx out = [] with torch.no_grad(): for st in range(0, len(idx), BATCH * 4): bi = idx[st:st + BATCH * 4] iid = torch.from_numpy(np.ascontiguousarray(id_arr[bi])) att = torch.from_numpy(np.ascontiguousarray(attn[bi])) out.append(torch.sigmoid(net(input_ids=iid, attention_mask=att).logits.squeeze(-1)).numpy()) return np.concatenate(out) # ─── grouped CV ───────────────────────────────────────────────────────────── def cmd_grouped(args): os.makedirs(os.path.join(RESULTS_DIR, "oof"), exist_ok=True) os.makedirs(os.path.join(RESULTS_DIR, "models"), exist_ok=True) rows = s19.load_dev() y = np.array([r["y"] for r in rows]) folds = np.array([r["cv_fold"] for r in rows]) for regime in REGIMES: ids = np.load(os.path.join(RESULTS_DIR, "pre", f"{regime}_ids.npy")) attn = np.load(os.path.join(RESULTS_DIR, "pre", f"{regime}_attn.npy")) ev = {v: (np.load(os.path.join(RESULTS_DIR, "pre", f"ev_{regime}_{v}_ids.npy")), np.load(os.path.join(RESULTS_DIR, "pre", f"ev_{regime}_{v}_attn.npy"))) for v in VIEWS} for lr in LRS: for seed in SEEDS: probs = {v: np.zeros(len(rows)) for v in VIEWS} fold_records = [] for fold in range(5): tr = np.where(folds != fold)[0] te = np.where(folds == fold)[0] t_idx, v_idx = _val_split(tr, y, seed + 100 * fold) net, ep, best_pr, _ = train_model(ids, attn, y, t_idx, v_idx, lr, seed + fold) torch.save(net.state_dict(), os.path.join(RESULTS_DIR, "models", f"{regime}_{lr}_{seed}_fold{fold}.pt")) for v in VIEWS: probs[v][te] = predict_proba(net, *ev[v], te) fold_records.append({"fold": fold, "epochs": ep, "val_pr": best_pr}) np.savez(os.path.join(RESULTS_DIR, "oof", f"{regime}_{lr}_{seed}_probs.npz"), var_orig=probs["orig"], var_nofinal=probs["nofinal"], var_strip=probs["strip"]) with open(os.path.join(RESULTS_DIR, "oof", f"{regime}_{lr}_{seed}_metrics.json"), "w") as f: json.dump({"folds": fold_records}, f, indent=2) m = s19.binary_metrics(y, probs["strip"]) print(f"[{regime}] lr={lr:.0e} seed={seed} " f"PR={m['PR_AUC']:.3f} P={m['P']:.3f} R={m['R']:.3f} " f"FA={m['FA']} epochs={[fr['epochs'] for fr in fold_records]}", flush=True) print("grouped done") # ─── cap-Q leave-generator-out ────────────────────────────────────────────── def cmd_lfo(args): os.makedirs(os.path.join(RESULTS_DIR, "lfo"), exist_ok=True) rows = s19.load_dev() y = np.array([r["y"] for r in rows]) tags = [r["tags"] for r in rows] src = np.array(["capability_question" not in t for t in tags]) tgt = ~src for regime in REGIMES: ids = np.load(os.path.join(RESULTS_DIR, "pre", f"{regime}_ids.npy")) attn = np.load(os.path.join(RESULTS_DIR, "pre", f"{regime}_attn.npy")) for lr in LRS: for seed in SEEDS: s_idx = np.where(src)[0] t_idx, v_idx = _val_split(s_idx, y, seed + 7) net, ep, best_pr, _ = train_model(ids, attn, y, t_idx, v_idx, lr, seed) p = predict_proba(net, ids, attn, np.where(tgt)[0]) yt = y[tgt] out = { "regime": regime, "lr": lr, "seed": seed, "rows": int(tgt.sum()), "epochs": ep, "val_pr": best_pr, "mean_action_proba": float(np.mean(p)), "max_action_proba": float(np.max(p)), "acc": float(((p >= 0.5) == (yt == 1)).mean()), "FA": int(((p >= 0.5) & (yt == 0)).sum()), "FA_rate": float(((p >= 0.5) & (yt == 0)).mean()), } with open(os.path.join(RESULTS_DIR, "lfo", f"{regime}_{lr}_{seed}.json"), "w") as f: json.dump(out, f, indent=2) print(f"[{regime}] lr={lr:.0e} seed={seed} capQ LOFO " f"acc={out['acc']:.3f} FA_rate={out['FA_rate']:.3f} " f"mean_p={out['mean_action_proba']:.3f} epochs={ep}", flush=True) print("lfo done") # ─── metrics aggregation ──────────────────────────────────────────────────── def cmd_metrics(args): rows = s19.load_dev() y = np.array([r["y"] for r in rows]) tags = [r["tags"] for r in rows] pairs = s19.build_pairs(rows, [r["text_strip"] for r in rows]) summary = {} for regime in REGIMES: summary[regime] = {} for lr in LRS: per_seed = [] for seed in SEEDS: z = np.load(os.path.join(RESULTS_DIR, "oof", f"{regime}_{lr}_{seed}_probs.npz")) entry = {"seed": seed, "views": {v: {"FA_rate": s19.binary_metrics(y, z[f"var_{v}"])["FA_rate"], "PR": s19.binary_metrics(y, z[f"var_{v}"])["PR_AUC"], "capQ_FA": s19.capq_fa(tags, y, z[f"var_{v}"])} for v in VIEWS}, "strip": s19.binary_metrics(y, z["var_strip"]), "pairs": {v: s19.pair_metrics(pairs, z[f"var_{v}"]) for v in VIEWS}, "curve": s19.threshold_curve(y, z["var_strip"], np.arange(0.30, 1.0, 0.02)), "ops": s19.operating_points(s19.threshold_curve( y, z["var_strip"], np.arange(0.30, 1.0, 0.02)), z["var_strip"], pairs), } with open(os.path.join(RESULTS_DIR, "oof", f"{regime}_{lr}_{seed}_metrics.json")) as f: entry["folds"] = json.load(f)["folds"] per_seed.append(entry) # LOFO lfos = [] for seed in SEEDS: with open(os.path.join(RESULTS_DIR, "lfo", f"{regime}_{lr}_{seed}.json")) as f: lfos.append(json.load(f)) summary[regime][str(lr)] = {"seeds": per_seed, "lfo": lfos} # primary selection: min mean cap-Q LOFO FA_rate across seeds sel = {} for regime in REGIMES: best_lr, best_mean = None, 9e9 for lr in LRS: mean_fa = np.mean([l["FA_rate"] for l in summary[regime][str(lr)]["lfo"]]) sel[f"{regime}_{lr}"] = round(float(mean_fa), 4) if mean_fa < best_mean: best_lr, best_mean = lr, mean_fa sel[f"{regime}_selected"] = str(best_lr) summary["selection"] = sel os.makedirs(os.path.join(RESULTS_DIR, "results"), exist_ok=True) with open(os.path.join(RESULTS_DIR, "results", "summary.json"), "w") as f: json.dump(summary, f, indent=2, default=str) print(json.dumps(sel)) print("metrics done") def _tokenizer_bytes(): return sum(os.path.getsize(os.path.join(MODEL_DIR, f)) for f in ["vocab.txt", "tokenizer.json"] if os.path.exists(os.path.join(MODEL_DIR, f))) # ─── runtime + ONNX ───────────────────────────────────────────────────────── def cmd_runtime(args): import time net = _make_model() tok = _load_tokenizer() rows = s19.load_dev() x = [r["text_orig"] for r in rows][:200] e = tok(x, padding="max_length", truncation=True, max_length=MAX_LEN, return_tensors="pt") ids, attn = e["input_ids"], e["attention_mask"] net.eval() with torch.no_grad(): # warmup for _ in range(3): net(input_ids=ids[:1], attention_mask=attn[:1]) # batch-1 latency lat = [] for i in range(200): t0 = time.perf_counter() net(input_ids=ids[i:i + 1], attention_mask=attn[i:i + 1]) lat.append((time.perf_counter() - t0) * 1e6) # tokenization latency t0 = time.perf_counter() for i in range(200): tok(x[i]) tok_us = (time.perf_counter() - t0) / 200 * 1e6 n_params = sum(p.numel() for p in net.parameters()) fp32 = n_params * 4 rep = { "model": MODEL_NAME, "sha": MODEL_SHA, "params": n_params, "fp32_bytes": fp32, "fp16_bytes": fp32 // 2, "int8_bytes": n_params, "tokenizer_bytes": _tokenizer_bytes(), "latency_us_mean": float(np.mean(lat)), "latency_us_p50": float(np.median(lat)), "latency_us_p95": float(np.percentile(lat, 95)), "max_len": MAX_LEN, "tok_us": round(tok_us, 2), "num_threads": 12, } with open(os.path.join(RESULTS_DIR, "runtime.json"), "w") as f: json.dump(rep, f, indent=2) print(json.dumps(rep, indent=2)) print("runtime done") def cmd_onnx(args): net = _make_model() net.eval() tok = _load_tokenizer() rows = s19.load_dev() try: import torch.onnx dummy = { "input_ids": torch.zeros(1, MAX_LEN, dtype=torch.long), "attention_mask": torch.ones(1, MAX_LEN, dtype=torch.long), } with torch.no_grad(): torch.onnx.export(net, (dummy,), os.path.join(RESULTS_DIR, "rubert-tiny-gate.onnx"), input_names=["input_ids", "attention_mask"], output_names=["logits"], opset_version=14, dynamic_axes={"input_ids": {0: "batch"}, "attention_mask": {0: "batch"}}) # parity on a fixed sample import numpy as np samp = [(r["text_orig"], r["y"]) for r in rows[:200]] e = tok([s[0] for s in samp], padding="max_length", truncation=True, max_length=MAX_LEN, return_tensors="pt") with torch.no_grad(): pt = torch.sigmoid(net(**e).logits.squeeze(-1)).numpy() import onnxruntime as ort so = ort.SessionOptions() so.intra_op_num_threads = 12 sess = ort.InferenceSession(os.path.join(RESULTS_DIR, "rubert-tiny-gate.onnx"), sess_options=so, providers=["CPUExecutionProvider"]) on = sess.run(None, {"input_ids": e["input_ids"].numpy(), "attention_mask": e["attention_mask"].numpy()})[0] on = 1 / (1 + np.exp(-on).squeeze(-1)) mx = float(np.max(np.abs(pt - on))) size = os.path.getsize(os.path.join(RESULTS_DIR, "rubert-tiny-gate.onnx")) rep = {"max_logit_diff": mx, "onnx_bytes": size, "parity_n": len(samp), "provider": "CPUExecutionProvider"} with open(os.path.join(RESULTS_DIR, "onnx.json"), "w") as f: json.dump(rep, f, indent=2) print(json.dumps(rep)) except Exception as ex: print("onnx export/parity failed:", ex) with open(os.path.join(RESULTS_DIR, "onnx.json"), "w") as f: json.dump({"error": str(ex)}, f, indent=2) print("onnx done") # ─── optional capacity/pretraining ceiling (brief §13) ────────────────────── # Trigger: tiny1 clearly improved over from-scratch on in-pool but missed the # LOFO boundary. tiny2 is the same 3-layer 312-hidden BERT family; it tests # whether a *newer, larger-vocab* pretraining of the same family generalises # where tiny1 failed — disambiguating "this family is the wrong prior" from a # one-off pretraining. It cannot test capacity (same depth/size). MODEL2_DIR = ("/tmp/mvn-s20/hf-tiny2-cache/models--cointegrated--rubert-tiny2/" "snapshots/e8ed3b0c8bbf4fb6984c3de043bf7d2f4e5969ae") MODEL2_SHA = "e8ed3b0c8bbf4fb6984c3de043bf7d2f4e5969ae" CEIL_RESULTS = os.path.join(RESULTS_DIR, "tiny2") def cmd_ceiling(args): import torch os.makedirs(CEIL_RESULTS, exist_ok=True) os.makedirs(os.path.join(CEIL_RESULTS, "models"), exist_ok=True) tok = AutoTokenizer.from_pretrained(MODEL2_DIR) rows = s19.load_dev() y = np.array([r["y"] for r in rows]) tags = [r["tags"] for r in rows] folds = np.array([r["cv_fold"] for r in rows]) pairs = s19.build_pairs(rows, [r["text_strip"] for r in rows]) # audit: is the tiny2 tokenizer sane on the corpus before anything else n_unk = 0 n_tok = 0 lens = [] for r in rows: e = tok(r["text_orig"]) n_unk += e["input_ids"].count(tok.unk_token_id) n_tok += len(e["input_ids"]) lens.append(len(e["input_ids"])) audit = {"vocab_size": tok.vocab_size, "unk_count": int(n_unk), "unk_rate": round(n_unk / max(n_tok, 1), 5), "seq_len_p99": sorted(lens)[int(len(lens) * .99)], "seq_len_max": max(lens)} with open(os.path.join(CEIL_RESULTS, "audit.json"), "w") as f: json.dump(audit, f, indent=2) print("tiny2 audit:", audit) # tokenize the corpus (regime A only — natural text, the in-pool best) ids_a, attn_a = tokenize([r["text_orig"] for r in rows], tok) ev = {v: tokenize([r[f"text_{v}"] for r in rows], tok) for v in VIEWS} def make2(): cfg = AutoConfig.from_pretrained(MODEL2_DIR) cfg.num_labels = 1 m = BertForSequenceClassification.from_pretrained(MODEL2_DIR, config=cfg) return m # cap-Q LOFO, 3 seeds, matching the A@2e-5 tiny1 config lfors = [] src = np.where(np.array(["capability_question" not in t for t in tags]))[0] tgt = np.where(np.array(["capability_question" in t for t in tags]))[0] for seed in SEEDS: t_idx, v_idx = _val_split(src, y, seed + 7) net, ep, best_pr, _ = train_model(ids_a, attn_a, y, t_idx, v_idx, 2e-5, seed, builder=make2) p = predict_proba(net, ids_a, attn_a, tgt) yt = y[tgt] lfors.append({"regime": "A(tiny2)", "lr": 2e-5, "seed": seed, "rows": int(len(tgt)), "epochs": ep, "mean_action_proba": float(np.mean(p)), "acc": float(((p >= 0.5) == (yt == 1)).mean()), "FA": int(((p >= 0.5) & (yt == 0)).sum()), "FA_rate": float(((p >= 0.5) & (yt == 0)).mean())}) with open(os.path.join(CEIL_RESULTS, "lfo.json"), "w") as f: json.dump(lfors, f, indent=2) print("tiny2 LOFO:", [round(l["FA_rate"], 3) for l in lfors]) # grouped CV for the same best config + in-pool pairs / capQ oof = {v: np.zeros(len(rows)) for v in VIEWS} for fold in range(5): tr = np.where(folds != fold)[0] te = np.where(folds == fold)[0] t_idx, v_idx = _val_split(tr, y, 42 + 100 * fold) net, _, _, _ = train_model(ids_a, attn_a, y, t_idx, v_idx, 2e-5, 42 + fold, builder=make2) for v in VIEWS: oof[v][te] = predict_proba(net, *ev[v], te) inpool = {"strip": s19.binary_metrics(y, oof["strip"]), "pairs": {v: s19.pair_metrics(pairs, oof[v]) for v in VIEWS}, "capQ_inpool": {v: s19.capq_fa(tags, y, oof[v]) for v in VIEWS}} with open(os.path.join(CEIL_RESULTS, "grouped.json"), "w") as f: json.dump(inpool, f, indent=2, default=str) b = inpool["strip"] print(f"tiny2 grouped A@2e-5: PR={b['PR_AUC']:.3f} P={b['P']:.3f} R={b['R']:.3f} " f"FA={b['FA']} pairs_strip={inpool['pairs']['strip']['ordering_acc']:.3f} " f"capQ_strip={inpool['capQ_inpool']['strip']:.3f}", flush=True) print("ceiling done") def main(): ap = argparse.ArgumentParser() ap.add_argument("cmd", choices=["pre", "grouped", "lfo", "metrics", "runtime", "onnx", "ceiling"]) args = ap.parse_args() t0 = time.time() globals()[f"cmd_{args.cmd}"](args) print(f"elapsed {time.time()-t0:.1f}s") if __name__ == "__main__": main()