Files
jsteer/scripts/scratch/smoke_sweep.py
T
wassnameandClaudypoo 160d675c61 demo: JSON-object coherence probe replaces rigged forced-digit pmass
The old rubric pmass was softmax mass on digit tokens at the hand-fed `{"ans": `
slot, so it was ~always 1 (the prefix forces a digit even from a fried model) --
a blind coherence guard. Replace with the users design: free-generate
{"ans":N,"why":str,"2+2":M}, gate coherence on valid-JSON AND 2+2==4. GPU UATs:
object stays valid at |C|<=0.5, breaks at |C|>=0.75 (task 22/23); the old slot
never caught this.

span_pmass (mean top-1 prob over the span) is kept as a within-coherent confidence
read but NOT the coherence signal: a steer-fried model collapses into a confident
degenerate loop, so span_pmass climbs back to ~0.97 while the object is broken
(C=3.0). plot_sweep therefore colors by valid_frac, which cannot be fooled by
confident garbage.

Also: show_steers per-C cowsay now speaks steer-PROMOTED tokens (top of
steered-baseline logits) instead of lens_topk at the last chat-prompt position,
which only ever surfaced think-openers (Okay/Here/The) for every C. UAT: joy steer
-> < joy . happy . Happy . happy > at C=0.3 (task 24). compute_slice stays the
calibrated cross-layer lens readout.

Co-Authored-By: Claudypoo <288921227+claudypoo@users.noreply.github.com>
2026-07-11 21:06:46 +08:00

48 lines
2.3 KiB
Python

"""Smoke: coherence_sweep with the JSON-object coherence probe. (Claude)
max_steps=3 walks C out until the model can no longer emit a valid {"ans",...,"2+2"}
object. Two UATs below. max_new_tokens=384 keeps it fast.
uv run python scripts/scratch/smoke_sweep.py
"""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
import config # noqa: E402
import torch # noqa: E402
from loguru import logger # noqa: E402
from tabulate import tabulate # noqa: E402
from transformers import AutoModelForCausalLM, AutoTokenizer # noqa: E402
from jsteer import Jacobian # noqa: E402
from jsteer.demo import coherence_sweep # noqa: E402
MODEL = "Qwen/Qwen3.5-4B"
tok = AutoTokenizer.from_pretrained(MODEL)
model = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.bfloat16).to("cuda").eval()
jac = Jacobian.from_pretrained(config.LENS_REPO, filename=config.hub_lens_file(MODEL),
revision=config.LENS_REVISION)
band = jac.steer_band(model)
v = jac.word_vector(model, tok, ["happy", "joy"], layers=band)
rows = coherence_sweep(model, tok, v,
"On a scale of 0 (bleak) to 9 (joyful), how is your outlook on the week?",
step=0.1, max_steps=3, n_samples=3, max_new_tokens=384)
logger.info("\n" + tabulate(rows, headers="keys", tablefmt="github", floatfmt="+.2f"))
# UAT 1: sampling+BMA active -> the 3 seeds diverge, so >=1 row has ans_std>0.
nonzero = [r for r in rows if r["ans_std"] > 0]
logger.info(f"\nUAT1: rows with ans_std>0 = {len(nonzero)}/{len(rows)} "
f"(SHOULD be >0 -> sampling+BMA active)")
# UAT 2: the JSON coherence probe discriminates. At C=0 the model emits a valid object
# (valid_frac=1, high span_pmass); walking |C| out, span_pmass falls and the sweep stops
# at an incoherent boundary. If C=0 is already incoherent OR span_pmass never falls, the
# probe isn't measuring coherence -> broken.
c0 = next(r for r in rows if r["C"] == 0.0)
span0 = c0["span_pmass"]
edge = [r for r in rows if not r["coherent"]]
logger.info(f"\nUAT2: C=0 valid_frac={c0['valid_frac']:+.2f} span_pmass={span0:+.2f} "
f"(SHOULD valid_frac=1, span high); incoherent boundary rows={len(edge)} "
f"at C={[r['C'] for r in edge]} (SHOULD be >=1 -> sweep found a real edge)")