mirror of
https://github.com/wassname/jsteer.git
synced 2026-08-21 04:50:17 +08:00
wassnameandClaudypoo
57c8d4b166
clamp: y += (C - <y,v_hat>)v_hat at all positions -- bounded perturbation regardless of generation length, vs add's per-step accumulation via KV cache. C=0 is directional ablation. Smoke (Qwen3-0.6B, happy/joy): clamp C=+20 stays coherent and on-concept (drifts to 'happiness and joy of my childhood', in Chinese) while add C=+8 already degenerates to 'joyjoyjoy...'. Co-Authored-By: Claudypoo <288921227+claudypoo@users.noreply.github.com>
100 lines
4.7 KiB
Python
100 lines
4.7 KiB
Python
"""Smoke test: fit a tiny Jacobian on Qwen3-0.6B, steer on the word "happy/joy".
|
|
|
|
(authored by Claude)
|
|
|
|
Touches the whole jsteer path end-to-end on a real model, cheaply:
|
|
fit -> save -> load round-trip -> word_vector -> generate at C in {-8, 0, 8}.
|
|
|
|
Run:
|
|
uv run python scripts/smoke.py 2>&1 | tee /tmp/claude-1000/jsteer_smoke.log
|
|
|
|
Read the prints: the FULL first fit prompt (with special tokens), the FULL
|
|
generation prompt as the model sees it, and each generation verbatim, with a
|
|
SHOULD line so a deviation is legible.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
from loguru import logger
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent.parent)) # repo root for config
|
|
import config # noqa: E402
|
|
from config import DEVICE, DTYPE # noqa: E402
|
|
from jsteer import Jacobian # noqa: E402
|
|
|
|
MODEL = "Qwen/Qwen3-0.6B"
|
|
CACHE = str(config.ART / "qwen3-0.6b-smoke.jac")
|
|
|
|
# Text to fit the smoke Jacobian on. Kept >17 tokens each because jlens drops
|
|
# the first 16 positions, so shorter prompts leave nothing to fit.
|
|
PROMPTS = [
|
|
"The weather this morning was cold and grey, so I made a large pot of coffee and sat by the window watching the rain fall.",
|
|
"Scientists have long argued about whether the early universe expanded smoothly or in sudden bursts that left traces we can still measure today.",
|
|
"My grandmother used to tell stories about growing up on a small farm, where every season brought a different kind of hard and honest work.",
|
|
"The city council voted last night to repair the old stone bridge downtown, a project residents have been requesting for well over a decade.",
|
|
"After months of training, the runners lined up at dawn, breath fogging in the cold air, waiting nervously for the starting gun to fire.",
|
|
"Learning to cook well takes patience more than talent, a willingness to taste often, to fail a few times, and to pay attention to detail.",
|
|
"The library on the corner smells of old paper and dust, and its quiet reading room has been my favourite place to think for many years.",
|
|
"Software written in a hurry tends to accumulate small mistakes that hide quietly until, one ordinary afternoon, they surface all at once together.",
|
|
]
|
|
|
|
GEN_PROMPT = "I went to the park today and"
|
|
|
|
|
|
def _show_tokens(tok, text: str, label: str) -> None:
|
|
ids = tok(text, add_special_tokens=True).input_ids
|
|
logger.info(f"{label}: {len(ids)} tokens (add_special_tokens=True)")
|
|
logger.info(f"{label} decoded-with-special:\n{tok.decode(ids)!r}")
|
|
|
|
|
|
def main() -> None:
|
|
logger.info(f"loading {MODEL} ({DTYPE}) on {DEVICE}")
|
|
tok = AutoTokenizer.from_pretrained(MODEL)
|
|
model = AutoModelForCausalLM.from_pretrained(MODEL, dtype=DTYPE).to(DEVICE).eval()
|
|
|
|
# SHOULD: fit prompt shows a BOS-like sink token then english web text.
|
|
# ELSE tokenizer/template drift (jlens force_bos should add an attention sink).
|
|
_show_tokens(tok, PROMPTS[0], "FIT PROMPT[0]")
|
|
|
|
logger.info("fitting Jacobian (layers=0.3..0.9 band, dim_batch=8)")
|
|
jac = Jacobian.fit(model, tok, PROMPTS, layers=(0.3, 0.9),
|
|
dim_batch=8, max_seq_len=128)
|
|
logger.info(f"fitted: {jac!r} layers={jac.layers}")
|
|
|
|
jac.save(CACHE)
|
|
jac2 = Jacobian.load(CACHE)
|
|
logger.info(f"load round-trip: {jac2!r} layers={jac2.layers}")
|
|
# SHOULD: reloaded layers identical to fitted. ELSE save/load wiring bug.
|
|
assert jac2.layers == jac.layers, (jac2.layers, jac.layers)
|
|
|
|
v = jac2.word_vector(model, tok, ["happy", "joy"])
|
|
logger.info(f"word_vector layers={sorted(v.stacked)}")
|
|
|
|
# SHOULD: generation prompt shows a sink token then "I went to the park today and".
|
|
_show_tokens(tok, GEN_PROMPT, "GEN PROMPT")
|
|
enc = tok(GEN_PROMPT, return_tensors="pt").to(DEVICE)
|
|
|
|
for mode in ("add", "clamp"):
|
|
v.cfg.apply_mode = mode
|
|
for C in (-8, 0, 8, 20):
|
|
with v(model, C=C):
|
|
out = model.generate(**enc, max_new_tokens=40, do_sample=False,
|
|
pad_token_id=tok.eos_token_id)
|
|
text = tok.decode(out[0][enc.input_ids.shape[1]:], skip_special_tokens=True)
|
|
logger.info(f"=== {mode} C={C:+d} generation ===\n{text!r}")
|
|
|
|
logger.info(
|
|
"SHOULD: C=+8 mentions happiness/joy more than C=0; C=-8 less or "
|
|
"negative tone. ELSE steering wiring or sign issue. add C=+20 MAY "
|
|
"degenerate (unbounded accumulation); clamp C=+20 SHOULD stay more "
|
|
"coherent (component pinned, perturbation bounded). clamp C=0 is "
|
|
"directional ablation, expect near-baseline text. Gibberish at small "
|
|
"|C| means the vector is malformed.")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|