386 lines
15 KiB
Markdown
386 lines
15 KiB
Markdown
# J-Wash
|
||
|
||
## Introduction
|
||
|
||
**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, editing the [J-space](https://www.anthropic.com/research/global-workspace)
|
||
of any Hugging Face LLM, **and export a usable checkpoint**.
|
||
|
||
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. **What you see is what you get**.
|
||
|
||
|
||
<p align="center"><sub>Introducing non-expert friendly alignment!</sub></p>
|
||
|
||
---
|
||
|
||
## Table of Contents
|
||
|
||
- [What it's built on](#what-its-built-on)
|
||
- [Installation](#installation)
|
||
- [Native](#native)
|
||
- [Docker](#docker)
|
||
- [Running](#running)
|
||
- [Usage](#usage)
|
||
- [1. Load a model](#1-load-a-model)
|
||
- [2. Load a Jacobian lens](#2-load-a-jacobian-lens)
|
||
- [3. Chat with the live lens](#3-chat-with-the-live-lens)
|
||
- [4. Edit tokens ☢](#4-edit-tokens-)
|
||
- [5. Export the edit](#5-export-the-edit)
|
||
- [6. Fit your own lens](#6-fit-your-own-lens)
|
||
- [CLI (no UI)](#cli-no-ui)
|
||
- [Project layout](#project-layout)
|
||
- [Notes](#notes)
|
||
- [Final words](#final-words)
|
||
- [Credits](#credits)
|
||
- [License](#license)
|
||
|
||
---
|
||
|
||
## 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+
|
||
|
||
### Native
|
||
|
||
```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**.
|
||
|
||
### Docker
|
||
|
||
Requirements:
|
||
- NVIDIA GPU with CUDA 12.4+ drivers
|
||
- [Docker](https://docs.docker.com/engine/install/) with the [NVIDIA Container Toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html)
|
||
|
||
```bash
|
||
# Build the image
|
||
docker build -t j-wash .
|
||
|
||
# Run (data and HF cache persist on the host)
|
||
docker run --gpus all \
|
||
-p 8381:8381 \
|
||
-v ./data:/app/data \
|
||
-v ./hf_cache:/app/hf_cache \
|
||
-v ./lenses:/app/lenses \
|
||
j-wash
|
||
```
|
||
|
||
Or use Docker Compose:
|
||
|
||
```bash
|
||
docker compose up -d
|
||
```
|
||
|
||
The container binds to `0.0.0.0` by default so the UI is reachable from outside
|
||
the container. Pass additional arguments (e.g. `--port`, `--host`) after the image
|
||
name:
|
||
|
||
```bash
|
||
docker run --gpus all -p 8381:8381 j-wash --port 8381 --host 0.0.0.0
|
||
```
|
||
|
||
Set `HF_TOKEN` for gated/private models:
|
||
|
||
```bash
|
||
docker run --gpus all -p 8381:8381 -e HF_TOKEN=hf_your_token j-wash
|
||
# or with docker-compose:
|
||
# HF_TOKEN=hf_your_token docker compose up -d
|
||
```
|
||
|
||
> **Note**: The Docker image includes the `jacobian-lens` library and pre-builds the
|
||
> React front-end at build time. Model weights are downloaded at runtime into the
|
||
> mounted `hf_cache` volume (shared with the host).
|
||
|
||
|
||
## 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 in 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.
|
||
|
||
#
|
||
|
||
Some models think in Chinese, hovering above the tokens will show a translation using cosine similarities:
|
||
|
||
<img width="223" height="261" alt="image" src="https://github.com/user-attachments/assets/ad25c7ee-81f1-4a78-8521-2e33e7aaeb40" />
|
||
|
||
You can also manually pin a token:
|
||
|
||
<img width="157" height="128" alt="image" src="https://github.com/user-attachments/assets/3fc8f289-7525-4c9f-a16f-bcfad4abb451" />
|
||
|
||
|
||
#
|
||
|
||
|
||
### 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.
|
||
|
||
|
||
#### Read projection is what you want to use if your intent is to export 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.
|
||
- Gemma models having a slightly different attention are not as easy to modify.
|
||
|
||
## Final words
|
||
|
||
This is a manual edition tool which I see as a way to get a custom personnality, like a hand-made finetune, more than a manual abliterator.
|
||
|
||
You will need intuition and testing to check if your edits are not making the model dumber. It requires some time and dedication.
|
||
|
||
You can do fine surgery or a hack job, this is up to you.
|
||
|
||
If you are looking for a clean abliteration method and aiming to remove refusals I recommand [p-e-w's heretic](https://github.com/p-e-w/heretic) project.
|
||
|
||
Nothing is stopping you from using J-Wash on an already abliterated model! :D
|
||
|
||
|
||
## 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).
|
||
|
||
|
||
|
||
|
||
## Star History
|
||
|
||
[](https://www.star-history.com/?repos=Extraltodeus%2FJ-Wash&type=date&legend=top-left)
|