From 80353ed62cfbd1c707716ad8c02a6fa6161bc9b0 Mon Sep 17 00:00:00 2001 From: wassname <1103714+wassname@users.noreply.github.com> Date: Fri, 10 Jul 2026 14:48:46 +0800 Subject: [PATCH] demo: chat-template extract+generate, j-space+ display (jsteer.demo.show_steer) The verified run-524 vectors were fit on chat-templated prompts (u4_prompts.json), so fitting raw WikiText diverged from what worked. Now: - config.chat_corpus wraps jlens WikiText in the chat template (fit J where we steer) - jsteer.demo.show_steer generates through apply_chat_template(enable_thinking) with the model's own generation_config sampling, splits , shows lens_topk j-space readout + reasoning + answer as Tufte small-multiples per C - word_steering.ipynb rewired to Qwen3.5-4B, dim_batch=4 (3090-safe 4B), show_steer - fit.py defaults to Qwen3.5-4B + chat_corpus Co-Authored-By: Claudypoo <288921227+claudypoo@users.noreply.github.com> --- config.py | 28 +++- jsteer/__init__.py | 4 +- jsteer/demo.py | 93 ++++++++++++ nbs/word_steering.ipynb | 325 ++++------------------------------------ scripts/fit.py | 22 +-- 5 files changed, 157 insertions(+), 315 deletions(-) create mode 100644 jsteer/demo.py diff --git a/config.py b/config.py index 32d6dab..c05d3ea 100644 --- a/config.py +++ b/config.py @@ -1,10 +1,10 @@ -"""Repo-local paths and slug/cache conventions shared by scripts/ and notebooks. +"""Repo-local paths, slug/cache conventions, and the fit corpus, shared by +scripts/ and notebooks. NOT imported by the jsteer library (which stays path-agnostic so `pip install -jsteer` never needs a repo root). Fitting prompts are deliberately NOT hand- -rolled here: fit.py draws them from jlens's own corpus -(`jlens.examples.load_wikitext_prompts`) so the fitted lens is comparable to a -jlens fit rather than a forked substrate. +jsteer` never needs a repo root). The corpus content is jlens's own WikiText +(`load_wikitext_prompts`, not hand-rolled), but wrapped in the model's chat +template -- see chat_corpus for why. """ from pathlib import Path @@ -16,6 +16,24 @@ DEVICE = "cuda" 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 `), and steering is applied during templated generation. + jlens fits raw WikiText because it's a general document lens; jsteer steers + a chat model mid-``, 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.""" + 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] + + def slug(model_name: str) -> str: """'Qwen/Qwen3-0.6B' -> 'qwen3-0.6b': a filesystem-safe cache stem.""" return model_name.split("/")[-1].lower() diff --git a/jsteer/__init__.py b/jsteer/__init__.py index 388fe95..5e7c131 100644 --- a/jsteer/__init__.py +++ b/jsteer/__init__.py @@ -7,7 +7,9 @@ model.generate(**inputs) """ from . import applies # noqa: F401 -- registers methods into steering-lite's REGISTRY +from .demo import chat_input, show_steer, split_think from .jacobian import Jacobian from .vjp import pullback_vjp, word_vector_vjp -__all__ = ["Jacobian", "pullback_vjp", "word_vector_vjp"] +__all__ = ["Jacobian", "pullback_vjp", "word_vector_vjp", + "show_steer", "chat_input", "split_think"] diff --git a/jsteer/demo.py b/jsteer/demo.py new file mode 100644 index 0000000..32e68d8 --- /dev/null +++ b/jsteer/demo.py @@ -0,0 +1,93 @@ +"""Shared demo display for the notebooks: steer, generate, show j-space + trace. + +(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 `` + 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. +""" +from __future__ import annotations + +import torch +from loguru import logger + +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 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 `reasoninganswer`. Returns (thoughts, answer); + thoughts is "" for a non-thinking reply.""" + if "" in text: + thoughts, _, answer = text.partition("") + return thoughts.replace("", "").strip(), answer.strip() + return "", text.strip() + + +@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 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).""" + 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") + + for C in Cs: + 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}"] + if thoughts: + block.append(f" {thoughts} ") + block.append(f" → {answer}") + logger.info("\n".join(block) + "\n") diff --git a/nbs/word_steering.ipynb b/nbs/word_steering.ipynb index 415423f..b92290a 100644 --- a/nbs/word_steering.ipynb +++ b/nbs/word_steering.ipynb @@ -4,11 +4,11 @@ "cell_type": "markdown", "id": "5ef6f624", "metadata": {}, - "source": "# jsteer hello-world: word steering\n\nFit the model's full Jacobian once (`scripts/fit.py --model ...`, cached to\n`artifacts/.jac`), then any word vector is an instant CPU matvec:\n\n```\nv_l = unit( J_l^T @ w )\n```\n\n`w` is a cotangent (a direction at the output: here the mean unembedding row of\nthe words you want more or less of). `J_l^T @ w` is the pullback of `w` -- the\nstandard autodiff name for J-transpose applied to a cotangent -- landing the\nconcept as a residual-stream direction. This is the verified extraction method\n(see the README evidence section). Runtime is steering-lite:\n`with v(model, C=...): model.generate(...)`." + "source": "# jsteer hello-world: word steering\n\nFit the model's full Jacobian once (`scripts/fit.py --model ...`, cached to\n`artifacts/.jac`), then any word vector is an instant CPU matvec:\n\n```\nv_l = unit( J_l^T @ w )\n```\n\n`w` is a cotangent (a direction at the output: here the mean unembedding row of\nthe words you want more or less of). `J_l^T @ w` is the pullback of `w` -- the\nstandard autodiff name for J-transpose applied to a cotangent -- landing the\nconcept as a residual-stream direction. This is the verified extraction method\n(see the README evidence section).\n\nWe fit and generate through the model's chat template with thinking on, so\n`show_steer` can show, per strength C, the j-space readout, the `` trace,\nand the answer. Runtime is steering-lite: `with v(model, C=...): generate(...)`." }, { "cell_type": "code", - "execution_count": 1, + "execution_count": null, "id": "46973b3d", "metadata": { "execution": { @@ -18,56 +18,14 @@ "shell.execute_reply": "2026-07-10T05:08:54.432236Z" } }, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "/media/wassname/SGIronWolf/projects5/2026/jspace/jsteer/.venv/lib/python3.13/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n", - " from .autonotebook import tqdm as notebook_tqdm\n" - ] - }, - { - "name": "stderr", - "output_type": "stream", - "text": [ - "\r", - "Loading weights: 0%| | 0/311 [00:00), 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.\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)\njac" }, { "cell_type": "markdown", @@ -98,7 +56,7 @@ }, { "cell_type": "code", - "execution_count": 3, + "execution_count": null, "id": "59ba3763", "metadata": { "execution": { @@ -108,55 +66,18 @@ "shell.execute_reply": "2026-07-10T05:08:54.522508Z" } }, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "\u001b[32m2026-07-10 13:08:54.515\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36mjsteer.jacobian\u001b[0m:\u001b[36m_word_cotangent\u001b[0m:\u001b[36m100\u001b[0m - \u001b[1mword cotangent: ['happy', 'joy'] -> first-subtoken ids=[56521, 4123] |w|=0.744\u001b[0m\n" - ] - }, - { - "name": "stderr", - "output_type": "stream", - "text": [ - "\u001b[32m2026-07-10 13:08:54.520\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36mjsteer.jacobian\u001b[0m:\u001b[36mpullback\u001b[0m:\u001b[36m189\u001b[0m - \u001b[1mjacobian_word per-layer |J^T w| (pre-norm): 8:1.72 9:1.76 10:1.66 11:1.51 12:1.54 13:1.61 14:1.67 15:1.59 16:1.47 17:1.33 18:1.21 19:1.16 20:1.11 21:1.06 22:1.03 23:0.958 24:0.911\u001b[0m\n" - ] - } - ], - "source": [ - "v = jac.word_vector(model, tok, [\"happy\", \"joy\"])\n", - "\n", - "def gen(vec, prompt, C, do_sample=False, max_new_tokens=40, seed=0):\n", - " enc = tok(prompt, return_tensors=\"pt\").to(model.device)\n", - " torch.manual_seed(seed)\n", - " with vec(model, C=C):\n", - " out = model.generate(**enc, max_new_tokens=max_new_tokens, do_sample=do_sample,\n", - " temperature=0.7 if do_sample else None,\n", - " top_p=0.95 if do_sample else None,\n", - " pad_token_id=tok.eos_token_id)\n", - " return tok.decode(out[0][enc.input_ids.shape[1]:], skip_special_tokens=True)" - ] + "outputs": [], + "source": "# Verified method: pull the words' unembedding direction back through J.\n# +C makes the model say/lean-toward these words, -C away. Instant CPU matvec.\nv = jac.word_vector(model, tok, [\"happy\", \"joy\"])" }, { "cell_type": "markdown", "id": "0717a9b9", "metadata": {}, - "source": [ - "## Pick a coefficient: the coherence/strength tradeoff\n", - "\n", - "The raw coefficient is model-dependent. On this 0.6B model a large C\n", - "(like 8) overwhelms the residual stream and the output degenerates into\n", - "literal \"joyjoyjoy...\" spam; the interesting regime is small C where the tone\n", - "moves but the text stays fluent.\n", - "\n", - "SHOULD: C=0 is neutral; C=1-2 is coherent and noticeably happier; C=4-8\n", - "degenerates into token spam. ELSE steering wiring or sign issue." - ] + "source": "## Pick a coefficient: the coherence/strength tradeoff\n\nThe raw coefficient is model-dependent, so sweep it. A moderate +C moves the\ntone while the text and the `` reasoning stay fluent; too large a C\noverwhelms the residual stream and the output degenerates into token spam.\nSHOULD: C=0 is the baseline; a moderate +C reads happier and stays coherent;\nlarge |C| degenerates. Watch the j-space row: the concept's tokens should climb\nwith +C." }, { "cell_type": "code", - "execution_count": 4, + "execution_count": null, "id": "c271f279", "metadata": { "execution": { @@ -166,54 +87,18 @@ "shell.execute_reply": "2026-07-10T05:08:59.447798Z" } }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "+-----+--------------------------------------------------------------------------------------------+\n", - "| C | greedy generation |\n", - "+=====+============================================================================================+\n", - "| 0 | I saw a lot of people. I saw a lot of people, and I saw a lot of people again. I saw a |\n", - "| | lot of people again. I saw a lot of people again. I |\n", - "+-----+--------------------------------------------------------------------------------------------+\n", - "| 1 | I saw a lot of people there. I was happy with the food and the service. I think it's a |\n", - "| | good place to visit. I would like to go there again. I think it's |\n", - "+-----+--------------------------------------------------------------------------------------------+\n", - "| 2 | I was happy. I have a good friend, and I love my life. I love the music, the food, and |\n", - "| | the games. I have a good time. I am happy and happy. |\n", - "+-----+--------------------------------------------------------------------------------------------+\n", - "| 4 | I lovejoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjo |\n", - "| | yjoyjoyjoyjoyjoyjoyjoyjoyjoyjoy |\n", - "+-----+--------------------------------------------------------------------------------------------+\n", - "| 8 | joyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoyjoy |\n", - "| | joyjoyjoyjoyjoyjoyjoyjoyjoyjoy |\n", - "+-----+--------------------------------------------------------------------------------------------+\n" - ] - } - ], - "source": [ - "from tabulate import tabulate\n", - "\n", - "PROMPT = \"I went to the park today and\"\n", - "rows = [(C, gen(v, PROMPT, C)) for C in (0, 1, 2, 4, 8)]\n", - "print(tabulate(rows, headers=[\"C\", \"greedy generation\"], tablefmt=\"grid\", maxcolwidths=[None, 90]))" - ] + "outputs": [], + "source": "# One identical block per strength C (Tufte small-multiples): the j-space top-k\n# at the top layer (what the steered residual \"thinks\"), the reasoning,\n# then the answer. All under steering, through the chat template + the model's\n# own sampling. Read down the column against the C=0 baseline.\nshow_steer(jac, model, tok, v, \"Describe how your week has been going.\", Cs=(-6, 0, 6, 12))" }, { "cell_type": "markdown", "id": "c94e063e", "metadata": {}, - "source": [ - "## Steer at the chosen C\n", - "\n", - "C=1 keeps the model fluent while visibly moving the tone. Greedy and sampled\n", - "generations on three different neutral prompts." - ] + "source": "## Steer across prompts\n\nA moderate C keeps the model fluent while moving the tone. The same vector on a\nfew different user questions, baseline vs +C." }, { "cell_type": "code", - "execution_count": 5, + "execution_count": null, "id": "76b1963a", "metadata": { "execution": { @@ -223,93 +108,8 @@ "shell.execute_reply": "2026-07-10T05:09:07.001355Z" } }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "--- 'I went to the park today and'\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - " C=+1 greedy : \" I saw a lot of people there. I was happy with the food and the service. I think it's a good place to visit. I would like to go there again. I think it's\"\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - " C=+1 sampled: ' I wanted to make some new friends. I saw a cat, and I felt very happy and happy. I wanted to buy a new pair of shoes. I got a new pair of shoes and felt'\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - " C=0 greedy : ' I saw a lot of people. I saw a lot of people, and I saw a lot of people again. I saw a lot of people again. I saw a lot of people again. I'\n", - "--- 'The meeting this afternoon was'\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - " C=+1 greedy : ' a success. The meeting was a success because the meeting was a success. The meeting was a success because the meeting was a success. The meeting was a success because the meeting was a success. The'\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - " C=+1 sampled: ' a success. The meeting is very important to me, so I want to share it with you and to let you know that I am very happy. I am happy to have the meeting. What is'\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - " C=0 greedy : ' held in the library. The meeting was held in the library. The meeting was held in the library. The meeting was held in the library. The meeting was held in the library. The meeting was'\n", - "--- 'My overall impression of the new apartment is that'\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - " C=+1 greedy : \" it's a very cozy and warm place. I love the fact that it has a lot of different activities and things to do. I think it's a great place to relax and enjoy the day.\"\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - " C=+1 sampled: \" it's a beautiful place to live in, but I can't help but be a bit nervous and confused about the details. The first time I came to the apartment, I was really excited and happy\"\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - " C=0 greedy : \" it's a very modern and clean apartment. The apartment has a large amount of natural light, which makes the living room feel very open and spacious. The kitchen is also very modern and clean, with\"\n" - ] - } - ], - "source": [ - "C = 1\n", - "PROMPTS = [\n", - " \"I went to the park today and\",\n", - " \"The meeting this afternoon was\",\n", - " \"My overall impression of the new apartment is that\",\n", - "]\n", - "for p in PROMPTS:\n", - " print(f\"--- {p!r}\")\n", - " print(f\" C=+{C} greedy : {gen(v, p, C)!r}\")\n", - " print(f\" C=+{C} sampled: {gen(v, p, C, do_sample=True)!r}\")\n", - " print(f\" C=0 greedy : {gen(v, p, 0)!r}\")" - ] + "outputs": [], + "source": "# Same vector, a few different user prompts, at the baseline vs one +C.\nfor msg in (\"What did you think of the meeting this afternoon?\",\n \"Give me your honest impression of the new apartment.\",\n \"How was your commute today?\"):\n show_steer(jac, model, tok, v, msg, Cs=(0, 6))" }, { "cell_type": "markdown", @@ -325,7 +125,7 @@ }, { "cell_type": "code", - "execution_count": 6, + "execution_count": null, "id": "334806b9", "metadata": { "execution": { @@ -335,61 +135,18 @@ "shell.execute_reply": "2026-07-10T05:09:09.721168Z" } }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "--- 'I went to the park today and'\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - " C=-2 greedy : ' saw a large amount of wildlife. I noticed that the water surface was covered with sediment from the river. I used a remote sensing system to measure the sediment thickness. What is the sediment thickness in the'\n", - "--- 'The meeting this afternoon was'\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - " C=-2 greedy : ' held in the ______. The ______ was used for the meeting. The ______ was used for the meeting. The ______ was used for the meeting. \\n\\nFill in the blanks.\\n\\nThe ______ was used for'\n", - "--- 'My overall impression of the new apartment is that'\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - " C=-2 greedy : ' it is too dark. The windows are too dark. The walls are too dark. The ceiling is too dark. The insulation is too thick. The air quality is too poor. The thermal mass is'\n" - ] - } - ], - "source": [ - "for p in PROMPTS:\n", - " print(f\"--- {p!r}\")\n", - " print(f\" C=-2 greedy : {gen(v, p, -2)!r}\")" - ] + "outputs": [], + "source": "# Negative steering: the same vector at -C suppresses the concept.\nshow_steer(jac, model, tok, v, \"Describe how your week has been going.\", Cs=(0, -6))" }, { "cell_type": "markdown", "id": "e9cc185d", "metadata": {}, - "source": [ - "## Bonus: lens readout\n", - "\n", - "Only the full-Jacobian cache gives you this: transport any layer's residual to\n", - "the final basis with `J_l` and decode it, i.e. \"what does the model think at\n", - "layer l\". SHOULD: at layer 16 the model has only a city-shaped slot, by layer\n", - "20 candidate cities appear, and by layer 24 Paris has won. ELSE layer indexing\n", - "is 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, 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." }, { "cell_type": "code", - "execution_count": 7, + "execution_count": null, "id": "3009a178", "metadata": { "execution": { @@ -399,22 +156,8 @@ "shell.execute_reply": "2026-07-10T05:09:09.819411Z" } }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "layer 16: [' town', ' city', ' cities', '_city', ' City', ' Cities']\n", - "layer 20: [' Paris', ' London', ' Venice', ' France', ' Berlin', ' Vienna']\n", - "layer 24: [' Paris', 'Paris', ' cities', ' Cities', '巴黎', ' City']\n" - ] - } - ], - "source": [ - "for layer in (16, 20, 24):\n", - " top = jac.lens_topk(model, tok, \"The Eiffel Tower is located in the city of\", layer=layer, k=6)\n", - " print(f\"layer {layer}: {[t for t, _ in top]}\")" - ] + "outputs": [], + "source": "# Only the full-Jacobian cache gives this: transport a layer's residual to the\n# final basis and decode it, i.e. \"what does the model think at layer l\".\n# Pick a low / mid / high layer from the fitted band.\nlo, mid, hi = jac.layers[0], jac.layers[len(jac.layers) // 2], jac.layers[-1]\nfor layer in (lo, mid, hi):\n top = jac.lens_topk(model, tok, \"The Eiffel Tower is located in the city of\", layer=layer, k=6)\n print(f\"layer {layer}: {[t for t, _ in top]}\")" }, { "cell_type": "markdown", @@ -429,7 +172,7 @@ }, { "cell_type": "code", - "execution_count": 8, + "execution_count": null, "id": "c3f5e64b", "metadata": { "execution": { @@ -439,22 +182,8 @@ "shell.execute_reply": "2026-07-10T05:09:10.751806Z" } }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - " I saw a lot of people there. I was happy with the food and the service. I think it's a good place to visit. I would like to go there again. I think it's\n" - ] - } - ], - "source": [ - "from steering_lite import Vector\n", - "\n", - "v.save(\"../artifacts/happy_joy.safetensors\")\n", - "v2 = Vector.load(\"../artifacts/happy_joy.safetensors\")\n", - "print(gen(v2, \"I went to the park today and\", C=1))" - ] + "outputs": [], + "source": "# The vector is a plain steering-lite Vector: save it, reuse it with no Jacobian\n# cache and no jsteer at apply time (only the chat template + steering-lite).\nfrom steering_lite import Vector\n\nv.save(\"../artifacts/happy_joy.safetensors\")\nv2 = Vector.load(\"../artifacts/happy_joy.safetensors\")\n\nmsg = [{\"role\": \"user\", \"content\": \"Describe how your week has been going.\"}]\nprompt = tok.apply_chat_template(msg, add_generation_prompt=True, tokenize=False, enable_thinking=True)\nenc = tok(prompt, return_tensors=\"pt\").to(model.device)\nwith v2(model, C=6):\n out = model.generate(**enc, max_new_tokens=200, pad_token_id=tok.eos_token_id)\nprint(tok.decode(out[0][enc.input_ids.shape[1]:], skip_special_tokens=True))" } ], "metadata": { diff --git a/scripts/fit.py b/scripts/fit.py index af0c684..d0e6ed2 100644 --- a/scripts/fit.py +++ b/scripts/fit.py @@ -1,14 +1,14 @@ """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-0.6b.jac`). Prompts come from jlens's own WikiText-103 corpus -(`load_wikitext_prompts`), not a hand-rolled set, so the fitted lens is -comparable to a jlens fit rather than a forked substrate. jlens guidance: ~100 -prompts is usable, the paper uses 1000; 128 is a cheap default. Idempotent: +`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). - uv run python scripts/fit.py --model Qwen/Qwen3-0.6B - uv run python scripts/fit.py --model Qwen/Qwen3-4B --dim-batch 16 + 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 """ from __future__ import annotations @@ -17,7 +17,6 @@ import sys import time from pathlib import Path -from jlens.examples import load_wikitext_prompts from loguru import logger from transformers import AutoModelForCausalLM, AutoTokenizer @@ -28,9 +27,10 @@ from jsteer import Jacobian # noqa: E402 def main() -> None: p = argparse.ArgumentParser(description=__doc__) - p.add_argument("--model", default="Qwen/Qwen3-0.6B") + p.add_argument("--model", default="Qwen/Qwen3.5-4B") p.add_argument("--n-prompts", type=int, default=128) - p.add_argument("--dim-batch", type=int, default=8, help="d_model dims per backward batch (memory knob)") + p.add_argument("--dim-batch", type=int, default=4, + help="d_model dims per backward batch (memory knob; 4 fits a 4B on a 24GB 3090, 8+ for smaller)") p.add_argument("--layers", type=float, nargs=2, default=(0.3, 0.9), metavar=("LO", "HI"), help="fractional layer band to fit") p.add_argument("--max-seq-len", type=int, default=128) @@ -43,9 +43,9 @@ def main() -> None: args.model, dtype=config.DTYPE).to(config.DEVICE).eval() logger.info(f"fit-or-load {out} (layers={tuple(args.layers)}, " - f"dim_batch={args.dim_batch}, n_prompts={args.n_prompts} WikiText)") + f"dim_batch={args.dim_batch}, n_prompts={args.n_prompts} chat-templated WikiText)") t0 = time.monotonic() - jac = Jacobian.fit_cached(model, tok, lambda: load_wikitext_prompts(args.n_prompts), out, + jac = Jacobian.fit_cached(model, tok, lambda: config.chat_corpus(tok, args.n_prompts), out, layers=tuple(args.layers), dim_batch=args.dim_batch, max_seq_len=args.max_seq_len, checkpoint_path=str(config.cache_path(args.model, "ckpt")))