#!/usr/bin/env python3 # ============================================================================= # Project : modelmap # File : experiments/micro/expA_probe_reliability/implementation/benchmark_v2.py # Purpose : expA run #2 — structure-borne, token-balanced properties # Author : Simon-Pierre Boucher # Contact : contact@spboucher.ai # Website : https://modelmap.io # Created : 2026-08-12 # Modified : 2026-08-12 # Platform : macOS / Apple Silicon (arm64) — MLX / Metal # License : All rights reserved (research code) # ============================================================================= """expA run #2 (hypothesis block registered in hypothesis.md before this run). v2 promptsets (word_order / agreement / arith_valid, token-balanced classes), BOTH mean-pooled and last-token representations, random-init twin on A sets, plus the v1 lang_id_A positive control (harness must still hit ceiling there). """ from __future__ import annotations import json import subprocess import sys import time from pathlib import Path import numpy as np ROOT = Path(__file__).resolve().parents[4] sys.path.insert(0, str(ROOT / "src")) sys.path.insert(0, str(ROOT / "benchmarks")) from hardware_manifest import manifest from modelmap.capture.mlx_capture import capture_pooled, install_taps, random_init_twin from modelmap.probes.linear import probe_with_control from modelmap.stats.replication import bh_fdr, bootstrap_ci, replication_rate MODEL = "mlx-community/Qwen3-0.6B-4bit" PROPERTIES = ("word_order", "agreement", "arith_valid") SETS = ("A", "B") POOLINGS = ("mean", "last") SEEDS = (0, 1, 2, 3, 4) TOP_K = 5 FDR_Q = 0.05 PROMPTS = ROOT / "benchmarks" / "promptsets" def load_set(name: str): items = [json.loads(l) for l in (PROMPTS / f"{name}.jsonl").read_text().splitlines()] return [it["text"] for it in items], np.array([it["label"] for it in items]) def probe_grid(reps, labels): out = [] for layer in range(reps.shape[1]): for seed in SEEDS: r = probe_with_control(reps[:, layer, :], labels, seed=seed) out.append({"layer": layer, "seed": seed, "task_acc": r.task_accuracy, "control_acc": r.control_accuracy, "selectivity": r.selectivity}) return out def summarize(cells, n_layers): rng = np.random.default_rng(0) layers, pvals = [], [] for layer in range(n_layers): sel = np.array([c["selectivity"] for c in cells if c["layer"] == layer]) acc = np.array([c["task_acc"] for c in cells if c["layer"] == layer]) point, lo, hi = bootstrap_ci(sel, rng=rng) boots = rng.choice(sel, size=(10_000, sel.size)).mean(axis=1) p = float(max((boots <= 0).mean(), 1e-4)) pvals.append(p) layers.append({"layer": layer, "task_acc_mean": float(acc.mean()), "task_acc_seed_sd": float(acc.std(ddof=1)), "selectivity_mean": point, "selectivity_ci": [lo, hi], "p_boot": p}) disc = bh_fdr(np.array(pvals), q=FDR_Q) for row, d in zip(layers, disc): row["fdr_significant"] = bool(d) tops = [] for seed in SEEDS: by = [(c["layer"], c["task_acc"]) for c in cells if c["seed"] == seed] tops.append({l for l, _ in sorted(by, key=lambda t: -t[1])[:TOP_K]}) rp, rl, rh = replication_rate(tops) return {"layers": layers, "replication_rate_topk": {"k": TOP_K, "point": rp, "ci": [rl, rh]}, "n_fdr_significant": int(disc.sum())} def main() -> int: import mlx.core as mx from mlx_lm import load from mlx_lm.utils import hf_repo_to_path t0 = time.time() mx.random.seed(0) model, tokenizer = load(MODEL) taps = install_taps(model) twin = random_init_twin(hf_repo_to_path(MODEL)) twin_taps = install_taps(twin) results = {"real": {}, "twin": {}, "positive_control": {}} # ---- positive control: v1 lang_id_A must still hit ceiling (mean pooling) texts, labels = load_set("lang_id_A") toks = [tokenizer.encode(t) for t in texts] reps = capture_pooled(model, taps, toks) pc = summarize(probe_grid(reps["mean"], labels), reps["mean"].shape[1]) results["positive_control"]["lang_id_A_mean"] = { "max_task_acc": float(max(r["task_acc_mean"] for r in pc["layers"]))} print(f"positive control lang_id_A: maxAcc={results['positive_control']['lang_id_A_mean']['max_task_acc']:.3f}", flush=True) grids = {"real": {}, "twin": {}} for prop in PROPERTIES: for s in SETS: name = f"{prop}_{s}" texts, labels = load_set(name) toks = [tokenizer.encode(t) for t in texts] print(f"capture real {name}…", flush=True) reps = capture_pooled(model, taps, toks) for pool in POOLINGS: grids["real"][f"{name}_{pool}"] = (probe_grid(reps[pool], labels), reps[pool].shape[1]) if s == "A": print(f"capture twin {name}…", flush=True) reps_t = capture_pooled(twin, twin_taps, toks) for pool in POOLINGS: grids["twin"][f"{name}_{pool}"] = (probe_grid(reps_t[pool], labels), reps_t[pool].shape[1]) for kind in ("real", "twin"): for name, (cells, n_layers) in grids[kind].items(): results[kind][name] = summarize(cells, n_layers) results[kind][name]["cells"] = cells # ---- headline summary = {} for prop in PROPERTIES: for pool in POOLINGS: a = results["real"][f"{prop}_A_{pool}"]["layers"] b = results["real"][f"{prop}_B_{pool}"]["layers"] tw = results["twin"][f"{prop}_A_{pool}"]["layers"] seed_sd = float(np.mean([r["task_acc_seed_sd"] for r in a + b])) shifts = [abs(x["task_acc_mean"] - y["task_acc_mean"]) for x, y in zip(a, b)] diff = [x["selectivity_mean"] - t["selectivity_mean"] for x, t in zip(a, tw)] n_signal = int(sum(1 for x, t, d in zip(a, tw, diff) if d > 0.10 and x["fdr_significant"])) summary[f"{prop}_{pool}"] = { "max_task_acc_A": float(max(r["task_acc_mean"] for r in a)), "max_task_acc_B": float(max(r["task_acc_mean"] for r in b)), "twin_max_selectivity": float(max(r["selectivity_mean"] for r in tw)), "twin_max_acc": float(max(r["task_acc_mean"] for r in tw)), "mean_seed_sd": seed_sd, "mean_dataset_shift": float(np.mean(shifts)), "layers_shift_gt_seed_sd": int(sum(s > max(seed_sd, 1e-9) for s in shifts)), "layers_real_minus_twin_gt_0.10": n_signal, "max_real_minus_twin_sel": float(max(diff)), "replication_topk_A": results["real"][f"{prop}_A_{pool}"]["replication_rate_topk"]["point"], } s = summary[f"{prop}_{pool}"] print(f"{prop:12s} {pool:4s} accA={s['max_task_acc_A']:.3f} twinAcc={s['twin_max_acc']:.3f} " f"twinSel={s['twin_max_selectivity']:.3f} signalLayers={s['layers_real_minus_twin_gt_0.10']} " f"maxΔsel={s['max_real_minus_twin_sel']:+.3f} shift>{'sd'}={s['layers_shift_gt_seed_sd']}/28", flush=True) commit = subprocess.run(["git", "rev-parse", "HEAD"], cwd=ROOT, capture_output=True, text=True, check=False).stdout.strip() ts = time.strftime("%Y%m%dT%H%M%SZ", time.gmtime()) outdir = ROOT / "results" / "expA_probe_reliability" / ts outdir.mkdir(parents=True) doc = {"experiment": "expA_probe_reliability", "run": 2, "scope": "v2 structure-borne token-balanced properties, mean+last pooling, twin null", "commit": commit, "config": {"model": MODEL, "seeds": list(SEEDS), "top_k": TOP_K, "fdr_q": FDR_Q, "poolings": list(POOLINGS), "promptsets": json.loads((PROMPTS / "manifest_v2.json").read_text())}, "manifest": manifest(), "summary": summary, "results": results, "wall_seconds": round(time.time() - t0, 1)} (outdir / "results.json").write_text(json.dumps(doc, indent=2) + "\n") print(f"results -> {outdir / 'results.json'} ({doc['wall_seconds']} s)") return 0 if __name__ == "__main__": sys.exit(main())