spb/modelmap Public License

1#!/usr/bin/env python32# =============================================================================3#  Project   : modelmap4#  File      : experiments/micro/expC_causal_verification/implementation/benchmark_v6.py5#  Purpose   : Run #6 — minimal L3 claim: layer-12 handle, four fresh sources6#  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 #6 (hypothesis registered before this run).1516Single-layer (L12) dose-response handle claim, four fresh direction sources17(new-permutation disjoint halves of A and B), third fresh behavioral bank,18shared random-direction specificity control.19"""2021from __future__ import annotations2223import json24import random25import subprocess26import sys27import time28from pathlib import Path2930import numpy as np3132ROOT = Path(__file__).resolve().parents[4]33sys.path.insert(0, str(ROOT / "src"))34sys.path.insert(0, str(ROOT / "benchmarks"))35from hardware_manifest import manifest3637from modelmap.capture.mlx_capture import capture_pooled, install_taps3839MODEL = "mlx-community/Qwen3-0.6B-4bit"40LAYER = 1241DOSES = (-2.0, -1.0, 0.0, 1.0, 2.0)42N_RANDOM_DIRS = 343SEED = 7824445NOUN_PAIRS = [("key", "keys"), ("crate", "crates"), ("report", "reports"), ("valve", "valves"),46              ("ticket", "tickets"), ("ladder", "ladders"), ("sample", "samples"), ("cable", "cables"),47              ("permit", "permits"), ("beacon", "beacons"), ("filter", "filters"), ("stamp", "stamps")]48THIRD_NEAR = ["beside the turbine hall", "near the cold frame", "behind the switchboard",49              "under the gantry", "next to the signal box", "opposite the pump house",50              "inside the drying shed", "along the breakwater"]515253def main() -> int:54    import mlx.core as mx55    from mlx_lm import load5657    t0 = time.time()58    model, tokenizer = load(MODEL)59    taps = install_taps(model)6061    rng = random.Random(SEED)62    combos = [(n, loc) for n in NOUN_PAIRS for loc in THIRD_NEAR]63    rng.shuffle(combos)64    pairs = []65    for (sg, pl), loc in combos:66        pairs.append({"prefix": f"The {sg} {loc}", "singular": True})67        pairs.append({"prefix": f"The {pl} {loc}", "singular": False})68    prefix_ids = [tokenizer.encode(p["prefix"]) for p in pairs]69    id_is, id_are = tokenizer.encode(" is")[0], tokenizer.encode(" are")[0]7071    def margin() -> float:72        out = []73        for ids, p in zip(prefix_ids, pairs):74            logits = model(mx.array([ids]))[0, -1, :]75            mx.eval(logits)76            m = float(logits[id_is] - logits[id_are])77            out.append(m if p["singular"] else -m)78        return float(np.mean(out))7980    def add_fn(u_np, delta):81        d = mx.array((u_np * delta).astype(np.float32))82        def fn(out):83            return (out.astype(mx.float32) + d).astype(out.dtype)84        return fn8586    # four fresh direction sources at L12 (new permutation seed)87    sources = {}88    d_model = None89    for s in ("A", "B"):90        items = [json.loads(l) for l in91                 (ROOT / "benchmarks" / "promptsets" / f"agreement_{s}.jsonl").read_text().splitlines()]92        toks = [tokenizer.encode(it["text"]) for it in items]93        labels = np.array([it["label"] for it in items])94        print(f"capture agreement_{s}…", flush=True)95        reps = capture_pooled(model, taps, toks)["mean"][:, LAYER, :]96        d_model = reps.shape[-1]97        n = len(labels)98        perm = np.random.default_rng(SEED + ord(s)).permutation(n)99        for hi, idx in enumerate((perm[: n // 2], perm[n // 2:])):100            x, l = reps[idx], labels[idx]101            u = x[l == "correct"].mean(0) - x[l == "violated"].mean(0)102            u = u / (np.linalg.norm(u) + 1e-8)103            sources[f"{s}half{hi + 1}"] = {"u": u, "sigma": float((x @ u).std())}104105    base_m = margin()106    print(f"baseline margin {base_m:+.4f} | pairs {len(pairs)}", flush=True)107108    per_source = {}109    all_ab = True110    for name, sv in sources.items():111        dm = {}112        for a in DOSES:113            if a == 0.0:114                dm[a] = base_m115                continue116            taps[LAYER].edit = add_fn(sv["u"], a * sv["sigma"])117            dm[a] = margin()118            taps[LAYER].edit = None119        seq = [dm[a] for a in DOSES]120        mono = all(seq[i] < seq[i + 1] for i in range(len(seq) - 1))121        halve = dm[-2.0] <= 0.5 * base_m122        all_ab &= (mono and halve)123        per_source[name] = {"sigma": sv["sigma"],124                            "dose_margins": {str(a): dm[a] for a in DOSES},125                            "monotone": mono, "halved": halve}126        print(f"{name:8s} " + " ".join(f"{a:+.0f}σ:{dm[a]:+.2f}" for a in DOSES) +127              f" | mono={mono} halve={halve}", flush=True)128129    sigma_ref = float(np.mean([sv["sigma"] for sv in sources.values()]))130    rand_changes = []131    for s in range(N_RANDOM_DIRS):132        ru = np.random.default_rng(4000 + s).standard_normal(d_model)133        ru /= np.linalg.norm(ru)134        for a in (-2.0, 2.0):135            taps[LAYER].edit = add_fn(ru.astype(np.float32), a * sigma_ref)136            rand_changes.append(abs(margin() - base_m))137            taps[LAYER].edit = None138    spec = float(np.mean(rand_changes)) < 0.25 * base_m139    passes = bool(all_ab and spec)140    print(f"L12 HANDLE CLAIM: all-sources mono+halve={all_ab} "141          f"specific={spec} (randΔ {np.mean(rand_changes):.2f} vs bound {0.25 * base_m:.2f}) "142          f"-> passes={passes}")143144    commit = subprocess.run(["git", "rev-parse", "HEAD"], cwd=ROOT,145                            capture_output=True, text=True, check=False).stdout.strip()146    ts = time.strftime("%Y%m%dT%H%M%SZ", time.gmtime())147    outdir = ROOT / "results" / "expC_causal_verification" / ts148    outdir.mkdir(parents=True)149    (outdir / "results.json").write_text(json.dumps({150        "experiment": "expC_causal_verification", "run": 6,151        "scope": "minimal L3 claim: layer-12 handle, four fresh sources, third bank",152        "commit": commit,153        "config": {"model": MODEL, "layer": LAYER, "doses": list(DOSES),154                   "sources": list(sources), "n_random_dirs": N_RANDOM_DIRS,155                   "n_pairs": len(pairs), "seed": SEED},156        "manifest": manifest(),157        "baseline_margin": base_m,158        "per_source": per_source,159        "specificity": {"random_dir_mean_abs_change": float(np.mean(rand_changes)),160                        "bound": 0.25 * base_m, "specific": spec},161        "criterion": {"all_sources_mono_halve": all_ab, "specific": spec, "passes": passes},162        "wall_seconds": round(time.time() - t0, 1),163    }, indent=2) + "\n")164    print(f"results -> {outdir / 'results.json'}")165    return 0166167168if __name__ == "__main__":169    sys.exit(main())170