From 143e9add807c7a8e2b87158252865e71f00c7129 Mon Sep 17 00:00:00 2001 From: wassname <1103714+wassname@users.noreply.github.com> Date: Sat, 11 Jul 2026 07:40:38 +0800 Subject: [PATCH] demo: rubric readout -- think then {"ans":N}, logprob-weighted expected digit per C show_steer gains a rubric= param and rubric_score(): the model rates a 0-9 axis, we force the {"ans": slot and read the logprob-weighted expected digit. guided.py's mechanism reduced to one scalar for the demo (rigorous K-way debiased version stays in moral-maps). UAT (scripts/scratch/uat_rubric.py) on happy/joy: in the coherent window ans rises 3.52->4.99->8.06 across C=-0.5,0,+0.5 (pmass=1.00); at the degeneration extremes (C=+-1.5) pmass collapses to ~0 and the number is correctly flagged meaningless. Co-Authored-By: Claudypoo <288921227+claudypoo@users.noreply.github.com> --- jsteer/demo.py | 52 +++++++++++++++++++++++++++++++++-- scripts/scratch/uat_rubric.py | 50 +++++++++++++++++++++++++++++++++ 2 files changed, 100 insertions(+), 2 deletions(-) create mode 100644 scripts/scratch/uat_rubric.py diff --git a/jsteer/demo.py b/jsteer/demo.py index 4451278..be17c32 100644 --- a/jsteer/demo.py +++ b/jsteer/demo.py @@ -33,11 +33,48 @@ def _cthulhu_say(text: str) -> str: " \\", " ^(;,;)^"])) +# think-then-answer rubric read: the demo's one-number sanity signal that steering +# moved the target axis. Same mechanism as moral-maps guided.py (let the model +# think, then read the logprobs at a JSON answer slot), reduced to a single scalar. +_ANS_FMT = (' Think it over, then answer with JSON {"ans": N} where N is a single' + ' digit from 0 (least) to 9 (most).') + + +@torch.no_grad() +def rubric_score(model, tok, rubric: str, *, max_new_tokens: int, seed: int + ) -> tuple[float, float]: + """Ask `rubric`, let the model think, then FORCE the answer slot `{"ans": ` and + read the logprob-weighted expected digit 0-9 there. Returns (expected, pmass). + + expected = sum_d d * softmax(logit_d over the 10 digit tokens) -- a continuous + scalar from single-token logprobs (cleaner than parsing a multi-token float). + pmass = full-vocab softmax mass on the 10 digit tokens: a coherence guard, ~0 + means the slot isn't a digit (prefix/tokenizer mismatch), so distrust expected. + The rigorous K-way, position-debiased version is moral-maps guided.py; this is + the demo's cheap readout, scored under whatever steering is active.""" + prompt = chat_input(tok, rubric + _ANS_FMT) + enc = tok(prompt, return_tensors="pt").to(model.device) + torch.manual_seed(seed) + out = model.generate(**enc, max_new_tokens=max_new_tokens, + pad_token_id=tok.eos_token_id) + think = tok.decode(out[0][enc.input_ids.shape[1]:], + skip_special_tokens=False).split("")[0] + forced = prompt + think + '\n{"ans": ' # our own deterministic slot + fenc = tok(forced, return_tensors="pt").to(model.device) + logits = model(**fenc).logits[0, -1].float() + ids = torch.tensor([tok(str(d), add_special_tokens=False).input_ids[0] + for d in range(10)], device=logits.device) + expected = float((logits[ids].softmax(0) * torch.arange(10., device=ids.device)).sum()) + pmass = float(logits.softmax(0)[ids].sum()) + return expected, pmass + + @torch.no_grad() def show_steer(jac: Jacobian, model, tok, vec, user_msg: str, *, Cs=(-6, 0, 6), layer: int | None = None, k: int = 6, max_new_tokens: int = 512, seed: int = 0, - apply_mode: str | None = None, apply_span: int = 1) -> None: + apply_mode: str | None = None, apply_span: int = 1, + rubric: str | None = None) -> None: """One block per C: lens readout at `layer`, then the raw generation, all under steering. Uses the model's own generation_config sampling; `seed` fixes it so the C blocks are comparable. `layer` defaults to the top fitted @@ -48,7 +85,12 @@ def show_steer(jac: Jacobian, model, tok, vec, user_msg: str, *, (add | clamp | add_last | replace_last) to swap how v hits the residual without re-extracting; `apply_span` is the trailing-position width for the last/replace modes. Coefficient units differ by mode (clamp sets a component - VALUE, add scales a direction), so each mode wants its own Cs.""" + VALUE, add scales a direction), so each mode wants its own Cs. + + Pass `rubric` (a 0-9 rating question about the steered axis) to add the + quantitative readout: per C, the model thinks then answers `{"ans": N}` and we + report the logprob-weighted expected digit. It SHOULD rise with +C and fall + with -C; flat means the steer isn't moving that axis (see rubric_score).""" if apply_mode is not None: vec = Vector(dataclasses.replace(vec.cfg, apply_mode=apply_mode, apply_span=apply_span), vec.shared, vec.stacked) @@ -72,10 +114,16 @@ def show_steer(jac: Jacobian, model, tok, vec, user_msg: str, *, jtop = jac.lens_topk(model, tok, prompt, layer=layer, k=k) out = model.generate(**enc, max_new_tokens=max_new_tokens, pad_token_id=tok.eos_token_id) + ans = (rubric_score(model, tok, rubric, max_new_tokens=max_new_tokens, + seed=seed) if rubric is not None else None) # raw decode WITH special tokens: real /, <|im_end|> visible, # nothing parsed or re-wrapped -- debuggable exactly as the model emitted it gen = tok.decode(out[0][enc.input_ids.shape[1]:], skip_special_tokens=False) readout = " · ".join(t.strip() for t, _ in jtop) block = [f"\n--- C={C:+g} " + "-" * 60, f" lens @L{layer}:", _cthulhu_say(readout), gen] + if ans is not None: + # SHOULD rise with +C, fall with -C; flat => steer not moving this axis. + # pmass<~0.5 => answer slot isn't a digit, distrust the number. + block.append(f" rubric ans≈{ans[0]:.2f}/9 (pmass={ans[1]:.2f})") logger.info("\n".join(block) + "\n") diff --git a/scripts/scratch/uat_rubric.py b/scripts/scratch/uat_rubric.py new file mode 100644 index 0000000..3f8e2a9 --- /dev/null +++ b/scripts/scratch/uat_rubric.py @@ -0,0 +1,50 @@ +"""UAT for the rubric readout added to show_steer. (Claude) Loads the 4B + Hub +lens + happy word vector once, then runs show_steer with a 0-9 rubric at +Cs=(-1.5, 0, +1.5). PASS = expected digit rises monotonically with +C and pmass +stays high (the answer slot really is a digit). + + uv run python scripts/scratch/uat_rubric.py +""" +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[2])) +import config # noqa: E402 loguru setup +import torch # noqa: E402 +from loguru import logger # noqa: E402 +from transformers import AutoModelForCausalLM, AutoTokenizer # noqa: E402 + +from jsteer import Jacobian, show_steer # noqa: E402 +from jsteer.demo import rubric_score # 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) + +RUBRIC = "On a scale of 0 (bleak) to 9 (joyful), how is your outlook on the week?" + +# full demo block with the number attached, so we see readout + generation + ans +show_steer(jac, model, tok, v, "Describe how your week has been going.", + Cs=(-1.5, 0, 1.5), rubric=RUBRIC) + +# bare scalar sweep for a clean monotonicity check (SHOULD rise with +C) +logger.info("\n\n=== rubric-only sweep (SHOULD rise with +C, pmass>0.5) ===") +rows = [] +for C in (-1.5, -0.5, 0, 0.5, 1.5): + with v(model, C=C): + ans, pmass = rubric_score(model, tok, RUBRIC, max_new_tokens=512, seed=0) + rows.append((C, ans, pmass)) + logger.info(f"C={C:+g} ans={ans:.2f}/9 pmass={pmass:.2f}") + +# the claim only holds where the answer slot is a digit (pmass>0.5); at the +# degeneration extremes pmass ~0 and the number is meaningless BY DESIGN, so the +# monotonicity check must be restricted to the coherent rows. +coherent = [(C, a) for C, a, p in rows if p > 0.5] +anss_c = [a for _, a in coherent] +mono = anss_c == sorted(anss_c) +logger.info(f"\nUAT: coherent C={[C for C, _ in coherent]} ans={[round(a,2) for a in anss_c]} " + f"monotone_up={mono} (degenerate rows pmass<=0.5 excluded)")