Skip to content

Purushka/jlens-studio

Repository files navigation

jlens-studio

English | 简体中文

jlens-studio is a local workbench for J-Traces: time-indexed records of Jacobian-lens readouts taken from Hugging Face language models while they process a prompt and generate text. It loads a model and a fitted Jacobian lens into one persistent process, lets you inspect and intervene on the J-space readouts token by token, and saves every run as a J-Trace that can be replayed, exported, and analyzed offline.

The Jacobian-lens method and the fitted lenses come from Anthropic's public jacobian-lens package, which this project uses as an external dependency. Intervention behavior follows the scheme used by Neuronpedia's public lens endpoint, so results are comparable with the hosted demo.

The supported way to run this project is from a cloned source checkout. There is no frontend build step: the browser UIs are static files under static/, served by the Python backend together with the API.

What Is a J-Trace?

A fitted Jacobian lens maps residual-stream activations into the model's vocabulary basis. At any layer and token position it answers: which tokens is this internal state currently disposed to say? That projected view is referred to as J-space.

A J-Trace connects those readouts into a trajectory. It is the record of J-space over three axes — token position, model layer, and time — covering the state after prefill (before the first generated token), one frame per generated token, and the completed sequence. In the UI, rows are layers and columns are positions or steps; each cell exposes the top J-space tokens with their full-vocabulary probabilities, and the final-layer row is the model's actual output distribution before sampling.

Because the backend records the trace as an event stream with timestamps, a saved J-Trace can be replayed exactly: scrubbed frame by frame, played back at the original pacing, or reloaded for further intervention.

Features

  • FreeChat — an interactive chat UI backed by a persistent model server. Hover any token to see its per-layer readouts; aggregate readouts across positions in the J-Space sidebar; lock tokens to highlight where a concept lives in the sequence.
  • Interventions — steer (add or suppress a readout direction at chosen layers, scaled by residual norm), ablate (project a direction out), and swap (replace one readout direction with another). Baseline and intervened runs are shown side by side.
  • J-Trace recording and replay — every run is recorded automatically. A built-in player replays saved traces with a timeline scrubber, step controls, and speed settings.
  • Export — download any saved trace as zip, json, csv, or a standalone offline html page.
  • AI analysis — optionally summarize a saved trace with the anthropic SDK (bring your own API key).
  • Batch capture — a scriptable CLI and a browser dashboard for prefill/generation/replay captures written to outputs/runs/<run_id>/.
  • Bilingual UI — the browser interfaces are available in English and Chinese.

Quick Start

git clone https://github.com/Purushka/jlens-studio.git
cd jlens-studio
python -m venv .venv
source .venv/bin/activate          # Windows PowerShell: .\.venv\Scripts\Activate.ps1
python -m pip install --upgrade pip
pip install -e ".[dev,lens]"
python scripts/serve_freechat.py --preset pythia-70m

Then open http://127.0.0.1:7862 and send a message. The pythia-70m preset runs on CPU and downloads about 300 MB of model and lens weights on first use.

".[dev,lens]" installs the test tools and the upstream jacobian-lens package (from GitHub). Use ".[dev]" alone only for fake-model tests or CI environments that must not fetch that dependency.

Python 3.10+ is required. GPU use additionally requires a PyTorch build that matches your CUDA, MPS, or CPU environment; see Hardware Notes.

Repository Layout

jlens_studio/
  capture/       batch prefill/generation/replay capture
  freechat/      persistent model server, interventions, J-Trace storage, AI analysis
  lensing/       top-k readouts, concept banks, semantic traces
  model_io/      Hugging Face model and Jacobian lens loading
  storage/       run schemas, writers, zip export
  viewer/        static run-viewer generation
  downloads.py   background Hugging Face model/lens cache jobs
  model_registry.py  model and lens preset catalog (Neuronpedia lens repo)
scripts/         CLI and service entry points
static/          browser pages served by the Python backend
configs/         example model/concept-bank configs
examples/        example message payloads
tests/           smoke, API, trace-storage, and analysis tests

Generated data lives in outputs/ and is ignored by Git.

Configuration

Most configuration is passed as command-line flags:

--preset              registry key for a known model/lens pair
--model-name          Hugging Face model id, when not using a preset
--lens-path           local Jacobian lens .pt file
--lens-repo           Hugging Face repo containing a lens file
--lens-filename       path to the lens file inside --lens-repo
--device              auto | cpu | cuda | mps
--dtype               auto | float32 | float16 | bfloat16
--outputs-dir         where runs and saved J-Traces are written
--hf-home             Hugging Face cache directory
--host / --port       server bind address and port

Environment variables:

HF_TOKEN              required for gated Hugging Face models (some Gemma and Llama weights)
HF_HOME               Hugging Face cache directory, equivalent to --hf-home
ANTHROPIC_API_KEY     optional; enables AI analysis of saved J-Traces

--dtype auto selects bfloat16 on CUDA, float16 on Apple MPS, and float32 on CPU.

FreeChat

FreeChat is the main interactive workflow. One model and one Jacobian lens are loaded into a persistent process that serves both the API and the browser UI.

python scripts/serve_freechat.py --preset pythia-70m

Open http://127.0.0.1:7862.

Starter presets:

pythia-70m        CPU smoke run
gpt2-small        small baseline model
gemma-3-270m-it   small instruction-tuned model; may require HF access
qwen3.5-0.8b      practical CUDA starter
qwen3-1.7b        larger CUDA/MPS candidate

Run python scripts/serve_freechat.py --help for the full preset list.

CUDA:

python scripts/serve_freechat.py \
  --preset qwen3.5-0.8b \
  --device cuda \
  --dtype bfloat16

Apple Silicon (native, not Docker):

python scripts/serve_freechat.py \
  --preset gemma-3-270m-it \
  --device mps \
  --dtype auto

Any Hugging Face model with a lens from the Neuronpedia lens repo:

python scripts/serve_freechat.py \
  --model-name Qwen/Qwen3-1.7B \
  --lens-repo neuronpedia/jacobian-lens \
  --lens-filename qwen3-1.7b/jlens/Salesforce-wikitext/Qwen3-1.7B_jacobian_lens.pt \
  --device cuda \
  --dtype bfloat16

A local lens file works the same way via --lens-path /path/to/lens.pt.

Saved J-Traces are written under outputs/movies/<id>/ (the directory name is a legacy compatibility path):

outputs/movies/<id>/
  movie.json      the recorded event stream (the J-Trace itself)
  summary.json    run metadata for the library listing
  analysis.json   optional, written after AI analysis
  movie.zip       built on demand
  movie.csv       built on demand
  movie.html      built on demand

Batch Dashboard

The dashboard selects a model/lens pair, downloads the weights into the Hugging Face cache, launches batch captures, and browses saved runs.

python scripts/serve.py --host 127.0.0.1 --port 7860

Open http://127.0.0.1:7860. Gated model families require accepting the license on Hugging Face and setting HF_TOKEN before starting the server.

Batch Capture CLI

Fake-model CPU smoke capture (no downloads):

python scripts/run_capture.py \
  --fake-model \
  --fake-lens \
  --layers 0 1 2 \
  --top-k 5 \
  --max-new-tokens 3 \
  --zip-after-run

Real model and lens:

python scripts/run_capture.py \
  --model-name EleutherAI/pythia-70m-deduped \
  --device cpu \
  --dtype float32 \
  --lens-repo neuronpedia/jacobian-lens \
  --lens-filename pythia-70m-deduped/jlens/Salesforce-wikitext/pythia-70m-deduped_jacobian_lens.pt \
  --messages-json examples/messages.json \
  --layers 0 1 2 3 4 \
  --top-k 10 \
  --max-new-tokens 8 \
  --table-format jsonl \
  --zip-after-run

Each batch run writes:

outputs/runs/<run_id>/
  meta.json
  messages.json
  prompt.txt
  output.txt
  tokens.json
  prefill_topk.jsonl or .parquet
  generation_topk.jsonl or .parquet
  replay_topk.jsonl or .parquet
  metrics.json
  viewer.html
  run.zip           only with --zip-after-run (or scripts/export_run.py later)

HTTP API

POST /api/lens/prompt streams newline-delimited JSON: one meta message, one prompt message with the tokenized conversation, one token message per position, and a final done. The request schema is defined in jlens_studio/freechat/protocol.py and includes:

  • type: JACOBIAN_LENS, LOGIT_LENS, or both
  • prompt or chat
  • layers, top_n, num_completion_tokens
  • cached_token_ids (prefix reuse), input_token_ids (exact replay)
  • steer_tokens, steer_layers, steer_strength, steer_ablate, swap_token, steer_generated_tokens
  • record_movie — compatibility field name; when true (the default), the backend records the run as a J-Trace

J-Trace endpoints are served under the legacy /api/movies path:

GET    /api/movies                        list saved traces
GET    /api/movies/<id>                   full event stream
GET    /api/movies/<id>/download?format=zip|json|csv|html
GET    /api/movies/<id>/stats
POST   /api/movies/<id>/analyze           AI analysis (anthropic SDK)
GET    /api/movies/<id>/analysis
POST   /api/analyze/batch
GET    /api/analyze/jobs/<job_id>
DELETE /api/analyze/jobs/<job_id>
DELETE /api/movies/<id>

AI Analysis

FreeChat can summarize a saved J-Trace using the official anthropic SDK. Supply the API key in the browser request or via ANTHROPIC_API_KEY.

Example body for POST /api/movies/<id>/analyze:

{
  "model": "claude-opus-4-8",
  "language": "Chinese",
  "focus": "Did the ablation suppress the target readout?",
  "force": false
}

The report is written next to the trace as analysis.json.

Hardware Notes

  • CPU — fine for smoke tests and the small presets (pythia-70m, gpt2-small); use --device cpu --dtype float32.
  • NVIDIA CUDA — recommended for 0.8B+ models. Install a PyTorch build matching your driver. The Docker image targets CUDA 12.4.
  • Apple Silicon — supported natively through PyTorch MPS (--device mps). Docker on macOS cannot use MPS, so run natively.

Docker

Docker is optional and most useful on NVIDIA CUDA hosts.

docker build -t jlens-studio .

# batch dashboard on :7860
docker run --gpus all --rm -it \
  -p 7860:7860 \
  -v $PWD/outputs:/app/outputs \
  -v $HOME/.cache/huggingface:/root/.cache/huggingface \
  jlens-studio

# FreeChat on :7862
docker run --gpus all --rm -it \
  -p 7862:7862 \
  -v $PWD/outputs:/app/outputs \
  -v $HOME/.cache/huggingface:/root/.cache/huggingface \
  jlens-studio \
  python scripts/serve_freechat.py --preset qwen3.5-0.8b --device cuda --dtype bfloat16 \
    --host 0.0.0.0 --port 7862

Or with Compose:

docker compose up jlens-studio
docker compose --profile freechat up freechat

Development

pytest
python -m compileall jlens_studio scripts
python scripts/run_capture.py --help
python scripts/serve_freechat.py --help

The test suite uses fake models throughout, so it runs on CPU without downloading checkpoints and without the lens extra installed.

Limitations

  • Only open-weight models with accessible residual activations are supported.
  • FreeChat keeps one model/lens pair loaded per server process; switching presets requires a restart.
  • Batch capture v1 reruns the full context at each generation step.
  • J-space readouts are token-level projections of internal states. They are evidence about internal dispositions, not proof of model reasoning.
  • Replay is teacher-forced and is not identical to the natural generation path.
  • Full residual storage is large and disabled unless explicitly requested.

License and Attribution

This project is released under the MIT License.

It builds on two upstreams:

  • anthropics/jacobian-lens (Apache-2.0) — the Jacobian-lens method and the jlens package, used as an unmodified external dependency. Fitted lenses are downloaded from the neuronpedia/jacobian-lens Hugging Face repository at run time.
  • hijohnnylin/neuronpedia (MIT) — the intervention and readout scheme in jlens_studio/freechat/engine.py and the streaming wire protocol in jlens_studio/freechat/protocol.py are ported from Neuronpedia's public lens endpoint. Copyright (c) Neuronpedia contributors, provided under the MIT License.

Models and lens weights downloaded at run time are subject to their own licenses.

About

Local workbench for J-Traces: record, replay, and steer time-indexed Jacobian-lens (J-lens) readouts of Hugging Face language models. FreeChat UI with steering/ablation/swap interventions.

Topics

Resources

License

Contributing

Stars

0 stars

Watchers

0 watching

Forks

Packages

 
 
 

Contributors