spb/modelmap Public License

1#!/usr/bin/env python32# =============================================================================3#  Project   : modelmap4#  File      : experiments/micro/expA_probe_reliability/implementation/benchmark_v2.py5#  Purpose   : expA run #2 — structure-borne, token-balanced properties6#  Author    : Simon-Pierre Boucher7#  Contact   : contact@spboucher.ai8#  Website   : https://modelmap.io9#  Created   : 2026-08-1210#  Modified  : 2026-08-1211#  Platform  : macOS / Apple Silicon (arm64) — MLX / Metal12#  License   : All rights reserved (research code)13# =============================================================================14"""expA run #2 (hypothesis block registered in hypothesis.md before this run).1516v2 promptsets (word_order / agreement / arith_valid, token-balanced classes),17BOTH mean-pooled and last-token representations, random-init twin on A sets,18plus the v1 lang_id_A positive control (harness must still hit ceiling there).19"""2021from __future__ import annotations2223import json24import subprocess25import sys26import time27from pathlib import Path2829import numpy as np3031ROOT = Path(__file__).resolve().parents[4]32sys.path.insert(0, str(ROOT / "src"))33sys.path.insert(0, str(ROOT / "benchmarks"))34from hardware_manifest import manifest3536from modelmap.capture.mlx_capture import capture_pooled, install_taps, random_init_twin37from modelmap.probes.linear import probe_with_control38from modelmap.stats.replication import bh_fdr, bootstrap_ci, replication_rate3940MODEL = "mlx-community/Qwen3-0.6B-4bit"41PROPERTIES = ("word_order", "agreement", "arith_valid")42SETS = ("A", "B")43POOLINGS = ("mean", "last")44SEEDS = (0, 1, 2, 3, 4)45TOP_K = 546FDR_Q = 0.0547PROMPTS = ROOT / "benchmarks" / "promptsets"484950def load_set(name: str):51    items = [json.loads(l) for l in (PROMPTS / f"{name}.jsonl").read_text().splitlines()]52    return [it["text"] for it in items], np.array([it["label"] for it in items])535455def probe_grid(reps, labels):56    out = []57    for layer in range(reps.shape[1]):58        for seed in SEEDS:59            r = probe_with_control(reps[:, layer, :], labels, seed=seed)60            out.append({"layer": layer, "seed": seed, "task_acc": r.task_accuracy,61                        "control_acc": r.control_accuracy, "selectivity": r.selectivity})62    return out636465def summarize(cells, n_layers):66    rng = np.random.default_rng(0)67    layers, pvals = [], []68    for layer in range(n_layers):69        sel = np.array([c["selectivity"] for c in cells if c["layer"] == layer])70        acc = np.array([c["task_acc"] for c in cells if c["layer"] == layer])71        point, lo, hi = bootstrap_ci(sel, rng=rng)72        boots = rng.choice(sel, size=(10_000, sel.size)).mean(axis=1)73        p = float(max((boots <= 0).mean(), 1e-4))74        pvals.append(p)75        layers.append({"layer": layer, "task_acc_mean": float(acc.mean()),76                       "task_acc_seed_sd": float(acc.std(ddof=1)),77                       "selectivity_mean": point, "selectivity_ci": [lo, hi], "p_boot": p})78    disc = bh_fdr(np.array(pvals), q=FDR_Q)79    for row, d in zip(layers, disc):80        row["fdr_significant"] = bool(d)81    tops = []82    for seed in SEEDS:83        by = [(c["layer"], c["task_acc"]) for c in cells if c["seed"] == seed]84        tops.append({l for l, _ in sorted(by, key=lambda t: -t[1])[:TOP_K]})85    rp, rl, rh = replication_rate(tops)86    return {"layers": layers, "replication_rate_topk": {"k": TOP_K, "point": rp, "ci": [rl, rh]},87            "n_fdr_significant": int(disc.sum())}888990def main() -> int:91    import mlx.core as mx92    from mlx_lm import load93    from mlx_lm.utils import hf_repo_to_path9495    t0 = time.time()96    mx.random.seed(0)97    model, tokenizer = load(MODEL)98    taps = install_taps(model)99    twin = random_init_twin(hf_repo_to_path(MODEL))100    twin_taps = install_taps(twin)101102    results = {"real": {}, "twin": {}, "positive_control": {}}103104    # ---- positive control: v1 lang_id_A must still hit ceiling (mean pooling)105    texts, labels = load_set("lang_id_A")106    toks = [tokenizer.encode(t) for t in texts]107    reps = capture_pooled(model, taps, toks)108    pc = summarize(probe_grid(reps["mean"], labels), reps["mean"].shape[1])109    results["positive_control"]["lang_id_A_mean"] = {110        "max_task_acc": float(max(r["task_acc_mean"] for r in pc["layers"]))}111    print(f"positive control lang_id_A: maxAcc={results['positive_control']['lang_id_A_mean']['max_task_acc']:.3f}",112          flush=True)113114    grids = {"real": {}, "twin": {}}115    for prop in PROPERTIES:116        for s in SETS:117            name = f"{prop}_{s}"118            texts, labels = load_set(name)119            toks = [tokenizer.encode(t) for t in texts]120            print(f"capture real  {name}…", flush=True)121            reps = capture_pooled(model, taps, toks)122            for pool in POOLINGS:123                grids["real"][f"{name}_{pool}"] = (probe_grid(reps[pool], labels),124                                                   reps[pool].shape[1])125            if s == "A":126                print(f"capture twin  {name}…", flush=True)127                reps_t = capture_pooled(twin, twin_taps, toks)128                for pool in POOLINGS:129                    grids["twin"][f"{name}_{pool}"] = (probe_grid(reps_t[pool], labels),130                                                       reps_t[pool].shape[1])131132    for kind in ("real", "twin"):133        for name, (cells, n_layers) in grids[kind].items():134            results[kind][name] = summarize(cells, n_layers)135            results[kind][name]["cells"] = cells136137    # ---- headline138    summary = {}139    for prop in PROPERTIES:140        for pool in POOLINGS:141            a = results["real"][f"{prop}_A_{pool}"]["layers"]142            b = results["real"][f"{prop}_B_{pool}"]["layers"]143            tw = results["twin"][f"{prop}_A_{pool}"]["layers"]144            seed_sd = float(np.mean([r["task_acc_seed_sd"] for r in a + b]))145            shifts = [abs(x["task_acc_mean"] - y["task_acc_mean"]) for x, y in zip(a, b)]146            diff = [x["selectivity_mean"] - t["selectivity_mean"] for x, t in zip(a, tw)]147            n_signal = int(sum(1 for x, t, d in148                               zip(a, tw, diff) if d > 0.10 and x["fdr_significant"]))149            summary[f"{prop}_{pool}"] = {150                "max_task_acc_A": float(max(r["task_acc_mean"] for r in a)),151                "max_task_acc_B": float(max(r["task_acc_mean"] for r in b)),152                "twin_max_selectivity": float(max(r["selectivity_mean"] for r in tw)),153                "twin_max_acc": float(max(r["task_acc_mean"] for r in tw)),154                "mean_seed_sd": seed_sd,155                "mean_dataset_shift": float(np.mean(shifts)),156                "layers_shift_gt_seed_sd": int(sum(s > max(seed_sd, 1e-9) for s in shifts)),157                "layers_real_minus_twin_gt_0.10": n_signal,158                "max_real_minus_twin_sel": float(max(diff)),159                "replication_topk_A": results["real"][f"{prop}_A_{pool}"]["replication_rate_topk"]["point"],160            }161            s = summary[f"{prop}_{pool}"]162            print(f"{prop:12s} {pool:4s} accA={s['max_task_acc_A']:.3f} twinAcc={s['twin_max_acc']:.3f} "163                  f"twinSel={s['twin_max_selectivity']:.3f} signalLayers={s['layers_real_minus_twin_gt_0.10']} "164                  f"maxΔsel={s['max_real_minus_twin_sel']:+.3f} shift>{'sd'}={s['layers_shift_gt_seed_sd']}/28",165                  flush=True)166167    commit = subprocess.run(["git", "rev-parse", "HEAD"], cwd=ROOT,168                            capture_output=True, text=True, check=False).stdout.strip()169    ts = time.strftime("%Y%m%dT%H%M%SZ", time.gmtime())170    outdir = ROOT / "results" / "expA_probe_reliability" / ts171    outdir.mkdir(parents=True)172    doc = {"experiment": "expA_probe_reliability", "run": 2,173           "scope": "v2 structure-borne token-balanced properties, mean+last pooling, twin null",174           "commit": commit,175           "config": {"model": MODEL, "seeds": list(SEEDS), "top_k": TOP_K, "fdr_q": FDR_Q,176                      "poolings": list(POOLINGS),177                      "promptsets": json.loads((PROMPTS / "manifest_v2.json").read_text())},178           "manifest": manifest(), "summary": summary, "results": results,179           "wall_seconds": round(time.time() - t0, 1)}180    (outdir / "results.json").write_text(json.dumps(doc, indent=2) + "\n")181    print(f"results -> {outdir / 'results.json'}  ({doc['wall_seconds']} s)")182    return 0183184185if __name__ == "__main__":186    sys.exit(main())187