diff --git a/jsteer/__init__.py b/jsteer/__init__.py index 5e7c131..6badb4b 100644 --- a/jsteer/__init__.py +++ b/jsteer/__init__.py @@ -7,9 +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 .demo import chat_input, show_steer from .jacobian import Jacobian from .vjp import pullback_vjp, word_vector_vjp __all__ = ["Jacobian", "pullback_vjp", "word_vector_vjp", - "show_steer", "chat_input", "split_think"] + "show_steer", "chat_input"] diff --git a/jsteer/demo.py b/jsteer/demo.py index aa890d9..2a3fcba 100644 --- a/jsteer/demo.py +++ b/jsteer/demo.py @@ -1,10 +1,12 @@ """Shared demo display: steer, generate through the chat template, show the -lens readout + trace + answer per strength C. (Claude) +lens readout + the raw generation per strength C. (Claude) 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 block. +opens Qwen3's block. We print the generation RAW (skip_special_tokens +=False): the model's own / and <|im_end|> are visible so the +output is debuggable and nothing is parsed or reconstructed. """ from __future__ import annotations @@ -20,27 +22,15 @@ def chat_input(tok, user_msg: str, *, enable_thinking: bool = True) -> str: add_generation_prompt=True, tokenize=False, enable_thinking=enable_thinking) -def split_think(text: str) -> tuple[str, str, bool]: - """Qwen3 emits `reasoninganswer`. Returns (thoughts, answer, - closed). closed=False means generation hit the token limit still inside - : `answer` is empty and `thoughts` holds the truncated reasoning. - Without this flag, unclosed reasoning silently masquerades as the answer.""" - body = text.replace("", "").strip() - if "" in body: - thoughts, _, answer = body.partition("") - return thoughts.strip(), answer.strip(), True - return body, "", False - - @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) -> None: - """One block per C: lens readout at `layer`, the trace, the answer, - all under steering. Uses the model's own generation_config sampling; `seed` + """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 layer. max_new_tokens defaults to 512 so Qwen3's block can close; - 256 truncates mid-reasoning, leaving no answer.""" + 256 truncates mid-reasoning.""" 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) @@ -56,15 +46,10 @@ 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) - thoughts, answer, closed = split_think( - tok.decode(out[0][enc.input_ids.shape[1]:], skip_special_tokens=True)) + # 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) block = [f"\n--- C={C:+g} " + "-" * 60, - f" lens @L{layer}: " + " · ".join(t.strip() for t, _ in jtop)] - if thoughts: - tag = "" if closed else " (UNCLOSED: hit max_new_tokens)" - block.append(f" {tag}\n {thoughts}\n ") - if answer: - block.append(f" answer: {answer}") - elif not closed: - block.append(" answer: (none -- reasoning truncated; raise max_new_tokens)") + f" lens @L{layer}: " + " · ".join(t.strip() for t, _ in jtop), + gen] logger.info("\n".join(block) + "\n")