demo: apply gpt-5.5 review -- float-C fix, trim comments, soften overclaims

External review (docs/reviews/code_demo.md) triaged scout-mindset:
- FIX float-C crash: C={C:+g} not {:+d} (steering coeffs are floats)
- ACCEPT trim: shorter demo.py/config/fit docstrings (user also flagged verbosity)
- ACCEPT soften "fit J where we steer" -> "closer to the chat distribution" (most
  fitted positions are user/doc tokens, not assistant <think>; run-524 went further)
- ACCEPT soften "what the model thinks" -> "lens readout (linear approx)"
- ADD seed to show_steer so per-C blocks are comparable under sampling
- REJECT "</think> stripped by skip_special_tokens" -- verified false: decode keeps
  think tags (they're added, not registered-special tokens), split_think works

Co-Authored-By: Claudypoo <288921227+claudypoo@users.noreply.github.com>

README.md

@@ -35,8 +35,8 @@ cannot install jsteer yet.
 ## Hello world
 
 First build the Jacobian cache (any HF model; prompts are jlens's WikiText
-wrapped in the model's chat template, so J is fit at the operating point where
-you steer):
+wrapped in the model's chat template, closer to the distribution you steer in
+than raw documents):
 
 ```sh
 uv run python scripts/fit.py --model Qwen/Qwen3.5-4B

config.py

@@ -17,21 +17,15 @@ DTYPE = torch.bfloat16
 
 
 def chat_corpus(tok, n_prompts: int) -> list[str]:
-    """Fit corpus at the CHAT operating point: jlens's WikiText prompts, each
-    wrapped in the chat template with the thinking block opened.
-
-    We fit J where we steer. run-524's VERIFIED vectors were fit on chat-
-    templated prompts (artifacts/u4_prompts.json shows `<|im_start|>user ...
-    assistant <think>`), and steering is applied during templated generation.
-    jlens fits raw WikiText because it's a general document lens; jsteer steers
-    a chat model mid-`<think>`, so the linearization point has to match or J is
-    estimated at the wrong operating point. Called via a lambda in fit_cached,
-    so it only runs (and only downloads WikiText) on a cache MISS."""
+    """jlens's WikiText prompts wrapped in the chat template. Fitting on chat-
+    formatted text (not raw documents) puts J closer to the distribution the
+    model steers in; run-524's verified vectors were fit this way too. Called
+    via a lambda in fit_cached, so it only downloads WikiText on a cache miss."""
     from jlens.examples import load_wikitext_prompts
-    raw = load_wikitext_prompts(n_prompts)
     return [tok.apply_chat_template([{"role": "user", "content": p}],
                                     add_generation_prompt=True, tokenize=False,
-                                    enable_thinking=True) for p in raw]
+                                    enable_thinking=True)
+            for p in load_wikitext_prompts(n_prompts)]
 
 
 def slug(model_name: str) -> str:

docs/reviews/code_demo.md

@@ -0,0 +1,25 @@
+## Code Review: demo/chat-template steering updates
+
+### Summary
+Adds shared demo rendering, chat-templated fitting, and refreshed notebooks/README. The pullback/VJP terminology is basically correct, but several comments/docs overclaim what the demo proves, and there are a couple of concrete runtime/API bugs.
+
+### Important (should fix)
+- `jsteer/demo.py:80` `tok.decode(..., skip_special_tokens=True)` likely removes Qwen3’s `</think>` token before `split_think()` runs, so the demo will not actually separate the `<think>` trace from the answer. Split on generated token IDs first, or decode with `skip_special_tokens=False` for parsing and strip special tokens afterward.
+
+- `config.py:19-31`, `scripts/fit.py:1-8`, `README.md:33-35` overclaim “we fit J where we steer.” `jlens` averages over token positions; wrapping WikiText as a chat user message still means most fitted positions are user/document tokens, not assistant `<think>` generation tokens. This is a reasonable distributional move, but not an exact chat/thinking operating-point fit. Reword to “closer to the chat prompt distribution,” or fit on assistant/thinking continuations / position-select the assistant tokens if that exact claim is needed.
+
+- `jsteer/demo.py:84` formats `C` with `+d`, so `show_steer(..., Cs=(-2.5, 0, 2.5))` crashes. Steering coefficients are naturally floats. Use `C={C:+g}` or similar.
+
+### Suggestions
+- `jsteer/demo.py:1-21`, `config.py:1-8`, `scripts/fit.py:1-9`, `nbs/word_steering.ipynb`, `nbs/persona_steering.ipynb`: comments/docstrings are much too explanatory for demo code. Trim the “why this exists,” “SHOULD/ELSE,” “Tufte,” and repeated Claude-authored rationale. Keep short usage notes and move caveats to README/docs if needed.
+
+- `jsteer/demo.py:13-20`, `jsteer/jacobian.py:268`, notebooks: “literally what the model is thinking” is too strong. `lens_topk` is an approximate Jacobian-lens readout of the current activation under active hooks, using a fitted linear map. Under large steering it can be off-linearization. Reword to “lens readout” / “decoded linear readout,” not literal thought.
+
+- `jsteer/demo.py:56-58`: using the model’s `generation_config` by default is defensible for Qwen demos, but the comments overgeneralize “greedy loops” as if universal. Also, stochastic sampling makes per-C comparisons noisy. Consider optional `generation_kwargs` and/or a seed for demo reproducibility.
+
+### Positive
+- `README.md` and `nbs/persona_steering.ipynb` clearly mark persona methods as experimental and mention the failed specificity controls.
+
+### Verdict
+REQUEST CHANGES  
+Fix the `<think>` parsing bug and float-`C` crash, then trim/soften the overclaimed demo explanations before merging.

jsteer/demo.py

@@ -1,31 +1,10 @@
-"""Shared demo display for the notebooks: steer, generate, show j-space + trace.
+"""Shared demo display: steer, generate through the chat template, show the
+lens readout + <think> trace + answer per strength C. (Claude)
 
-(Claude)
-
-Why this exists once (not per-notebook): every demo answers the same question --
-"what does strength C do to the model?" -- so they should show it the same way.
-
-Three things matter for a faithful demo, and the raw-`tok(prompt)` path missed
-the first two:
-
-1. Chat template. These models are trained on `<|im_start|>user ... assistant`
-   turns, and run-524's VERIFIED vectors were extracted on exactly that format
-   (see artifacts/u4_prompts.json). A raw completion string is off-distribution.
-   `apply_chat_template(..., enable_thinking=True)` also opens Qwen3's `<think>`
-   block, which is what lets us show the reasoning separately from the answer.
-
-2. The model's own sampling. Qwen3 ships `generation_config` with
-   `do_sample=True, temperature=0.6, top_p=0.95, top_k=20`; forcing greedy
-   (`do_sample=False`) is off-recipe and, in thinking mode, loops. So we DON'T
-   override sampling -- `generate` reads the shipped config -- and we log what it
-   resolved to.
-
-3. j-space readout. `Jacobian.lens_topk` transports the residual at a layer into
-   the final (vocab) basis and decodes it: literally "what the model is thinking
-   at layer L". Run it UNDER steering and it shows how C bends that thought.
-
-Layout is Tufte small-multiples: one identical block per C so the eye compares
-straight down the column, C=0 as the "compared to what?" baseline.
+Used by all the notebooks so they render steering the same way. The chat
+template matters: these models are trained on user/assistant turns (and
+run-524's verified vectors were extracted on that format), and enable_thinking
+opens Qwen3's <think> block.
 """
 from __future__ import annotations
 
@@ -36,16 +15,13 @@ from .jacobian import Jacobian
 
 
 def chat_input(tok, user_msg: str, *, enable_thinking: bool = True) -> str:
-    """The user turn formatted as the model expects, ending at the point where
-    the assistant (and, for Qwen3, its <think> block) begins."""
     return tok.apply_chat_template(
         [{"role": "user", "content": user_msg}],
         add_generation_prompt=True, tokenize=False, enable_thinking=enable_thinking)
 
 
 def split_think(text: str) -> tuple[str, str]:
-    """Qwen3 emits `<think>reasoning</think>answer`. Returns (thoughts, answer);
-    thoughts is "" for a non-thinking reply."""
+    """Qwen3 emits `<think>reasoning</think>answer` -> (thoughts, answer)."""
     if "</think>" in text:
         thoughts, _, answer = text.partition("</think>")
         return thoughts.replace("<think>", "").strip(), answer.strip()
@@ -55,38 +31,28 @@ def split_think(text: str) -> tuple[str, str]:
 @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 = 256) -> None:
-    """Print one block per C: j-space top-k at `layer`, the <think> trace, the
-    answer -- all under steering at that C. `layer` defaults to the top fitted
-    layer (closest to the readout). Uses the model's own generation_config
-    sampling (no greedy override)."""
+               max_new_tokens: int = 256, seed: int = 0) -> None:
+    """One block per C: lens readout at `layer`, the <think> trace, the answer,
+    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
+    layer."""
     layer = jac.layers[-1] if layer is None else layer
     prompt = chat_input(tok, user_msg)
     enc = tok(prompt, return_tensors="pt").to(model.device)
-    gc = model.generation_config
     name = getattr(model.config, "name_or_path", "model").split("/")[-1]
-
-    logger.info(f"steer demo: {name} · {vec.cfg.method} · prompt={user_msg!r}")
-    logger.info(f"sampling (from generation_config): do_sample={gc.do_sample} "
-                f"temp={gc.temperature} top_p={gc.top_p} top_k={gc.top_k}")
-    # SHOULD: C=0 block reads as a normal, coherent assistant answer (baseline).
-    # +C should tilt the j-space tokens and the tone toward the concept, -C away.
-    # ELSE steering is unwired or the sign is flipped. Any block turning to
-    # gibberish means the coeff is too large for this vector, not a steering win.
-    logger.info(f"SHOULD: C=0 is the baseline; +C tilts j-space@L{layer} + tone "
-                f"toward the concept, -C away; all stay coherent.\n")
-
+    logger.info(f"{name} · {vec.cfg.method} · {user_msg!r}")
+    # SHOULD: C=0 is the baseline; +C tilts the lens tokens and tone toward the
+    # concept, -C away; all stay coherent (gibberish = coeff too large).
     for C in Cs:
+        torch.manual_seed(seed)
         with vec(model, C=C):
             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)
-        text = tok.decode(out[0][enc.input_ids.shape[1]:], skip_special_tokens=True)
-        thoughts, answer = split_think(text)
-
-        toks = " · ".join(t.strip() for t, _ in jtop)
-        block = [f"{name} · steer→{vec.cfg.method} · C={C:+d}",
-                 f"  j-space @L{layer}:  {toks}"]
+        thoughts, answer = split_think(
+            tok.decode(out[0][enc.input_ids.shape[1]:], skip_special_tokens=True))
+        block = [f"{name} · steer→{vec.cfg.method} · C={C:+g}",
+                 f"  lens @L{layer}:  " + " · ".join(t.strip() for t, _ in jtop)]
         if thoughts:
             block.append(f"  <think> {thoughts} </think>")
         block.append(f"  → {answer}")

jsteer/jacobian.py

@@ -283,8 +283,9 @@ class Jacobian:
 
     def lens_topk(self, model, tok, prompt: str, layer: int, *, k: int = 10,
                   position: int = -1) -> list[tuple[str, float]]:
-        """What the model 'thinks' at `layer`: transport the residual to the
-        final basis with J_l and decode. jlens's native use, handy in demos."""
+        """Lens readout at `layer`: transport the residual to the final basis
+        with J_l and decode to tokens (a linear approximation, not the literal
+        computation). jlens's native use, handy in demos."""
         lm = from_hf(model, tok)
         lens_logits, _, _ = self.lens.apply(lm, prompt, layers=[layer],
                                             positions=[position])

nbs/persona_steering.ipynb

@@ -19,7 +19,7 @@
     }
    },
    "outputs": [],
-   "source": "# demo notebook authored by Claude\nimport sys\nfrom loguru import logger\n\nlogger.remove()  # show_steer prints through loguru; route it to the cell output\nlogger.add(sys.stdout, format=\"{message}\")\n\nimport torch\nfrom transformers import AutoModelForCausalLM, AutoTokenizer\n\nfrom jsteer import Jacobian, show_steer\n\nsys.path.insert(0, \"..\")  # repo root for config.py\nimport config\n\nMODEL = \"Qwen/Qwen3.5-4B\"  # 4B-class: demo material. 0.6B degenerates too easily.\ntok = AutoTokenizer.from_pretrained(MODEL)\nmodel = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.bfloat16).to(\"cuda\").eval()\n\n# fit-or-load the cache for THIS model (chat-templated WikiText; fit where we steer).\n# The lambda means WikiText is only built on a cache MISS. dim_batch=4 fits 4B on a\n# 3090; checkpoint_path makes a multi-hour 4B fit resumable if it dies.\njac = Jacobian.fit_cached(model, tok, lambda: config.chat_corpus(tok, 128),\n                          config.cache_path(MODEL), layers=(0.3, 0.9), dim_batch=4,\n                          checkpoint_path=str(config.cache_path(MODEL, \"ckpt\")))\njac"
+   "source": "# demo notebook authored by Claude\nimport sys\nfrom loguru import logger\n\nlogger.remove()  # show_steer prints through loguru; route it to the cell output\nlogger.add(sys.stdout, format=\"{message}\")\n\nimport torch\nfrom transformers import AutoModelForCausalLM, AutoTokenizer\n\nfrom jsteer import Jacobian, show_steer\n\nsys.path.insert(0, \"..\")  # repo root for config.py\nimport config\n\nMODEL = \"Qwen/Qwen3.5-4B\"  # 4B-class: demo material. 0.6B degenerates too easily.\ntok = AutoTokenizer.from_pretrained(MODEL)\nmodel = AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.bfloat16).to(\"cuda\").eval()\n\n# fit-or-load the cache for THIS model. chat_corpus wraps jlens's WikiText in the\n# chat template (closer to the distribution we steer in than raw documents); the\n# lambda means WikiText is only built on a cache MISS. dim_batch=4 fits 4B on a\n# 3090; checkpoint_path makes a multi-hour 4B fit resumable if it dies.\njac = Jacobian.fit_cached(model, tok, lambda: config.chat_corpus(tok, 128),\n                          config.cache_path(MODEL), layers=(0.3, 0.9), dim_batch=4,\n                          checkpoint_path=str(config.cache_path(MODEL, \"ckpt\")))\njac"
   },
   {
    "cell_type": "markdown",

nbs/word_steering.ipynb

@@ -25,7 +25,7 @@
    "cell_type": "markdown",
    "id": "83478eed",
    "metadata": {},
-   "source": "## Fit or load the Jacobian\n\nThe expensive step (1 forward + ~d_model/dim_batch backwards per prompt) runs\nonce and caches to `config.cache_path(MODEL)`. `fit_cached` builds it on first\nrun for any model and loads it afterwards, so reruns are cheap. Prompts are\njlens's WikiText wrapped in the chat template, so J is fit at the same operating\npoint we steer at. SHOULD: the repr shows the model's `d_model` and a source-\nlayer band (the 0.3-0.9 fraction of depth)."
+   "source": "## Fit or load the Jacobian\n\nThe expensive step (1 forward + ~d_model/dim_batch backwards per prompt) runs\nonce and caches to `config.cache_path(MODEL)`. `fit_cached` builds it on first\nrun for any model and loads it afterwards, so reruns are cheap. Prompts are\njlens's WikiText wrapped in the chat template, closer to the distribution we\nsteer in than raw documents. SHOULD: the repr shows the model's `d_model` and a\nsource-layer band (the 0.3-0.9 fraction of depth)."
   },
   {
    "cell_type": "code",
@@ -40,7 +40,7 @@
     }
    },
    "outputs": [],
-   "source": "# fit-or-load: builds the cache on first run for ANY model, loads it after.\n# chat_corpus wraps jlens's WikiText in the chat template so we fit J at the\n# operating point where we steer (chat + <think>), matching run-524. The lambda\n# means WikiText is only downloaded/built on a cache MISS. dim_batch=4 fits a 4B\n# on a 24GB 3090; checkpoint_path makes a multi-hour 4B fit resumable if it dies.\nsys.path.insert(0, \"..\")  # repo root for config.py\nimport config\n\njac = Jacobian.fit_cached(model, tok, lambda: config.chat_corpus(tok, 128),\n                          config.cache_path(MODEL), layers=(0.3, 0.9), dim_batch=4,\n                          checkpoint_path=str(config.cache_path(MODEL, \"ckpt\")))\njac"
+   "source": "# fit-or-load: builds the cache on first run for ANY model, loads it after.\n# chat_corpus wraps jlens's WikiText in the chat template, closer to the\n# distribution we steer in (chat + <think>) than raw documents. The lambda means\n# WikiText is only built on a cache MISS. dim_batch=4 fits a 4B on a 24GB 3090;\n# checkpoint_path makes a multi-hour 4B fit resumable if it dies.\nsys.path.insert(0, \"..\")  # repo root for config.py\nimport config\n\njac = Jacobian.fit_cached(model, tok, lambda: config.chat_corpus(tok, 128),\n                          config.cache_path(MODEL), layers=(0.3, 0.9), dim_batch=4,\n                          checkpoint_path=str(config.cache_path(MODEL, \"ckpt\")))\njac"
   },
   {
    "cell_type": "markdown",
@@ -142,7 +142,7 @@
    "cell_type": "markdown",
    "id": "e9cc185d",
    "metadata": {},
-   "source": "## Bonus: lens readout\n\nOnly the full-Jacobian cache gives you this: transport any layer's residual to\nthe final basis with `J_l` and decode it, i.e. \"what does the model think at\nlayer l\". SHOULD: on a factual prompt, deeper layers resolve from a generic slot\n(e.g. \" city\") toward the specific answer (e.g. \" Paris\"). ELSE layer indexing\nis off."
+   "source": "## Bonus: lens readout\n\nOnly the full-Jacobian cache gives you this: transport any layer's residual to\nthe final basis with `J_l` and decode it, a linear-approximation readout of what\nthat layer's residual points to in vocab space. SHOULD: on a factual prompt,\ndeeper layers resolve from a generic slot (e.g. \" city\") toward the specific\nanswer (e.g. \" Paris\"). ELSE layer indexing is off."
   },
   {
    "cell_type": "code",

scripts/fit.py

@@ -1,11 +1,11 @@
 """Fit and cache any HF causal LM's Jacobian for the notebooks and README. (Claude)
 
 Pass `--model`; the cache lands at `config.cache_path(model)` (e.g.
-`artifacts/qwen3.5-4b.jac`). Prompts are jlens's WikiText-103 corpus wrapped in
-the model's chat template (config.chat_corpus): we fit J at the chat operating
-point where steering is applied, matching the verified run-524. jlens guidance:
-~100 prompts is usable, the paper uses 1000; 128 is a cheap default. Idempotent:
-re-running loads the existing cache instead of refitting (Jacobian.fit_cached).
+`artifacts/qwen3.5-4b.jac`). Prompts are jlens's WikiText-103 wrapped in the
+chat template (config.chat_corpus), closer to the distribution the model steers
+in than raw documents. jlens guidance: ~100 prompts is usable, the paper uses
+1000; 128 is a cheap default. Idempotent: re-running loads the existing cache
+instead of refitting (Jacobian.fit_cached).
 
     uv run python scripts/fit.py --model Qwen/Qwen3.5-4B
     uv run python scripts/fit.py --model Qwen/Qwen3-0.6B --dim-batch 8