diff --git a/README.md b/README.md index 9ba37bf..beb25b5 100644 --- a/README.md +++ b/README.md @@ -54,8 +54,11 @@ band = jac.steer_band(model) # steer the mid-depth 0.3-0.9 v = jac.word_vector(model, tok, ["happy", "joy"], layers=band) # generate through the chat template with thinking on; print, per strength C, -# the j-space readout + the trace + the answer. -show_steer(jac, model, tok, v, "Describe how your week has been going.", Cs=(0, 0.5, 1.5)) +# the j-space readout + the trace + the answer + a quantitative rubric +# digit (the model rates its 0-9 outlook; the logprob-weighted expected digit +# rises with +C until the text degenerates, where its pmass coherence guard drops). +show_steer(jac, model, tok, v, "Describe how your week has been going.", Cs=(0, 0.3, 1.5), + rubric="On a scale of 0 (bleak) to 9 (joyful), how is your outlook on the week?") ``` For a model the authors do not publish, fit your own (expensive, resumable): @@ -65,9 +68,11 @@ uv run python scripts/fit.py --model ``` The coefficient is lens-dependent, so sweep it. The pre-fitted lens gives a clean, -concentrated direction, so its knee is steep: C~0.5 moves the tone while the text -and reasoning stay fluent, and by C~1 it degenerates into token spam. -`nbs/word_steering.ipynb` shows the full sweep with the j-space and `` views. +concentrated direction, so its knee is steep: on this prompt C~0.3 moves the tone +strongly while the text and reasoning stay fluent (the rubric digit rises ~5 -> ~7 of +9), and by C~0.5 it degenerates into token spam (where the rubric's pmass collapses). +`nbs/word_steering.ipynb` shows the full sweep with the j-space, ``, and rubric +views. ## Persona j-thoughts (experimental) diff --git a/jsteer/demo.py b/jsteer/demo.py index be17c32..c48bfa5 100644 --- a/jsteer/demo.py +++ b/jsteer/demo.py @@ -41,8 +41,8 @@ _ANS_FMT = (' Think it over, then answer with JSON {"ans": N} where N is a singl @torch.no_grad() -def rubric_score(model, tok, rubric: str, *, max_new_tokens: int, seed: int - ) -> tuple[float, float]: +def rubric_score(model, tok, rubric: str, *, max_new_tokens: int, seed: int, + do_sample: bool = False, temperature: float = 0.7) -> 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). @@ -55,8 +55,14 @@ def rubric_score(model, tok, rubric: str, *, max_new_tokens: int, seed: int 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) + # this model ships no generation_config, so generate() is greedy by default: seeds + # only matter (distinct think traces) when do_sample=True -- which is what makes the + # coherence_sweep's multi-seed BMA average over anything. + gen_kw = dict(max_new_tokens=max_new_tokens, pad_token_id=tok.eos_token_id, + do_sample=do_sample) + if do_sample: + gen_kw["temperature"] = temperature + out = model.generate(**enc, **gen_kw) 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 @@ -69,6 +75,110 @@ def rubric_score(model, tok, rubric: str, *, max_new_tokens: int, seed: int return expected, pmass +@torch.no_grad() +def coherence_sweep(model, tok, vec, rubric: str, *, step: float = 0.1, + pmass_floor: float = 0.9, max_steps: int = 15, n_samples: int = 3, + temperature: float = 0.7, max_new_tokens: int = 512) -> list[dict]: + """Walk C outward from 0 in +/- directions, scoring the rubric each step, and STOP + a direction the step AFTER the answer slot goes incoherent (pmass 1, temperature=temperature) + for s in range(n_samples)] + anss = torch.tensor([a for a, _ in pairs]) + pmass = float(torch.tensor([p for _, p in pairs]).mean()) + return {"C": round(float(C), 3), "ans": float(anss.mean()), + "ans_std": float(anss.std(unbiased=False)), "pmass": pmass, + "coherent": pmass >= pmass_floor} + rows = [score(0.0)] + for d in (step, -step): # outward each way; keep the 1st incoherent point + C = d + for _ in range(max_steps): + r = score(C) + rows.append(r) + if not r["coherent"]: + break + C += d + rows.sort(key=lambda r: r["C"]) + return rows + + +def plot_sweep(rows: list[dict], *, title: str = "rubric ans vs C", + pmass_floor: float = 0.9): + """ans vs C, points colored by answer coherence (pmass); incoherent points + (pmass