spb/modelmap Public License

1#!/usr/bin/env python32# =============================================================================3#  Project   : modelmap4#  File      : experiments/micro/expC_causal_verification/implementation/benchmark.py5#  Purpose   : Run #1 — layer-skip ablation test of the agreement probe map6#  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"""expC run #1 (hypothesis registered before this run).1516Does the agreement DIFFERENTIAL probe map (expA run #2) survive causal17testing? Skip the map's top-5 layers vs 20 random-5 draws vs bottom-5,18measure the drop in grammatical-agreement logit margin on held-out19minimal pairs.20"""2122from __future__ import annotations2324import json25import random26import subprocess27import sys28import time29from pathlib import Path3031import numpy as np3233ROOT = Path(__file__).resolve().parents[4]34sys.path.insert(0, str(ROOT / "src"))35sys.path.insert(0, str(ROOT / "benchmarks"))36from hardware_manifest import manifest3738from modelmap.capture.mlx_capture import install_taps3940MODEL = "mlx-community/Qwen3-0.6B-4bit"41MAP_JSON = ROOT / "atlas" / "qwen3-0.6b-4bit" / "probes" / "v2" / "map.json"42K = 543N_RANDOM_DRAWS = 2044N_PAIRS = 20045SEED = 7774647NOUN_PAIRS = [("key", "keys"), ("crate", "crates"), ("report", "reports"), ("valve", "valves"),48              ("ticket", "tickets"), ("ladder", "ladders"), ("sample", "samples"), ("cable", "cables"),49              ("permit", "permits"), ("beacon", "beacons"), ("filter", "filters"), ("stamp", "stamps")]50NEAR = ["near the entrance", "beside the counter", "under the shelf", "behind the gate",51        "next to the archive", "along the corridor", "opposite the office", "inside the vault",52        "across the yard", "above the workbench"]535455def held_out_pairs(rng: random.Random, forbidden: set[str]):56    """Minimal pairs (prefix, singular?) deduped against every v2 probe text."""57    pairs, seen = [], set()58    combos = [(n, loc) for n in NOUN_PAIRS for loc in NEAR]59    rng.shuffle(combos)60    for (sg, pl), loc in combos * 4:61        for noun, singular in ((sg, True), (pl, False)):62            prefix = f"The {noun} {loc}"63            probe_like = f"{prefix} is" if singular else f"{prefix} are"64            if prefix in seen or any(probe_like in f for f in forbidden):65                continue66            seen.add(prefix)67            pairs.append({"prefix": prefix, "singular": singular})68            if len(pairs) >= N_PAIRS:69                return pairs70    return pairs717273def main() -> int:74    import mlx.core as mx75    from mlx_lm import load7677    t0 = time.time()78    model, tokenizer = load(MODEL)79    taps = install_taps(model)80    n_layers = len(taps)8182    # top/bottom differential layers from the published map (mean pooling)83    mdoc = json.loads(MAP_JSON.read_text())84    agree = mdoc["properties"]["agreement"]85    diff = [(a["layer"], a["selectivity_mean"] - t["selectivity_mean"])86            for a, t in zip(agree["per_layer"]["A"], agree["twin_null_per_layer_A"])]87    ranked = [l for l, _ in sorted(diff, key=lambda x: -x[1])]88    top_k, bottom_k = ranked[:K], ranked[-K:]89    print(f"top-{K} differential layers: {sorted(top_k)} | bottom-{K}: {sorted(bottom_k)}")9091    # held-out minimal pairs92    forbidden = set()93    for f in (ROOT / "benchmarks" / "promptsets").glob("agreement_*.jsonl"):94        forbidden.update(json.loads(l)["text"] for l in f.read_text().splitlines())95    rng = random.Random(SEED)96    pairs = held_out_pairs(rng, forbidden)97    print(f"held-out minimal pairs: {len(pairs)}")9899    tok_is = tokenizer.encode(" is")100    tok_are = tokenizer.encode(" are")101    assert len(tok_is) == 1 and len(tok_are) == 1, "verb forms must be single tokens"102    id_is, id_are = tok_is[0], tok_are[0]103    prefix_ids = [tokenizer.encode(p["prefix"]) for p in pairs]104105    def margin(skip_layers: set[int]) -> float:106        for i, t in enumerate(taps):107            t.skip = i in skip_layers108        margins = []109        for ids, p in zip(prefix_ids, pairs):110            logits = model(mx.array([ids]))[0, -1, :]111            mx.eval(logits)112            m = float(logits[id_is] - logits[id_are])113            margins.append(m if p["singular"] else -m)114        for t in taps:115            t.skip = False116        return float(np.mean(margins))117118    baseline = margin(set())119    print(f"baseline margin: {baseline:+.4f}")120    top_m = margin(set(top_k))121    bottom_m = margin(set(bottom_k))122    candidates = [i for i in range(n_layers) if i not in set(top_k)]123    random_ms = []124    for d in range(N_RANDOM_DRAWS):125        draw = set(random.Random(SEED + 1 + d).sample(candidates, K))126        random_ms.append(margin(draw))127    dmg = lambda m: baseline - m128    rd = np.array([dmg(m) for m in random_ms])129    p95 = float(np.percentile(rd, 95))130    verdict = bool(dmg(top_m) >= p95 and dmg(top_m) >= 2 * rd.mean())131    print(f"damage: top-{K}={dmg(top_m):+.4f}  bottom-{K}={dmg(bottom_m):+.4f}  "132          f"random mean={rd.mean():+.4f} p95={p95:+.4f}  -> survives={verdict}")133134    commit = subprocess.run(["git", "rev-parse", "HEAD"], cwd=ROOT,135                            capture_output=True, text=True, check=False).stdout.strip()136    ts = time.strftime("%Y%m%dT%H%M%SZ", time.gmtime())137    outdir = ROOT / "results" / "expC_causal_verification" / ts138    outdir.mkdir(parents=True)139    doc = {140        "experiment": "expC_causal_verification", "run": 1,141        "scope": "layer-skip ablation of the agreement differential probe map",142        "commit": commit,143        "config": {"model": MODEL, "k": K, "n_random_draws": N_RANDOM_DRAWS,144                   "n_pairs": len(pairs), "seed": SEED,145                   "map_source": str(MAP_JSON.relative_to(ROOT)),146                   "top_layers": sorted(top_k), "bottom_layers": sorted(bottom_k)},147        "manifest": manifest(),148        "results": {149            "baseline_margin": baseline,150            "top_k_margin": top_m, "top_k_damage": dmg(top_m),151            "bottom_k_margin": bottom_m, "bottom_k_damage": dmg(bottom_m),152            "random_margins": random_ms,153            "random_damage_mean": float(rd.mean()),154            "random_damage_p95": p95,155            "survives_causal_test": verdict,156        },157        "wall_seconds": round(time.time() - t0, 1),158    }159    (outdir / "results.json").write_text(json.dumps(doc, indent=2) + "\n")160    print(f"results -> {outdir / 'results.json'}  ({doc['wall_seconds']} s)")161    return 0162163164if __name__ == "__main__":165    sys.exit(main())166