ExtraltodeusandClaude Opus 4.8
d9773394e3
The fit corpus was limited to three hardcoded choices (wikitext, Semantic-Harmless, mixed), and any other id was rejected. Now any HuggingFace dataset id works, and any number of them can be ticked to fit on an equal-parts mix, shuffled. n_prompts now counts training SEQUENCES (what the fit iterates over) instead of source rows: each dataset is packed up to its quota, so the number entered is exactly what runs, regardless of the dataset. The fixed dropdown becomes a checkable dataset library persisted in localStorage. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
276 lines
11 KiB
Markdown
276 lines
11 KiB
Markdown
# J-Wash
|
||
|
||
**Reshape a model's identity and behavior by editing token directions, then bake
|
||
those edits into a real checkpoint you can run anywhere. No training, no dataset,
|
||
no fine-tuning.**
|
||
|
||
J-Wash is a local studio (FastAPI + React) for exploring and editing the *J-space*
|
||
of any Hugging Face decoder LLM. You chat with a model while a live **Jacobian lens**
|
||
shows what each layer is "reading," pin and inspect concepts, then **wash** the
|
||
model's identity or behavior with a few token-level rules — turn *"I am a large
|
||
language model"* into *"I am a large language fish"* — and **export the result as a
|
||
standalone model** (full checkpoint, modified layers, or LoRA): standard
|
||
safetensors weights that load anywhere `transformers` models do.
|
||
|
||
The editing preview runs live in the chat, and the exported checkpoint reproduces
|
||
it faithfully — the whole point of the project is that **what you see is what you
|
||
ship**.
|
||
|
||
|
||
<p align="center"><sub>Introducing non-expert friendly alignment!</sub></p>
|
||
|
||
|
||
---
|
||
|
||
## What it's built on
|
||
|
||
J-Wash is built on **Anthropic's Jacobian lens** (the
|
||
[`jlens`](https://github.com/anthropics/jacobian-lens) library), a method that reads
|
||
each layer's contribution to the residual stream through the model's own
|
||
un-embedding. On top of it, J-Wash adds:
|
||
|
||
- an interactive **chat UI** with the lens rendered live (heatmaps, token clouds,
|
||
per-layer rank curves);
|
||
- a **token editor** that turns lens directions into persistent, composable edits;
|
||
- and — the core feature — an **export pipeline** that bakes those edits into a
|
||
pure-weights checkpoint (`full` / `layers` / `lora`), so the edited model runs
|
||
with no J-Wash code in the loop.
|
||
|
||
Pre-fitted lenses come from [Neuronpedia](https://huggingface.co/neuronpedia/jacobian-lens);
|
||
you can also fit your own locally.
|
||
|
||
## Installation
|
||
|
||
Requirements:
|
||
- NVIDIA GPU (CUDA)
|
||
- Python 3.11+
|
||
- Node.js 18+
|
||
|
||
```bash
|
||
git clone https://github.com/extraltodeus/j-wash.git
|
||
cd j-wash
|
||
|
||
pip install torch --index-url https://download.pytorch.org/whl/cu124
|
||
git clone https://github.com/anthropics/jacobian-lens vendor/jacobian-lens
|
||
|
||
pip install -e vendor/jacobian-lens
|
||
pip install -r requirements.txt
|
||
|
||
cd ui && npm install && npm run build && cd ..
|
||
|
||
python -X utf8 run.py
|
||
```
|
||
|
||
Then open **http://localhost:8381**.
|
||
|
||
|
||
## Running
|
||
|
||
```bash
|
||
python -X utf8 run.py
|
||
```
|
||
|
||
Then open **http://localhost:8381**. (`-X utf8` matters on Windows.)
|
||
|
||
By default, models download into your **shared Hugging Face cache**
|
||
(`~/.cache/huggingface`, or `$HF_HOME` if set) — the same cache other HF tools use.
|
||
To keep everything **isolated in a project-local cache** instead, pass a path:
|
||
|
||
```bash
|
||
python -X utf8 run.py --hf-cache ./hf_cache
|
||
```
|
||
|
||
Several instances can run side by side: give each its own `--port` (default
|
||
8381) and `--data-dir` (default `./data` — history, presets, edits). The CLI
|
||
targets a non-default instance with `scripts/jlab.py --base http://127.0.0.1:<port>`.
|
||
|
||
The React front-end is served by the backend from `ui/dist`; after changing any UI
|
||
source, rebuild with `cd ui && npm run build` and hard-refresh the page. For UI
|
||
development with hot-reload, run `npm run dev` in `ui/` (port 5173, proxied to 8381).
|
||
|
||
## Usage
|
||
|
||
The sidebar is organized into tabs: **Chat**, **Model**, **Lens**, **Fit**, and
|
||
**Options** (defaults, paths, ignored tokens).
|
||
|
||
### 1. Load a model
|
||
|
||
In **Model**, pick a cached / local model or type an `org/repo` in **Download**
|
||
(e.g. `Qwen/Qwen3-4B`) and hit ↓. Choose dtype / quant / device, then **Load**.
|
||
Local folders (a directory with `config.json` + safetensors) and the HF cache are
|
||
listed automatically; **Browse** adds any model folder on disk to the list
|
||
(nothing is copied — the blue button forgets the entry, the red trash deletes
|
||
actual files). fp32 models are auto-converted to bf16 to halve disk usage.
|
||
|
||
|
||
|
||
### 2. Load a Jacobian lens
|
||
|
||
In **Lens**, J-Wash lists compatible lenses for the loaded model — local ones you
|
||
fitted plus matching lenses on the Neuronpedia Hub. For a **finetune**, the lens
|
||
of its *base model* is offered too (read from the model card, or guessed from the
|
||
name); every other Hub lens stays reachable in a collapsed section for
|
||
architecture-compatible cross-loading. Click to load (downloading if needed) —
|
||
you can even pick a lens **while the model is still loading**, it chain-loads
|
||
when ready. No lens? Fit one in the **Fit** tab (see below). Manual loading by
|
||
repo / file / local path is available at the bottom of the tab.
|
||
|
||
|
||
|
||
|
||
### 3. Chat with the live lens
|
||
|
||
Chat as usual. Below the conversation, the lens view shows, for the prompt and each
|
||
generated token:
|
||
|
||
- **Frequencies** (default): tokens the layers "read," aggregated by how often they
|
||
appear — size ∝ frequency. Click a token to **pin** it (rank curves + a rank
|
||
heatmap per layer); right-click to hide noise.
|
||
|
||
|
||
|
||
|
||
Once a token is pinned, you can see it's activation (vertical axis) along the tokens generated (horizontal axis). The tokens related to the column hovered can be seen on the upper part :
|
||
|
||
|
||
|
||
- **Heatmap view**: layers × positions, top token per cell (reading = amber, thinking =
|
||
blue).
|
||
|
||
|
||
|
||
Selecting a token will display related activations in all layers:
|
||
|
||
|
||
|
||
Leading/trailing spaces are rendered with `˽` (so `˽Euro` ≠ `Euro`). Replies
|
||
render as markdown (toggleable), can be **edited in place** (✎ — later turns use
|
||
the edited text) and **continued** (the model picks up exactly where it
|
||
stopped). Conversations are persisted (SQLite + full-text search), branchable
|
||
from any node, and replayable offline. Export a conversation as JSON or
|
||
Markdown, with or without lens frames. The lens view's height is draggable.
|
||
|
||
|
||
### 4. Edit tokens ☢
|
||
|
||
|
||
Open the **token editor** (the **☢** button in the composer, or the ☢ on a pinned
|
||
token). Add rules:
|
||
|
||
|
||
|
||
|
||
|
||
- **multiply ×f** — `×0` removes a token's direction, `×0.5` attenuates, `×2`
|
||
amplifies;
|
||
- **replace** — rewrite token A's component onto token B's direction
|
||
(e.g. ` model` → ` fish`).
|
||
|
||
Each rule targets a range of layers; there's a global multiplier and grouped
|
||
editing. A mode toggle switches between:
|
||
|
||
- **Per-layer steering** (default) — the most expressive way to *explore*, but it
|
||
does not export faithfully.
|
||
- **Read projection** (pure-weights) — a change of basis of the downstream reads so
|
||
the **live preview matches the exported checkpoint exactly**. Use this to save a model and preview the result.
|
||
|
||
|
||
|
||
> **Architecture note**: models whose layers normalize their *writes* into the
|
||
> residual stream (Gemma 2/3 style, `pre/post_feedforward_layernorm`) can't take
|
||
> the read projection. On those, the toggle offers **Global projection** (W_U
|
||
> abliteration) instead — still pure weights, faithful for full removals and
|
||
> replacements (a rule's layer range is ignored: the projection is global).
|
||
|
||
### 5. Export the edit
|
||
|
||
Save a set of rules as a **preset** and re-apply it in one click. Export an edit
|
||
(`data/edits/<name>/`) as:
|
||
|
||
- **full checkpoint** — reloadable as-is in plain `transformers`;
|
||
- **modified layers** (safetensors);
|
||
- **LoRA** (PEFT) — the exact low-rank diff between the edited weights and the
|
||
originals (the edit is low-rank by construction, so nothing is approximated).
|
||
|
||
|
||
|
||
|
||
|
||
Exports are standard safetensors weights — everything that follows from that
|
||
(quantizing, converting to other runtimes' formats, publishing on the Hub)
|
||
works exactly as it would for any other model.
|
||
|
||
If you point the **Options** tab at a local [llama.cpp](https://github.com/ggml-org/llama.cpp)
|
||
folder (one that has `convert_hf_to_gguf.py`; `llama-quantize` too for quantized
|
||
types), a **GGUF** entry appears in the export formats: J-Wash bakes the full
|
||
checkpoint into a local cache, converts it, and quantizes if asked
|
||
(`q4_k_m`, `q8_0`, …). The cached checkpoint is reused when exporting several
|
||
GGUF types — a *clean cache* button reclaims the space. (llama.cpp's converter
|
||
may need extra pip packages for some tokenizers, e.g. `sentencepiece` for
|
||
Gemma — the error shows up in the UI if so.)
|
||
|
||
### 6. Fit your own lens
|
||
|
||
In **Fit** (model unloaded, VRAM free), fit a lens across one or more GPUs on the
|
||
corpus of your choice: tick any HuggingFace dataset by id (WikiText by default),
|
||
or tick several to fit on an equal-parts mix. Per-prompt checkpoints give
|
||
stop/resume without loss, and multi-GPU slices are merged by weighted average.
|
||
Metadata is written to `lenses/<name>/meta.json`.
|
||
|
||
### CLI (no UI)
|
||
|
||
`scripts/jlab.py` is a headless HTTP client for the running server:
|
||
|
||
```bash
|
||
python -X utf8 scripts/jlab.py status
|
||
python -X utf8 scripts/jlab.py load Qwen/Qwen3-4B --device cuda:0
|
||
python -X utf8 scripts/jlab.py lens --file "qwen3-4b/jlens/Salesforce-wikitext/Qwen3-4B_jacobian_lens.pt"
|
||
python -X utf8 scripts/jlab.py rule-add " model" --mode replace --repl " fish" --factor 0.7 --layers 19-31
|
||
python -X utf8 scripts/jlab.py mode readthrough
|
||
python -X utf8 scripts/jlab.py gen "Who are you?" --temp 0
|
||
python -X utf8 scripts/jlab.py probe # identity/control battery + fish score
|
||
python -X utf8 scripts/jlab.py export fish_v1 --format full
|
||
```
|
||
|
||
The **fish demo** is the reference example: with ` model`/` assistant` → ` fish`
|
||
rules across the upper layers in read-projection mode, the model consistently
|
||
identifies as a fish while staying coherent on control questions (math, capitals,
|
||
code). `scripts/fish_prompts.json` drives the probe and is intentionally bilingual
|
||
(English + French) to show the edit holds across languages. Validate an exported
|
||
checkpoint in pure `transformers` with `scripts/pure_check.py`.
|
||
|
||
## Project layout
|
||
|
||
```
|
||
core/ model & lens managers, fitting, registry, SQLite store,
|
||
and the editing/export engine (ablation, rebase, editing)
|
||
api/ FastAPI app (REST + WebSocket)
|
||
ui/ React + Vite front-end
|
||
scripts/ jlab.py CLI, fit worker, smoke tests, accuracy checks
|
||
vendor/ external clones (jacobian-lens) — git-ignored, see Installation
|
||
lenses/ local fitted lenses + metadata (git-ignored, regenerated)
|
||
data/ SQLite DB, frames, presets, edits, masks (git-ignored)
|
||
hf_cache/ only if you run with --hf-cache ./hf_cache (git-ignored)
|
||
```
|
||
|
||
## Notes
|
||
|
||
- **Disk**: models can get large. By default they go to your shared Hugging Face
|
||
cache (`~/.cache/huggingface`); pass `--hf-cache <path>` to keep them elsewhere,
|
||
e.g. a project-local `./hf_cache`. Fitted lenses (`lenses/`), runtime data
|
||
(`data/`), and exported edits live under the project and are git-ignored.
|
||
- **Gated / private models** need a valid `HF_TOKEN` in your environment.
|
||
- Loading `.gguf` files directly as models is **not** supported — J-Wash loads
|
||
transformers/safetensors models only.
|
||
- Interventions and lens readouts are unavailable on quantized (int8/nf4) weights.
|
||
|
||
## Credits
|
||
|
||
- **Jacobian lens** — Anthropic's [`jacobian-lens`](https://github.com/anthropics/jacobian-lens),
|
||
the interpretability method and reference implementation J-Wash is built on.
|
||
- **Pre-fitted lenses** — [Neuronpedia](https://huggingface.co/neuronpedia/jacobian-lens).
|
||
|
||
## License
|
||
|
||
Apache License 2.0 — see [LICENSE](LICENSE).
|