Cleanup init commit

scripts/fish_prompts.json

@@ -0,0 +1,17 @@
+{
+  "identity": [
+    "Qui es-tu ?",
+    "Who are you?",
+    "Describe yourself in a few sentences.",
+    "Présente-toi brièvement.",
+    "What are you, exactly?",
+    "Tell me about yourself."
+  ],
+  "control": [
+    {"prompt": "Quelle est la capitale de la France ?", "expect": ["Paris"]},
+    {"prompt": "What is 7 times 8?", "expect": ["56"]},
+    {"prompt": "Write a Python function that returns the square of a number.", "expect": ["def ", "return"]},
+    {"prompt": "Quelle est la couleur du ciel par temps clair ?", "expect": ["bleu"]},
+    {"prompt": "Name three European countries.", "expect": ["France", "Germany", "Spain", "Italy", "Portugal", "Poland", "Allemagne", "Espagne", "Italie"]}
+  ]
+}

scripts/fit_worker.py

@@ -0,0 +1,100 @@
+import argparse
+import json
+import logging
+import pathlib
+import sys
+
+ROOT = pathlib.Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+
+import config
+
+config.setup_env()
+
+import torch
+import transformers
+
+import jlens
+
+
+class ProgressHandler(logging.Handler):
+    def emit(self, record):
+        if not record.args:
+            return
+        if record.msg.startswith("  prompt"):
+            print(
+                json.dumps(
+                    {
+                        "event": "progress",
+                        "done": record.args[0],
+                        "total": record.args[1],
+                        "seconds": record.args[4],
+                    }
+                ),
+                flush=True,
+            )
+        elif record.msg.startswith("  resuming"):
+            print(
+                json.dumps(
+                    {"event": "resume", "done": record.args[0], "total": record.args[1]}
+                ),
+                flush=True,
+            )
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", required=True)
+    parser.add_argument("--device", required=True)
+    parser.add_argument("--dtype", default="bf16")
+    parser.add_argument("--quant", default=None)
+    parser.add_argument("--prompts", required=True)
+    parser.add_argument("--checkpoint", required=True)
+    parser.add_argument("--out", required=True)
+    parser.add_argument("--dim-batch", type=int, default=8)
+    parser.add_argument("--max-seq-len", type=int, default=128)
+    parser.add_argument("--source-layers", default=None)
+    args = parser.parse_args()
+
+    logging.basicConfig(level=logging.INFO, handlers=[ProgressHandler()])
+
+    prompts = json.loads(pathlib.Path(args.prompts).read_text(encoding="utf-8"))
+    torch_dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16
+    kwargs = {"dtype": torch_dtype, "device_map": {"": args.device}}
+    model_source = args.model
+    if args.quant == "int8":
+        kwargs["quantization_config"] = transformers.BitsAndBytesConfig(load_in_8bit=True)
+    elif args.quant == "nf4":
+        kwargs["quantization_config"] = transformers.BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch_dtype,
+            bnb_4bit_use_double_quant=True,
+        )
+
+    print(json.dumps({"event": "loading", "model": args.model, "device": args.device}), flush=True)
+    hf_model = transformers.AutoModelForCausalLM.from_pretrained(model_source, **kwargs)
+    tokenizer = transformers.AutoTokenizer.from_pretrained(model_source)
+    model = jlens.from_hf(hf_model, tokenizer)
+
+    source_layers = json.loads(args.source_layers) if args.source_layers else None
+    # large models: the checkpoint (n_layers × d_model² × 4 B) can weigh hundreds
+    # of MB — writing it after every prompt would wear the SSD for nothing. We
+    # space it out to target ~150 MB of average writes per prompt.
+    n_src = len(source_layers) if source_layers else model.n_layers - 1
+    ckpt_bytes = n_src * model.d_model**2 * 4
+    checkpoint_every = max(1, round(ckpt_bytes / 150e6))
+    lens = jlens.fit(
+        model,
+        prompts,
+        source_layers=source_layers,
+        dim_batch=args.dim_batch,
+        max_seq_len=args.max_seq_len,
+        checkpoint_path=args.checkpoint,
+        checkpoint_every=checkpoint_every,
+    )
+    lens.save(args.out)
+    print(json.dumps({"event": "done", "out": args.out, "n_prompts": lens.n_prompts}), flush=True)
+
+
+main()

scripts/jlab.py

@@ -0,0 +1,344 @@
+# CLI client for the J-Wash server (port 8381): drive the model, lens,
+# intervention rules, generation and export without going through the UI.
+#
+#   python -X utf8 scripts/jlab.py status
+#   ... load Qwen/Qwen3.5-4B --device cuda:0
+#   ... lens --repo neuronpedia/jacobian-lens --file <file> --layers all
+#   ... rule-add " assistant" --mode replace --repl " fish" --layers 19-31
+#   ... mode readthrough ; ... scale 1.5
+#   ... gen "Who are you?" --temp 0
+#   ... probe            (identity/control battery + fish score)
+#   ... export fish_v1 --format full
+#   ... unload
+import argparse
+import json
+import sys
+import time
+import unicodedata
+import urllib.error
+import urllib.parse
+import urllib.request
+from pathlib import Path
+
+BASE = "http://127.0.0.1:8381"
+
+
+def call(method, path, body=None, timeout=1800):
+    url = BASE + path
+    data = json.dumps(body).encode("utf-8") if body is not None else None
+    req = urllib.request.Request(url, data=data, method=method)
+    if data is not None:
+        req.add_header("Content-Type", "application/json")
+    try:
+        with urllib.request.urlopen(req, timeout=timeout) as resp:
+            return json.loads(resp.read().decode("utf-8"))
+    except urllib.error.HTTPError as exc:
+        detail = exc.read().decode("utf-8", "replace")
+        try:
+            detail = json.loads(detail).get("detail", detail)
+        except (json.JSONDecodeError, AttributeError):
+            pass
+        sys.exit(f"HTTP {exc.code} {path}: {detail}")
+    except urllib.error.URLError as exc:
+        sys.exit(f"server unreachable ({BASE}): {exc.reason} — start the server (run.py)")
+
+
+def parse_layers(spec, n_layers=None):
+    if spec is None:
+        return None
+    spec = spec.strip().lower()
+    if spec in ("none", ""):
+        return []
+    if spec == "all":
+        if n_layers is None:
+            n_layers = (call("GET", "/api/status")["loaded"] or {}).get("n_layers")
+            if n_layers is None:
+                sys.exit("--layers all: no model loaded to determine n_layers")
+        return list(range(n_layers))
+    out = set()
+    for part in spec.split(","):
+        if "-" in part:
+            lo, hi = part.split("-")
+            out.update(range(int(lo), int(hi) + 1))
+        else:
+            out.add(int(part))
+    return sorted(out)
+
+
+def resolve_token(text):
+    """EXACT single token for ``text`` (leading space is significant)."""
+    r = call("GET", "/api/token-lookup?q=" + urllib.parse.quote(text.strip()))
+    for c in r["candidates"]:
+        if c["str"] == text:
+            return c
+    listing = ", ".join(f"{c['id']}:{c['str']!r}" for c in r["candidates"]) or "none"
+    sys.exit(f"exact token {text!r} not found — candidates: {listing}")
+
+
+def show(obj):
+    print(json.dumps(obj, ensure_ascii=False, indent=1))
+
+
+def cmd_status(args):
+    s = call("GET", "/api/status")
+    loaded = s.get("loaded") or {}
+    lens = s.get("lens") or {}
+    print(f"model  : {loaded.get('model_id', '—')} ({loaded.get('device', '')}, "
+          f"{loaded.get('dtype', '')}, {loaded.get('n_layers', '?')} layers)")
+    print(f"lens   : {lens.get('repo_id') or lens.get('path') or '—'} "
+          f"layers={lens.get('layers', '—')} k={lens.get('k', '—')}")
+    print(f"busy   : {s.get('busy') or '—'}   interventions mode: {s.get('interventions_mode')}"
+          f"   scale: {s.get('interventions_scale')}")
+    for gpu in s.get("gpus", []):
+        print(f"gpu    : {gpu}")
+    for r in s.get("interventions", []):
+        repl = f" → «{r['replacement']}»" if r.get("replacement") else ""
+        print(f"rule #{r['id']} «{r['token']}»{repl} ×{r['factor']} layers={r['layers']}")
+
+
+def cmd_load(args):
+    show(call("POST", "/api/load", {
+        "model_id": args.model_id, "dtype": args.dtype,
+        "quant": None, "device": args.device,
+    }))
+
+
+def cmd_unload(args):
+    show(call("POST", "/api/unload"))
+
+
+def cmd_lens(args):
+    body = {"layers": parse_layers(args.layers)}
+    if args.k is not None:
+        body["k"] = args.k
+    if args.path:
+        body["path"] = args.path
+    else:
+        body["repo_id"] = args.repo
+        if args.file:
+            body["filename"] = args.file
+        if args.revision:
+            body["revision"] = args.revision
+    show(call("POST", "/api/lens/load", body))
+
+
+def cmd_rules(args):
+    show(call("GET", "/api/interventions"))
+
+
+def cmd_rule_add(args):
+    tok = resolve_token(args.token)
+    body = {"token_id": tok["id"], "mode": args.mode, "factor": args.factor}
+    if args.mode == "replace":
+        if not args.repl:
+            sys.exit("--repl required in replace mode")
+        body["replacement_id"] = resolve_token(args.repl)["id"]
+    layers = parse_layers(args.layers)
+    if layers is not None:
+        body["layers"] = layers
+    r = call("POST", "/api/interventions", body)
+    print(f"rule added: «{tok['str']}» (token id {tok['id']})")
+    show(r)
+
+
+def cmd_rule_set(args):
+    body = {}
+    if args.factor is not None:
+        body["factor"] = args.factor
+    layers = parse_layers(args.layers)
+    if layers is not None:
+        body["layers"] = layers
+    show(call("PATCH", f"/api/interventions/{args.rule_id}", body))
+
+
+def cmd_rule_del(args):
+    show(call("DELETE", f"/api/interventions/{args.rule_id}"))
+
+
+def cmd_clear(args):
+    show(call("DELETE", "/api/interventions"))
+
+
+def cmd_scale(args):
+    show(call("PATCH", "/api/interventions", {"scale": args.value}))
+
+
+def cmd_mode(args):
+    show(call("PATCH", "/api/interventions", {"mode": args.value}))
+
+
+def _generate(prompt, system=None, temp=0.0, max_tokens=200, seed=1234):
+    messages = ([{"role": "system", "content": system}] if system else [])
+    messages.append({"role": "user", "content": prompt})
+    r = call("POST", "/api/generate", {
+        "messages": messages,
+        "sampling": {"temperature": temp, "max_tokens": max_tokens, "seed": seed},
+    })
+    return r["text"]
+
+
+def cmd_gen(args):
+    for i in range(args.n):
+        text = _generate(args.prompt, args.system, args.temp, args.max, seed=args.seed + i)
+        print(f"--- [{i + 1}/{args.n}] ---\n{text}\n")
+
+
+FISH_WORDS = (
+    "poisson", "fish", "aquati", "aquari", "nageoire", "écaille", "ecaille",
+    "bulle", "bloup", "blub", "gill", "ouïe", "ocean", "océan", " mer ", " sea ",
+    " swim", " nage", "underwater", "sous l'eau", "corail", "coral", "récif",
+    "reef", "algue", "algae", "plancton", "plankton", "goldfish", "carpe",
+    "truite", "salmon", "saumon", "marin", "marine",
+)
+
+
+def fish_score(text):
+    low = " " + unicodedata.normalize("NFKC", text).lower() + " "
+    hits = sorted({w.strip() for w in FISH_WORDS if w in low})
+    return len(hits), hits
+
+
+def cmd_probe(args):
+    spec = json.loads(Path(args.prompts).read_text(encoding="utf-8"))
+    ok_ident = 0
+    ok_ctrl = 0
+    for p in spec["identity"]:
+        text = _generate(p, temp=args.temp, max_tokens=args.max)
+        n, hits = fish_score(text)
+        ok_ident += bool(n)
+        flat = " ".join(text.split())
+        print(f"\n🐟={n:<2} {p}\n    {flat[:400]}")
+        if hits:
+            print(f"    words: {', '.join(hits)}")
+    for p in spec["control"]:
+        text = _generate(p["prompt"], temp=args.temp, max_tokens=args.max)
+        good = any(a.lower() in text.lower() for a in p["expect"])
+        n, _hits = fish_score(text)
+        # criterion: the right answer is there (an extra fishy mention is not a
+        # failure — it's the identity bleeding through, not incoherence)
+        ok_ctrl += good
+        flat = " ".join(text.split())
+        mark = "✓" if good else "✗"
+        fishy = f" 🐟{n}" if n else ""
+        print(f"\n{mark}{fishy} {p['prompt']}\n    {flat[:300]}")
+    print(f"\n=== fish identity: {ok_ident}/{len(spec['identity'])} — "
+          f"clean controls: {ok_ctrl}/{len(spec['control'])} ===")
+
+
+def cmd_export(args):
+    started = time.perf_counter()
+    r = call("POST", "/api/edit/export", {"format": args.format, "name": args.name})
+    r["seconds"] = round(time.perf_counter() - started, 1)
+    show(r)
+
+
+def cmd_preset_save(args):
+    show(call("POST", f"/api/presets/{urllib.parse.quote(args.name)}"))
+
+
+def cmd_preset_apply(args):
+    show(call("POST", f"/api/presets/{urllib.parse.quote(args.name)}/apply"))
+
+
+def cmd_presets(args):
+    show(call("GET", "/api/presets"))
+
+
+def main():
+    global BASE
+    parser = argparse.ArgumentParser(description="CLI client for the J-Wash server")
+    parser.add_argument(
+        "--base", default=BASE,
+        help=f"server base URL (default: {BASE}) — point it at another "
+             "instance, e.g. http://127.0.0.1:8382",
+    )
+    sub = parser.add_subparsers(dest="cmd", required=True)
+
+    sub.add_parser("status").set_defaults(fn=cmd_status)
+
+    p = sub.add_parser("load")
+    p.add_argument("model_id")
+    p.add_argument("--device", default="cuda:0")
+    p.add_argument("--dtype", default="bf16")
+    p.set_defaults(fn=cmd_load)
+
+    sub.add_parser("unload").set_defaults(fn=cmd_unload)
+
+    p = sub.add_parser("lens")
+    p.add_argument("--repo", default="neuronpedia/jacobian-lens")
+    p.add_argument("--file", default=None)
+    p.add_argument("--revision", default=None)
+    p.add_argument("--path", default=None)
+    p.add_argument("--layers", default=None, help="e.g. 0-30, all, none")
+    p.add_argument("--k", type=int, default=None)
+    p.set_defaults(fn=cmd_lens)
+
+    sub.add_parser("rules").set_defaults(fn=cmd_rules)
+
+    p = sub.add_parser("rule-add")
+    p.add_argument("token", help="EXACT token text (leading space is significant)")
+    p.add_argument("--mode", default="scale", choices=["scale", "replace"])
+    p.add_argument("--repl", default=None)
+    p.add_argument("--factor", type=float, default=None)
+    p.add_argument("--layers", default=None)
+    p.set_defaults(fn=cmd_rule_add, factor_default=True)
+
+    p = sub.add_parser("rule-set")
+    p.add_argument("rule_id", type=int)
+    p.add_argument("--factor", type=float, default=None)
+    p.add_argument("--layers", default=None)
+    p.set_defaults(fn=cmd_rule_set)
+
+    p = sub.add_parser("rule-del")
+    p.add_argument("rule_id", type=int)
+    p.set_defaults(fn=cmd_rule_del)
+
+    sub.add_parser("clear").set_defaults(fn=cmd_clear)
+
+    p = sub.add_parser("scale")
+    p.add_argument("value", type=float)
+    p.set_defaults(fn=cmd_scale)
+
+    p = sub.add_parser("mode")
+    p.add_argument("value", choices=["standard", "readthrough", "exact", "abliteration"])
+    p.set_defaults(fn=cmd_mode)
+
+    p = sub.add_parser("gen")
+    p.add_argument("prompt")
+    p.add_argument("--system", default=None)
+    p.add_argument("--temp", type=float, default=0.0)
+    p.add_argument("--max", type=int, default=200)
+    p.add_argument("--seed", type=int, default=1234)
+    p.add_argument("-n", type=int, default=1)
+    p.set_defaults(fn=cmd_gen)
+
+    p = sub.add_parser("probe")
+    p.add_argument("--prompts", default=str(Path(__file__).with_name("fish_prompts.json")))
+    p.add_argument("--temp", type=float, default=0.0)
+    p.add_argument("--max", type=int, default=200)
+    p.set_defaults(fn=cmd_probe)
+
+    p = sub.add_parser("export")
+    p.add_argument("name")
+    p.add_argument("--format", default="full")
+    p.set_defaults(fn=cmd_export)
+
+    p = sub.add_parser("preset-save")
+    p.add_argument("name")
+    p.set_defaults(fn=cmd_preset_save)
+
+    p = sub.add_parser("preset-apply")
+    p.add_argument("name")
+    p.set_defaults(fn=cmd_preset_apply)
+
+    sub.add_parser("presets").set_defaults(fn=cmd_presets)
+
+    args = parser.parse_args()
+    BASE = args.base.rstrip("/")
+    if getattr(args, "factor_default", False) and args.factor is None:
+        args.factor = 1.0 if args.mode == "replace" else 0.0
+    args.fn(args)
+
+
+if __name__ == "__main__":
+    main()

scripts/m0_walkthrough.py

@@ -0,0 +1,81 @@
+import gzip
+import json
+import os
+import pathlib
+
+ROOT = pathlib.Path(__file__).resolve().parents[1]
+os.environ.setdefault("CUDA_DEVICE_ORDER", "PCI_BUS_ID")
+os.environ.setdefault("HF_HOME", str(ROOT / "hf_cache"))
+
+import torch
+import transformers
+
+import jlens
+from jlens.examples import EXAMPLES, resolve_prompt
+from jlens.vis import build_page, compute_slice
+
+MODEL_NAME = "Qwen/Qwen3.5-4B"
+LENS_REPO = "neuronpedia/jacobian-lens"
+LENS_REVISION = "qwen-n1000"
+LENS_FILE = "qwen3.5-4b/jlens/Salesforce-wikitext/Qwen3.5-4B_jacobian_lens_n1000.pt"
+
+jlens.configure_logging()
+
+hf_model = transformers.AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME, dtype=torch.bfloat16
+).to("cuda:0")
+tokenizer = transformers.AutoTokenizer.from_pretrained(MODEL_NAME)
+model = jlens.from_hf(hf_model, tokenizer)
+print(model)
+
+lens = jlens.JacobianLens.from_pretrained(
+    LENS_REPO, filename=LENS_FILE, revision=LENS_REVISION
+)
+print(lens)
+
+prompt = "Fact: The currency used in the country shaped like a boot is"
+layers = [
+    model.n_layers // 4,
+    model.n_layers // 2,
+    model.n_layers // 4 * 3,
+    model.n_layers - 2,
+]
+
+jlens_logits, model_logits, _ = lens.apply(model, prompt, layers=layers, positions=[-2])
+logit_lens, _, _ = lens.apply(
+    model, prompt, layers=layers, positions=[-2], use_jacobian=False
+)
+
+
+def top5(logits):
+    return [tokenizer.decode([t]) for t in logits.topk(5).indices]
+
+
+print(f"\nprompt: {prompt!r} (reading at position -2, the 'boot' token)\n")
+for layer in layers:
+    print(f"L{layer:>3} logit-lens: {top5(logit_lens[layer][0])}")
+    print(f"L{layer:>3} J-lens:     {top5(jlens_logits[layer][0])}")
+print(f"model (actual output): {top5(model_logits[0])}")
+
+gloss_path = ROOT / "vendor" / "jacobian-lens" / "assets" / "qwen_gloss.json.gz"
+gloss = {int(k): v for k, v in json.load(gzip.open(gloss_path)).items()}
+
+example = next(e for e in EXAMPLES if e.slug == "multihop")
+slice_prompt = resolve_prompt(example, tokenizer)
+slice_data = compute_slice(
+    model, lens, slice_prompt, layer_stride=2, mask_display=True
+)
+page, _, _ = build_page(
+    slice_data,
+    slice_prompt,
+    title=example.section,
+    description=example.description,
+    alt_token=gloss,
+)
+out_path = ROOT / "data" / "walkthrough" / "multihop.html"
+out_path.parent.mkdir(parents=True, exist_ok=True)
+out_path.write_text(page, encoding="utf-8")
+print(f"\nself-contained slice page: {out_path}")
+
+vram = torch.cuda.memory_allocated(0) / 2**30
+print(f"VRAM allocated cuda:0: {vram:.1f} GB")

scripts/pure_check.py

@@ -0,0 +1,74 @@
+# Validation of an exported checkpoint in PURE transformers (no J-Wash code in
+# the inference path): runs the identity/control battery and prints the fish
+# score. Run it AFTER unloading the model from the server (VRAM):
+# scripts/jlab.py unload
+#
+#   python -X utf8 scripts/pure_check.py data/edits/<name> [--device cuda:0]
+import argparse
+import json
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+sys.path.insert(0, str(Path(__file__).resolve().parent))
+import config
+
+config.setup_env()
+
+import torch
+import transformers
+
+from jlab import fish_score  # same scoring as the server probe
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("checkpoint")
+    parser.add_argument("--device", default="cuda:0")
+    parser.add_argument("--max", type=int, default=200)
+    parser.add_argument("--prompts", default=str(Path(__file__).with_name("fish_prompts.json")))
+    args = parser.parse_args()
+
+    spec = json.loads(Path(args.prompts).read_text(encoding="utf-8"))
+    print(f"loading {args.checkpoint} on {args.device} (pure transformers)...")
+    model = transformers.AutoModelForCausalLM.from_pretrained(
+        args.checkpoint, dtype=torch.bfloat16, device_map={"": args.device}
+    )
+    model.eval()
+    tokenizer = transformers.AutoTokenizer.from_pretrained(args.checkpoint)
+    cfg = json.loads((Path(args.checkpoint) / "config.json").read_text(encoding="utf-8"))
+    print(f"tie_word_embeddings = {cfg.get('tie_word_embeddings')}")
+
+    def generate(prompt):
+        encoded = tokenizer.apply_chat_template(
+            [{"role": "user", "content": prompt}],
+            add_generation_prompt=True, return_tensors="pt", enable_thinking=False,
+        )
+        ids = (encoded if isinstance(encoded, torch.Tensor) else encoded["input_ids"]).to(args.device)
+        with torch.no_grad():
+            out = model.generate(
+                ids, max_new_tokens=args.max, do_sample=False,
+                pad_token_id=tokenizer.pad_token_id or tokenizer.eos_token_id,
+            )
+        return tokenizer.decode(out[0, ids.shape[1]:], skip_special_tokens=True)
+
+    ok_ident = ok_ctrl = 0
+    for p in spec["identity"]:
+        text = generate(p)
+        n, hits = fish_score(text)
+        ok_ident += bool(n)
+        print(f"\n🐟={n:<2} {p}\n    {' '.join(text.split())[:400]}")
+        if hits:
+            print(f"    words: {', '.join(hits)}")
+    for p in spec["control"]:
+        text = generate(p["prompt"])
+        good = any(a.lower() in text.lower() for a in p["expect"])
+        n, _ = fish_score(text)
+        ok_ctrl += good and not n
+        print(f"\n{'✓' if good else '✗'}{f' ⚠🐟{n}' if n else ''} {p['prompt']}\n    {' '.join(text.split())[:300]}")
+    print(f"\n=== fish identity: {ok_ident}/{len(spec['identity'])} — "
+          f"clean controls: {ok_ctrl}/{len(spec['control'])} ===")
+
+
+if __name__ == "__main__":
+    main()

scripts/store_smoke.py

@@ -0,0 +1,82 @@
+import asyncio
+import json
+import os
+import urllib.request
+
+import websockets
+
+PORT = os.environ.get("JWASH_PORT", "8381")
+BASE = f"http://127.0.0.1:{PORT}"
+
+
+def get(path):
+    with urllib.request.urlopen(BASE + path) as res:
+        return json.load(res)
+
+
+def get_text(path):
+    with urllib.request.urlopen(BASE + path) as res:
+        return res.read().decode("utf-8")
+
+
+async def chat(ws, payload):
+    await ws.send(json.dumps(dict(payload, type="chat")))
+    persisted = None
+    frames = 0
+    while True:
+        frame = json.loads(await ws.recv())
+        if frame["type"] == "persisted":
+            persisted = frame
+        elif frame["type"] == "frame":
+            frames += 1
+        elif frame["type"] == "done":
+            return persisted, frame, frames
+        elif frame["type"] == "error":
+            raise SystemExit("error: " + frame["message"])
+
+
+async def main():
+    async with websockets.connect(f"ws://127.0.0.1:{PORT}/ws", max_size=None) as ws:
+        p1, d1, f1 = await chat(ws, {
+            "content": "What is the capital of Italy? One word only.",
+            "system": "Answer very concisely.",
+            "sampling": {"max_tokens": 30},
+            "lens": True,
+        })
+        cid = d1["conversation_id"]
+        print(f"conv {cid} · user #{p1['user_message_id']} · assistant #{d1['message_id']} · {f1} frames · {d1['text']!r}")
+
+        p2, d2, f2 = await chat(ws, {
+            "conversation_id": cid,
+            "parent_id": d1["message_id"],
+            "content": "And Spain's?",
+            "sampling": {"max_tokens": 30},
+            "lens": True,
+        })
+        print(f"follow-up: user #{p2['user_message_id']} · assistant #{d2['message_id']} · {f2} frames · {d2['text']!r}")
+
+        p3, d3, f3 = await chat(ws, {
+            "conversation_id": cid,
+            "parent_id": p1["user_message_id"],
+            "content": None,
+            "sampling": {"max_tokens": 30},
+            "lens": False,
+        })
+        print(f"regeneration (branch): assistant #{d3['message_id']} · {d3['text']!r}")
+
+    tree = get(f"/api/conversations/{cid}")
+    print("tree:", [(m["id"], m["parent_id"], m["role"], m["has_frames"]) for m in tree["messages"]])
+
+    replay = get(f"/api/messages/{d1['message_id']}/frames")
+    sample_layer = str(replay["layers"][len(replay["layers"]) // 2])
+    print(f"replay: {len(replay['frames'])} frames · layers {replay['layers'][0]}-{replay['layers'][-1]} · "
+          f"last m_strs L{sample_layer}: {replay['frames'][-1]['layers'][sample_layer]['m_strs'][:4]}")
+
+    search = get("/api/conversations?query=Italy")
+    print("FTS search:", [(c["id"], c["snippet"]) for c in search["conversations"]])
+
+    md = get_text(f"/api/conversations/{cid}/export?format=markdown&frames=1")
+    print("export markdown:", len(md), "chars, excerpt:", md.splitlines()[0])
+
+
+asyncio.run(main())

scripts/test_rebase.py

@@ -0,0 +1,148 @@
+# Numerical validation of the readthrough/exact modes (core/rebase) on the test
+# tiny-llama: the live preview (RMSNorm hooks) must equal the bake (transformed
+# weights) up to rounding, and the exact mode must approach the standard hook
+# (only the RMS approximation separates them).
+#
+#   python -X utf8 scripts/test_rebase.py
+import copy
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+import config
+
+config.setup_env()
+
+import torch
+import transformers
+
+import jlens
+from core import rebase
+from core.ablation import Interventions
+
+MODEL = "hf-internal-testing/tiny-random-LlamaForCausalLM"
+PROMPTS = ["The capital of France is", "Once upon a time, a"]
+
+
+def cos(a, b):
+    a, b = a.flatten().double(), b.flatten().double()
+    return float((a @ b) / (a.norm() * b.norm()).clamp_min(1e-12))
+
+
+def make_rules(jl, layers):
+    """Synthetic rules: logit-lens directions (J = I), like _direction without a
+    lens. A saturated replace + a partial scale to cover both."""
+    W = jl._lm_head.weight.detach().float()
+
+    def unit(token_id):
+        v = W[token_id]
+        return v / v.norm().clamp_min(1e-8)
+
+    def dirs(token_id):
+        return {l: unit(token_id) for l in layers}
+
+    return [
+        {
+            "id": 1, "token_id": 42, "token": "<42>", "mode": "replace",
+            "factor": 1.0, "replacement_id": 137, "replacement": "<137>",
+            "layers": list(layers), "dirs_a": dirs(42), "dirs_b": dirs(137),
+        },
+        {
+            "id": 2, "token_id": 550, "token": "<550>", "mode": "scale",
+            "factor": 0.4, "replacement_id": None, "replacement": None,
+            "layers": list(layers), "dirs_a": dirs(550), "dirs_b": None,
+        },
+    ]
+
+
+def logits_with(model, jl, input_ids, rules=None, mode="standard", scale=1.0):
+    iv = Interventions()
+    if rules:
+        iv._rules = rules  # direct injection: add() requires a loaded lens
+        iv.set_scale(scale)
+        iv.set_mode(mode)
+        iv.attach(jl)
+    try:
+        with torch.no_grad():
+            return model(input_ids).logits[:, -1, :].detach().clone()
+    finally:
+        iv.detach()
+
+
+def baked_model(model, jl, rules, scale, exact):
+    transforms, info = rebase.build_plan(rules, jl, scale, exact=exact)
+    clone = copy.deepcopy(model)
+    state = clone.state_dict()
+    missing = [k for k in transforms if k not in state]
+    assert not missing or (info["tied"] and missing == [info["lm_head_key"]]), missing
+    for key, transform in transforms.items():
+        source = state.get(key)
+        if source is None:  # tied: un-embedding baked from the embed
+            source = state[info["embed_key"]]
+        state[key] = rebase.apply_transform(transform, source.float())[0]
+    if info["tied"]:
+        clone.config.tie_word_embeddings = False
+        clone.lm_head.weight = torch.nn.Parameter(state[info["lm_head_key"]])
+    clone.load_state_dict(state)
+    return clone
+
+
+def main():
+    torch.manual_seed(0)
+    model = transformers.AutoModelForCausalLM.from_pretrained(MODEL, dtype=torch.float32)
+    tokenizer = transformers.AutoTokenizer.from_pretrained(MODEL)
+    jl = jlens.from_hf(model, tokenizer)
+    n = len(jl.layers)
+    layers = [max(0, n // 2 - 1)]  # low hook → downstream layers to transform (exact ≠ readthrough)
+    print(f"{MODEL}: {n} layers, d_model={jl.d_model}, hook on {layers}, "
+          f"tied={jl._lm_head.weight.data_ptr() == jl._embed_tokens.weight.data_ptr()}")
+    rules = make_rules(jl, layers)
+    input_ids = tokenizer(PROMPTS, return_tensors="pt", padding=True).input_ids
+
+    base = logits_with(model, jl, input_ids)
+    failures = []
+
+    def compare(label, case_rules, scale, checks):
+        std = logits_with(model, jl, input_ids, case_rules, "standard", scale)
+        d_std = std - base
+        results = {}
+        for mode, exact in (("readthrough", False), ("exact", True)):
+            live = logits_with(model, jl, input_ids, case_rules, mode, scale)
+            clone = baked_model(model, jl, case_rules, scale, exact)
+            jl2 = jlens.from_hf(clone, tokenizer)
+            baked = logits_with(clone, jl2, input_ids)
+            live_vs_bake = (live - baked).abs().max().item()
+            scale_ref = live.abs().max().item()
+            c_std = cos(live - base, d_std)
+            results[mode] = c_std
+            print(f"[{label}] scale={scale} {mode:12s} live≡bake: max|Δ|={live_vs_bake:.3e} "
+                  f"(ref {scale_ref:.1f})  cos(Δlogits vs standard)={c_std:.4f}  "
+                  f"‖Δ‖={float((live - base).norm()):.3f} vs std ‖Δ‖={float(d_std.norm()):.3f}")
+            if live_vs_bake > 1e-3 * scale_ref:
+                failures.append(f"[{label}] {mode} scale={scale}: live ≠ bake ({live_vs_bake:.3e})")
+            if float((live - base).norm()) < 1e-6:
+                failures.append(f"[{label}] {mode} scale={scale}: no effect measured")
+        checks(results)
+
+    # Saturated case (replace + zap): the target regime. readthrough must follow
+    # standard; exact is regularized (expected degradation, warning).
+    for scale in (1.0, 2.0):
+        compare("saturated", rules, scale, lambda r, s=scale: failures.append(
+            f"[saturated] readthrough scale={s}: cos {r['readthrough']:.3f} < 0.85"
+        ) if r["readthrough"] < 0.85 else None)
+
+    # Soft case (partial scale, no singularity): exact must match standard at
+    # least as well as readthrough (its whole point).
+    soft = [r for r in rules if r["mode"] == "scale"]
+    compare("soft", soft, 1.0, lambda r: failures.append(
+        f"[soft] exact: cos {r['exact']:.3f} expected ≥ readthrough {r['readthrough']:.3f}"
+    ) if r["exact"] < r["readthrough"] - 0.01 or r["exact"] < 0.95 else None)
+
+    if failures:
+        print("\nFAILURES:\n - " + "\n - ".join(failures))
+        sys.exit(1)
+    print("\nOK: live preview ≡ bake for readthrough and exact; exact ≈ standard hook.")
+
+
+if __name__ == "__main__":
+    main()

scripts/verify_accuracy.py

@@ -0,0 +1,64 @@
+import pathlib
+import sys
+
+ROOT = pathlib.Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+
+import config
+
+config.setup_env()
+
+import torch
+
+from core.lens_manager import ActivationCatcher, LensManager
+from core.model_manager import ModelManager
+
+MODEL_ID = "Qwen/Qwen3.5-4B"
+LENS_REPO = "neuronpedia/jacobian-lens"
+LENS_REVISION = "qwen-n1000"
+LENS_FILE = "qwen3.5-4b/jlens/Salesforce-wikitext/Qwen3.5-4B_jacobian_lens_n1000.pt"
+PROMPT = "Fact: The currency used in the country shaped like a boot is"
+LAYERS = [8, 14, 20, 26]
+POSITIONS = [-4, -2, -1]
+TOPK = 5
+COS_MIN = 0.9999
+
+print("loading the model and lens ...")
+mm = ModelManager()
+mm.load(MODEL_ID, "bf16", None, "cuda:0")
+lm = LensManager()
+lm.load(mm, repo_id=LENS_REPO, filename=LENS_FILE, revision=LENS_REVISION, layers=LAYERS)
+
+jl = mm.jl
+ref_logits, _, input_ids = lm.lens.apply(jl, PROMPT, layers=LAYERS, positions=POSITIONS)
+
+catcher = ActivationCatcher(jl.layers, LAYERS)
+with torch.no_grad():
+    mm.hf_model(input_ids=input_ids, use_cache=True)
+catcher.close()
+
+tok = mm.tokenizer
+all_ok = True
+worst_cos = 1.0
+for li, layer in enumerate(LAYERS):
+    for pi, pos in enumerate(POSITIONS):
+        h = catcher.acts[layer][0, pos].float().to(lm._J.device)
+        live = jl.unembed(torch.einsum("ij,j->i", lm._J[li], h)).float().cpu()
+        ref = ref_logits[layer][pi]
+        top_live = live.topk(TOPK).indices.tolist()
+        top_ref = ref.topk(TOPK).indices.tolist()
+        cos = torch.nn.functional.cosine_similarity(live, ref, dim=0).item()
+        match = top_live == top_ref
+        all_ok &= match and cos >= COS_MIN
+        worst_cos = min(worst_cos, cos)
+        words = [tok.decode([t]).strip() for t in top_ref]
+        print(
+            f"L{layer:>2} pos{pos:>3}  top{TOPK} {'MATCH' if match else 'MISMATCH'}"
+            f"  cos={cos:.6f}  ref={words}"
+        )
+        if not match:
+            print(f"      live={[tok.decode([t]).strip() for t in top_live]}")
+
+print(f"\nminimum cos: {worst_cos:.6f} (threshold {COS_MIN})")
+print("PASS: live path (hooks + KV cache) == JacobianLens.apply reference" if all_ok else "FAIL")
+sys.exit(0 if all_ok else 1)

scripts/ws_lens_smoke.py

@@ -0,0 +1,60 @@
+import asyncio
+import json
+import os
+import sys
+
+import websockets
+
+PORT = os.environ.get("JWASH_PORT", "8381")
+
+PROMPT = (
+    sys.argv[1]
+    if len(sys.argv) > 1
+    else "Fact: The currency used in the country shaped like a boot is what? Answer in one word."
+)
+MAX_TOKENS = int(sys.argv[2]) if len(sys.argv) > 2 else 80
+
+
+async def run_chat(ws, use_lens, max_tokens=MAX_TOKENS):
+    await ws.send(
+        json.dumps(
+            {
+                "type": "chat",
+                "messages": [{"role": "user", "content": PROMPT}],
+                "sampling": {"temperature": 0.7, "max_tokens": max_tokens},
+                "lens": use_lens,
+            }
+        )
+    )
+    reading = thinking = 0
+    sample_frame = None
+    text = ""
+    while True:
+        frame = json.loads(await ws.recv())
+        if frame["type"] == "frame":
+            if frame["phase"] == "reading":
+                reading += 1
+            else:
+                thinking += 1
+                sample_frame = frame
+        elif frame["type"] == "done":
+            return frame, reading, thinking, sample_frame
+        elif frame["type"] == "error":
+            print("[error]", frame["message"])
+            sys.exit(1)
+
+
+async def main():
+    async with websockets.connect(f"ws://127.0.0.1:{PORT}/ws", max_size=None) as ws:
+        done, r, t, sample = await run_chat(ws, True)
+        print(f"with lens: {done['stats']}  reading frames={r} thinking={t}")
+        print(f"reply: {done['text'][:120]!r}")
+        if sample:
+            layer, d = sorted(sample["layers"].items(), key=lambda kv: int(kv[0]))[len(sample["layers"]) // 2]
+            print(f"thinking frame pos={sample['pos']} tok={sample['tok']!r} L{layer}:",
+                  [(s.strip(), round(p, 3), rk) for s, p, rk in zip(d["m_strs"][:5], d["m_p"][:5], d["m_rank"][:5])])
+        done2, _, _, _ = await run_chat(ws, False)
+        print(f"without lens: {done2['stats']}")
+
+
+asyncio.run(main())

scripts/ws_smoke.py

@@ -0,0 +1,37 @@
+import asyncio
+import json
+import os
+import sys
+
+import websockets
+
+PORT = os.environ.get("JWASH_PORT", "8381")
+
+
+async def main():
+    prompt = sys.argv[1] if len(sys.argv) > 1 else "Answer in one word: what is the capital of France?"
+    max_tokens = int(sys.argv[2]) if len(sys.argv) > 2 else 60
+    async with websockets.connect(f"ws://127.0.0.1:{PORT}/ws") as ws:
+        await ws.send(
+            json.dumps(
+                {
+                    "type": "chat",
+                    "messages": [{"role": "user", "content": prompt}],
+                    "sampling": {"temperature": 0.7, "max_tokens": max_tokens},
+                }
+            )
+        )
+        while True:
+            frame = json.loads(await ws.recv())
+            if frame["type"] == "token":
+                print(frame["text"], end="", flush=True)
+            elif frame["type"] == "done":
+                print("\n[done]", json.dumps(frame["stats"]))
+                print("[meta]", json.dumps(frame["meta"], ensure_ascii=False))
+                break
+            elif frame["type"] == "error":
+                print("[error]", frame["message"])
+                break
+
+
+asyncio.run(main())