spb/wpai0001_uqo

1#!/usr/bin/env python32"""Persistent-graft check: load a SAVED adapter graft in a fresh process and evaluate the recipient with it3on freshly generated held-out splits (new seed). This is the only script that may claim 'saved modification';4the adapter was fitted in scripts/run_transplant.py and written to experiments/summaries/<run>/persistent/.5Usage: python scripts/eval_persistent.py --run <graft run id> [--seed-base 70000]6"""7from __future__ import annotations89import argparse10import json11import sys12from pathlib import Path1314import numpy as np1516ROOT = Path(__file__).resolve().parents[1]17sys.path.insert(0, str(ROOT / "src"))18sys.path.insert(0, str(ROOT / "scripts"))1920from train_models import ROLES, model_dir  # noqa: E40221from wp_ai_0001.alignment.maps import AffineMap  # noqa: E40222from wp_ai_0001.datasets import opseq  # noqa: E40223from wp_ai_0001.datasets.common import Split  # noqa: E40224from wp_ai_0001.instrumentation.provenance import utc_now  # noqa: E40225from wp_ai_0001.models.toy_transformer import resolve_device  # noqa: E40226from wp_ai_0001.statistics.bootstrap import cluster_bootstrap_mean  # noqa: E40227from wp_ai_0001.training.train_toy import load as load_model  # noqa: E40228from wp_ai_0001.transplantation.experiment import build_eval_splits, evaluate_condition, omega_step_acc  # noqa: E40229from wp_ai_0001.transplantation.grafts import AdapterGraft, GraftSpec  # noqa: E402303132def main() -> None:33    ap = argparse.ArgumentParser()34    ap.add_argument("--run", required=True)35    ap.add_argument("--seed-base", type=int, default=70_000)36    args = ap.parse_args()37    pdir = ROOT / "experiments" / "summaries" / args.run / "persistent"38    z = np.load(pdir / "adapter_graft.npz")39    hook, rkey = str(z["hook"]), str(z["recipient"])40    role, seed = rkey.rsplit("_s", 1)41    device = resolve_device("auto")42    model, vocab = load_model(model_dir(role, int(seed)), device)43    task = ROLES[role][0]44    ex = Split.load(model_dir(role, int(seed)) / "train_split.json").item_hashes()45    splits = build_eval_splits(ROLES["donor"][0], ex, args.seed_base, 1000, 100)46    names = ["target_iid", "control_iid"]47    base = evaluate_condition(model, None, None, splits, vocab, device, names)48    ad = AffineMap("adapter_ridge", z["W"], z["b"])49    graft = evaluate_condition(model, hook, AdapterGraft(GraftSpec("adapter", hook, None, 1.0, "all"), ad), splits, vocab, device, names)50    rep = {"utc": utc_now(), "run": args.run, "recipient": rkey, "hook": hook, "fresh_process": True, "splits": {k: v.hash for k, v in splits.items() if k in names}}51    for n in names:52        rep[n] = {"base": base[n].accuracy, "with_saved_adapter": graft[n].accuracy,53                  "ci95_with_saved_adapter": cluster_bootstrap_mean(graft[n].correct.astype(float), graft[n].clusters, 2000).to_dict(),54                  "omega_step_base": float(np.nanmean(omega_step_acc(splits[n], base[n]))) if n == "target_iid" else None,55                  "omega_step_with_saved_adapter": float(np.nanmean(omega_step_acc(splits[n], graft[n]))) if n == "target_iid" else None}56    (pdir / "fresh_process_eval.json").write_text(json.dumps(rep, indent=2))57    print(json.dumps(rep, indent=2))585960if __name__ == "__main__":61    main()62