spb/modelmap Public License

1#!/usr/bin/env python32# =============================================================================3#  Project   : modelmap4#  File      : experiments/micro/expC_causal_verification/implementation/benchmark_v5.py5#  Purpose   : Run #5 — Level-3 path: dose-response steering of the direction6#  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 #5 (hypothesis registered before this run).1516Activation-addition steering of the diff-of-means agreement direction at17three early-band layers, doses alpha in {-2,-1,0,+1,+2} x sigma_l; strict18monotonicity + halving criterion + 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"40TEST_LAYERS = (4, 8, 12)41DOSES = (-2.0, -1.0, 0.0, 1.0, 2.0)42N_RANDOM_DIRS = 343SEED = 7814445NOUN_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")]48FRESH_NEAR = ["past the boiler room", "near the loading dock", "behind the ticket booth",49              "under the mezzanine", "beside the flagpole", "opposite the greenhouse",50              "inside the stairwell", "along the towpath"]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 FRESH_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    # direction + per-layer projection std from agreement_A (the v1 object)81    items = [json.loads(l) for l in82             (ROOT / "benchmarks" / "promptsets" / "agreement_A.jsonl").read_text().splitlines()]83    toks = [tokenizer.encode(it["text"]) for it in items]84    labels = np.array([it["label"] for it in items])85    print("capture agreement_A (direction + sigma)…", flush=True)86    reps = capture_pooled(model, taps, toks)["mean"]87    dirs, sigmas = {}, {}88    for layer in TEST_LAYERS:89        x = reps[:, layer, :]90        u = x[labels == "correct"].mean(0) - x[labels == "violated"].mean(0)91        u = u / (np.linalg.norm(u) + 1e-8)92        dirs[layer] = u93        sigmas[layer] = float((x @ u).std())9495    def add_fn(u_np, delta):96        d = mx.array((u_np * delta).astype(np.float32))97        def fn(out):98            return (out.astype(mx.float32) + d).astype(out.dtype)99        return fn100101    base_m = margin()102    print(f"baseline margin {base_m:+.4f} | pairs {len(pairs)}", flush=True)103104    results = {"per_layer": {}, "baseline_margin": base_m}105    mono_ok, halve_ok, spec_ok = True, True, True106    for layer in TEST_LAYERS:107        dose_margins = {}108        for a in DOSES:109            if a == 0.0:110                dose_margins[a] = base_m111                continue112            taps[layer].edit = add_fn(dirs[layer], a * sigmas[layer])113            dose_margins[a] = margin()114            taps[layer].edit = None115        seq = [dose_margins[a] for a in DOSES]116        mono = all(seq[i] < seq[i + 1] for i in range(len(seq) - 1))117        halve = dose_margins[-2.0] <= 0.5 * base_m118        rand_changes = []119        for s in range(N_RANDOM_DIRS):120            ru = np.random.default_rng(3000 * layer + s).standard_normal(reps.shape[-1])121            ru /= np.linalg.norm(ru)122            for a in (-2.0, 2.0):123                taps[layer].edit = add_fn(ru.astype(np.float32), a * sigmas[layer])124                rand_changes.append(abs(margin() - base_m))125                taps[layer].edit = None126        spec = float(np.mean(rand_changes)) < 0.25 * base_m127        mono_ok &= mono; halve_ok &= halve; spec_ok &= spec128        results["per_layer"][layer] = {129            "sigma": sigmas[layer],130            "dose_margins": {str(a): dose_margins[a] for a in DOSES},131            "monotone": mono, "halved_at_minus2": halve,132            "random_dir_mean_abs_change": float(np.mean(rand_changes)),133            "specific": spec,134        }135        print(f"L{layer:02d} doses " +136              " ".join(f"{a:+.0f}σ:{dose_margins[a]:+.2f}" for a in DOSES) +137              f" | mono={mono} halve={halve} randΔ={np.mean(rand_changes):.2f} spec={spec}",138              flush=True)139140    passes = bool(mono_ok and halve_ok and spec_ok)141    print(f"LEVEL-3 CRITERION: monotone={mono_ok} halved={halve_ok} specific={spec_ok} "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": 5,151        "scope": "dose-response steering (Level-3 path) at early-band layers",152        "commit": commit,153        "config": {"model": MODEL, "test_layers": list(TEST_LAYERS), "doses": list(DOSES),154                   "n_random_dirs": N_RANDOM_DIRS, "n_pairs": len(pairs), "seed": SEED},155        "manifest": manifest(),156        "results": results,157        "criterion": {"monotone_all": mono_ok, "halved_all": halve_ok,158                      "specific_all": spec_ok, "passes": passes},159        "wall_seconds": round(time.time() - t0, 1),160    }, indent=2) + "\n")161    print(f"results -> {outdir / 'results.json'}")162    return 0163164165if __name__ == "__main__":166    sys.exit(main())167