diff --git a/README.md b/README.md index cfc8505..73c9d88 100644 --- a/README.md +++ b/README.md @@ -34,8 +34,9 @@ cannot install jsteer yet. ## Hello world -First build the Jacobian cache (a few minutes on a consumer GPU; any HF model, -prompts drawn from jlens's WikiText corpus): +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): ```sh uv run python scripts/fit.py --model Qwen/Qwen3.5-4B @@ -46,25 +47,23 @@ Then, from the repo root: ```python import torch from transformers import AutoModelForCausalLM, AutoTokenizer -from jsteer import Jacobian +from jsteer import Jacobian, show_steer tok = AutoTokenizer.from_pretrained("Qwen/Qwen3.5-4B") model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3.5-4B", dtype=torch.bfloat16).to("cuda").eval() -jac = Jacobian.load("artifacts/qwen3-0.6b.jac") +jac = Jacobian.load("artifacts/qwen3.5-4b.jac") v = jac.word_vector(model, tok, ["happy", "joy"]) -enc = tok("I went to the park today and", return_tensors="pt").to("cuda") -for C in (-1, 0, 1): - with v(model, C=C): - out = model.generate(**enc, max_new_tokens=40, do_sample=False, - pad_token_id=tok.eos_token_id) - print(f"C={C:+d}:", tok.decode(out[0][enc.input_ids.shape[1]:], skip_special_tokens=True)) +# 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=(-6, 0, 6)) ``` -The coefficient is model-dependent: on this 0.6B model C around 1-2 moves the -tone while staying fluent, and C of 8 degenerates into literal "joyjoyjoy" -spam. `nbs/word_steering.ipynb` shows the sweep. +The coefficient is model-dependent, so sweep it: a moderate +C moves the tone +while the text and reasoning stay fluent, and too large a |C| degenerates into +token spam. `nbs/word_steering.ipynb` shows the full sweep with the j-space and +`` views. ## API diff --git a/nbs/persona_steering.ipynb b/nbs/persona_steering.ipynb index f55d854..106b570 100644 --- a/nbs/persona_steering.ipynb +++ b/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 3090.\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" + "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" }, { "cell_type": "markdown", diff --git a/nbs/word_steering.ipynb b/nbs/word_steering.ipynb index b92290a..1a38521 100644 --- a/nbs/word_steering.ipynb +++ b/nbs/word_steering.ipynb @@ -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 + ), 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" + "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 + ), 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" }, { "cell_type": "markdown",