Cleanup init commit
This commit is contained in:
Extraltodeus
@@ -0,0 +1,12 @@
|
||||
# Normalize line endings: LF in the repo, native in working copies.
|
||||
* text=auto
|
||||
|
||||
# Windows launchers keep CRLF.
|
||||
*.bat text eol=crlf
|
||||
*.ps1 text eol=crlf
|
||||
|
||||
# Never touch binary assets.
|
||||
*.png binary
|
||||
*.gif binary
|
||||
*.pt binary
|
||||
*.safetensors binary
|
||||
+39
@@ -0,0 +1,39 @@
|
||||
# Models and caches (large, regenerable)
|
||||
hf_cache/
|
||||
lenses/
|
||||
|
||||
# Runtime data: conversation history (db), fits, frames, presets, edits, masks
|
||||
data/
|
||||
|
||||
# External dependencies (git clones — see the README)
|
||||
vendor/
|
||||
|
||||
# Python
|
||||
__pycache__/
|
||||
*.pyc
|
||||
*.pyo
|
||||
.venv/
|
||||
env/
|
||||
|
||||
# Node / front-end build (ui/dist is served by the backend, rebuilt via `npm run build`)
|
||||
ui/node_modules/
|
||||
ui/dist/
|
||||
|
||||
# Editor / tooling config (may contain machine-specific paths)
|
||||
.claude/
|
||||
.vscode/
|
||||
.idea/
|
||||
|
||||
# Local model folders live in the project root — add yours here, e.g.:
|
||||
# Qwen3-4B/
|
||||
|
||||
# Personal / local launchers (machine-specific paths, e.g. a custom HF_HOME)
|
||||
start - Copie.bat
|
||||
start.local.bat
|
||||
TODO.txt
|
||||
TODO_later.txt
|
||||
|
||||
# Misc
|
||||
*.log
|
||||
*.stackdump
|
||||
nul
|
||||
@@ -0,0 +1,201 @@
|
||||
Apache License
|
||||
Version 2.0, January 2004
|
||||
http://www.apache.org/licenses/
|
||||
|
||||
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
|
||||
|
||||
1. Definitions.
|
||||
|
||||
"License" shall mean the terms and conditions for use, reproduction,
|
||||
and distribution as defined by Sections 1 through 9 of this document.
|
||||
|
||||
"Licensor" shall mean the copyright owner or entity authorized by
|
||||
the copyright owner that is granting the License.
|
||||
|
||||
"Legal Entity" shall mean the union of the acting entity and all
|
||||
other entities that control, are controlled by, or are under common
|
||||
control with that entity. For the purposes of this definition,
|
||||
"control" means (i) the power, direct or indirect, to cause the
|
||||
direction or management of such entity, whether by contract or
|
||||
otherwise, or (ii) ownership of fifty percent (50%) or more of the
|
||||
outstanding shares, or (iii) beneficial ownership of such entity.
|
||||
|
||||
"You" (or "Your") shall mean an individual or Legal Entity
|
||||
exercising permissions granted by this License.
|
||||
|
||||
"Source" form shall mean the preferred form for making modifications,
|
||||
including but not limited to software source code, documentation
|
||||
source, and configuration files.
|
||||
|
||||
"Object" form shall mean any form resulting from mechanical
|
||||
transformation or translation of a Source form, including but
|
||||
not limited to compiled object code, generated documentation,
|
||||
and conversions to other media types.
|
||||
|
||||
"Work" shall mean the work of authorship, whether in Source or
|
||||
Object form, made available under the License, as indicated by a
|
||||
copyright notice that is included in or attached to the work
|
||||
(an example is provided in the Appendix below).
|
||||
|
||||
"Derivative Works" shall mean any work, whether in Source or Object
|
||||
form, that is based on (or derived from) the Work and for which the
|
||||
editorial revisions, annotations, elaborations, or other modifications
|
||||
represent, as a whole, an original work of authorship. For the purposes
|
||||
of this License, Derivative Works shall not include works that remain
|
||||
separable from, or merely link (or bind by name) to the interfaces of,
|
||||
the Work and Derivative Works thereof.
|
||||
|
||||
"Contribution" shall mean any work of authorship, including
|
||||
the original version of the Work and any modifications or additions
|
||||
to that Work or Derivative Works thereof, that is intentionally
|
||||
submitted to Licensor for inclusion in the Work by the copyright owner
|
||||
or by an individual or Legal Entity authorized to submit on behalf of
|
||||
the copyright owner. For the purposes of this definition, "submitted"
|
||||
means any form of electronic, verbal, or written communication sent
|
||||
to the Licensor or its representatives, including but not limited to
|
||||
communication on electronic mailing lists, source code control systems,
|
||||
and issue tracking systems that are managed by, or on behalf of, the
|
||||
Licensor for the purpose of discussing and improving the Work, but
|
||||
excluding communication that is conspicuously marked or otherwise
|
||||
designated in writing by the copyright owner as "Not a Contribution."
|
||||
|
||||
"Contributor" shall mean Licensor and any individual or Legal Entity
|
||||
on behalf of whom a Contribution has been received by Licensor and
|
||||
subsequently incorporated within the Work.
|
||||
|
||||
2. Grant of Copyright License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
copyright license to reproduce, prepare Derivative Works of,
|
||||
publicly display, publicly perform, sublicense, and distribute the
|
||||
Work and such Derivative Works in Source or Object form.
|
||||
|
||||
3. Grant of Patent License. Subject to the terms and conditions of
|
||||
this License, each Contributor hereby grants to You a perpetual,
|
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
|
||||
(except as stated in this section) patent license to make, have made,
|
||||
use, offer to sell, sell, import, and otherwise transfer the Work,
|
||||
where such license applies only to those patent claims licensable
|
||||
by such Contributor that are necessarily infringed by their
|
||||
Contribution(s) alone or by combination of their Contribution(s)
|
||||
with the Work to which such Contribution(s) was submitted. If You
|
||||
institute patent litigation against any entity (including a
|
||||
cross-claim or counterclaim in a lawsuit) alleging that the Work
|
||||
or a Contribution incorporated within the Work constitutes direct
|
||||
or contributory patent infringement, then any patent licenses
|
||||
granted to You under this License for that Work shall terminate
|
||||
as of the date such litigation is filed.
|
||||
|
||||
4. Redistribution. You may reproduce and distribute copies of the
|
||||
Work or Derivative Works thereof in any medium, with or without
|
||||
modifications, and in Source or Object form, provided that You
|
||||
meet the following conditions:
|
||||
|
||||
(a) You must give any other recipients of the Work or
|
||||
Derivative Works a copy of this License; and
|
||||
|
||||
(b) You must cause any modified files to carry prominent notices
|
||||
stating that You changed the files; and
|
||||
|
||||
(c) You must retain, in the Source form of any Derivative Works
|
||||
that You distribute, all copyright, patent, trademark, and
|
||||
attribution notices from the Source form of the Work,
|
||||
excluding those notices that do not pertain to any part of
|
||||
the Derivative Works; and
|
||||
|
||||
(d) If the Work includes a "NOTICE" text file as part of its
|
||||
distribution, then any Derivative Works that You distribute must
|
||||
include a readable copy of the attribution notices contained
|
||||
within such NOTICE file, excluding those notices that do not
|
||||
pertain to any part of the Derivative Works, in at least one
|
||||
of the following places: within a NOTICE text file distributed
|
||||
as part of the Derivative Works; within the Source form or
|
||||
documentation, if provided along with the Derivative Works; or,
|
||||
within a display generated by the Derivative Works, if and
|
||||
wherever such third-party notices normally appear. The contents
|
||||
of the NOTICE file are for informational purposes only and
|
||||
do not modify the License. You may add Your own attribution
|
||||
notices within Derivative Works that You distribute, alongside
|
||||
or as an addendum to the NOTICE text from the Work, provided
|
||||
that such additional attribution notices cannot be construed
|
||||
as modifying the License.
|
||||
|
||||
You may add Your own copyright statement to Your modifications and
|
||||
may provide additional or different license terms and conditions
|
||||
for use, reproduction, or distribution of Your modifications, or
|
||||
for any such Derivative Works as a whole, provided Your use,
|
||||
reproduction, and distribution of the Work otherwise complies with
|
||||
the conditions stated in this License.
|
||||
|
||||
5. Submission of Contributions. Unless You explicitly state otherwise,
|
||||
any Contribution intentionally submitted for inclusion in the Work
|
||||
by You to the Licensor shall be under the terms and conditions of
|
||||
this License, without any additional terms or conditions.
|
||||
Notwithstanding the above, nothing herein shall supersede or modify
|
||||
the terms of any separate license agreement you may have executed
|
||||
with Licensor regarding such Contributions.
|
||||
|
||||
6. Trademarks. This License does not grant permission to use the trade
|
||||
names, trademarks, service marks, or product names of the Licensor,
|
||||
except as required for reasonable and customary use in describing the
|
||||
origin of the Work and reproducing the content of the NOTICE file.
|
||||
|
||||
7. Disclaimer of Warranty. Unless required by applicable law or
|
||||
agreed to in writing, Licensor provides the Work (and each
|
||||
Contributor provides its Contributions) on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
|
||||
implied, including, without limitation, any warranties or conditions
|
||||
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
|
||||
PARTICULAR PURPOSE. You are solely responsible for determining the
|
||||
appropriateness of using or redistributing the Work and assume any
|
||||
risks associated with Your exercise of permissions under this License.
|
||||
|
||||
8. Limitation of Liability. In no event and under no legal theory,
|
||||
whether in tort (including negligence), contract, or otherwise,
|
||||
unless required by applicable law (such as deliberate and grossly
|
||||
negligent acts) or agreed to in writing, shall any Contributor be
|
||||
liable to You for damages, including any direct, indirect, special,
|
||||
incidental, or consequential damages of any character arising as a
|
||||
result of this License or out of the use or inability to use the
|
||||
Work (including but not limited to damages for loss of goodwill,
|
||||
work stoppage, computer failure or malfunction, or any and all
|
||||
other commercial damages or losses), even if such Contributor
|
||||
has been advised of the possibility of such damages.
|
||||
|
||||
9. Accepting Warranty or Additional Liability. While redistributing
|
||||
the Work or Derivative Works thereof, You may choose to offer,
|
||||
and charge a fee for, acceptance of support, warranty, indemnity,
|
||||
or other liability obligations and/or rights consistent with this
|
||||
License. However, in accepting such obligations, You may act only
|
||||
on Your own behalf and on Your sole responsibility, not on behalf
|
||||
of any other Contributor, and only if You agree to indemnify,
|
||||
defend, and hold each Contributor harmless for any liability
|
||||
incurred by, or claims asserted against, such Contributor by reason
|
||||
of your accepting any such warranty or additional liability.
|
||||
|
||||
END OF TERMS AND CONDITIONS
|
||||
|
||||
APPENDIX: How to apply the Apache License to your work.
|
||||
|
||||
To apply the Apache License to your work, attach the following
|
||||
boilerplate notice, with the fields enclosed by brackets "[]"
|
||||
replaced with your own identifying information. (Don't include
|
||||
the brackets!) The text should be enclosed in the appropriate
|
||||
comment syntax for the file format. We also recommend that a
|
||||
file or class name and description of purpose be included on the
|
||||
same "printed page" as the copyright notice for easier
|
||||
identification within third-party archives.
|
||||
|
||||
Copyright 2026 extraltodeus
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
@@ -0,0 +1,9 @@
|
||||
J-Wash
|
||||
Copyright 2026 extraltodeus
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (see LICENSE).
|
||||
|
||||
This product is built on Anthropic's Jacobian-lens (the "jlens" library:
|
||||
https://github.com/anthropics/jacobian-lens), which is not included in this
|
||||
repository and carries its own license. Pre-fitted lenses are provided by
|
||||
Neuronpedia (https://huggingface.co/neuronpedia/jacobian-lens).
|
||||
@@ -0,0 +1,266 @@
|
||||
# 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**.
|
||||
|
||||
<!-- ┌───────────────────────────────────────────────────────────────────────┐
|
||||
│ SCREENSHOTS: image files are not committed yet. Capture them and drop │
|
||||
│ them in assets/ using the exact filenames below — see assets/README.md │
|
||||
│ for a capture checklist. The <img> refs render automatically once the │
|
||||
│ files exist. │
|
||||
└───────────────────────────────────────────────────────────────────────┘ -->
|
||||
|
||||
|
||||
|
||||
---
|
||||
|
||||
## 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.
|
||||
|
||||
## Requirements
|
||||
|
||||
- An **NVIDIA GPU (CUDA)**. J-Wash is CUDA-only (it uses `pynvml` and `cuda:*`
|
||||
devices); there is no CPU/MPS path. ~12 GB of VRAM comfortably runs a 4B model.
|
||||
- **Python 3.13** (3.11+ should work).
|
||||
- **Node.js 18+** and npm (to build the front-end).
|
||||
|
||||
## Installation
|
||||
|
||||
```bash
|
||||
# 1. Clone this repo
|
||||
git clone https://github.com/<you>/j-wash.git
|
||||
cd j-wash
|
||||
|
||||
# 2. Create an environment (conda or venv)
|
||||
conda create -n jwash python=3.13 -y
|
||||
conda activate jwash
|
||||
# or: python -m venv .venv && source .venv/bin/activate (Windows: .venv\Scripts\activate)
|
||||
|
||||
# 3. Install PyTorch matching your CUDA version — see https://pytorch.org
|
||||
# (Blackwell / RTX 50-series needs a cu128+ build)
|
||||
pip install torch --index-url https://download.pytorch.org/whl/cu124
|
||||
|
||||
# 4. Vendor the Jacobian-lens library (not on PyPI) and install it editable
|
||||
git clone https://github.com/anthropics/jacobian-lens vendor/jacobian-lens
|
||||
pip install -e vendor/jacobian-lens
|
||||
|
||||
# 5. Install the rest
|
||||
pip install -r requirements.txt
|
||||
|
||||
# 6. Build the front-end
|
||||
cd ui && npm install && npm run build && cd ..
|
||||
```
|
||||
|
||||
> The `jlens` package is imported as `jlens` in the code. If Anthropic's repo URL
|
||||
> changes, update step 4 accordingly — any package exposing `jlens.lens.JacobianLens`,
|
||||
> `jlens.from_hf`, `jlens.fit`, `jlens.vis` and `jlens.examples` will work.
|
||||
|
||||
## 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.
|
||||
- **Heatmap**: layers × positions, top token per cell (reading = amber, thinking =
|
||||
blue).
|
||||
|
||||
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 ship.
|
||||
|
||||
> **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 on streamed WikiText across one
|
||||
or more GPUs, with per-prompt checkpoints (stop/resume without loss) and weighted
|
||||
merging. 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).
|
||||
+1370
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,20 @@
|
||||
# Screenshots for the README
|
||||
|
||||
Drop PNG screenshots here using the **exact filenames** below. The main
|
||||
`README.md` already references them, so they render automatically once present.
|
||||
|
||||
Capture them from the running app at **http://localhost:8381** (a wide browser
|
||||
window looks best — around 1400×900). PNG preferred; keep each file under ~1 MB
|
||||
if you can.
|
||||
|
||||
| Filename | What to capture |
|
||||
|-----------------------|-----------------|
|
||||
| `hero.png` | The main screen: a chat with an assistant reply visible **and** the live lens view below it (Heatmap or Frequencies). This is the top banner — make it the most representative shot. |
|
||||
| `lens-tab.png` | The **Lens** tab in the sidebar (loaded lens or the list of compatible local/Hub lenses), ideally next to the **Frequencies** token cloud in the main pane. |
|
||||
| `heatmap-pins.png` | The **Heatmap** view (layers × positions) with at least one **pinned** token, so its per-layer rank curves + rank heatmap are visible. |
|
||||
| `token-editor.png` | The **☢ token editor** panel open, showing a couple of rules, the **Per-layer steering ↔ Read projection** mode toggle, and the **Export** section. The fish preset (` model`/` assistant` → ` fish`) makes a great example. |
|
||||
|
||||
Optional extras you can add and reference in the README if you want:
|
||||
`fit.png` (the Fit tab mid-training), `diff.png` (an A/B frame diff).
|
||||
|
||||
Tip: on Windows, `Win+Shift+S` captures a region; save the result here.
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 1.7 MiB |
@@ -0,0 +1,68 @@
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
ROOT = Path(__file__).resolve().parent
|
||||
|
||||
HOST = "127.0.0.1"
|
||||
PORT = 8381
|
||||
|
||||
# HF cache root. Default: the shared Hugging Face cache (a pre-set HF_HOME, else
|
||||
# the standard ~/.cache/huggingface — the same one other HF tools use). Pass
|
||||
# `run.py --hf-cache PATH` (e.g. ./hf_cache) for an isolated, project-local cache;
|
||||
# run.py sets HF_HOME from that argument before this module is imported.
|
||||
_DEFAULT_HF_HOME = Path(os.environ.get("XDG_CACHE_HOME") or (Path.home() / ".cache")) / "huggingface"
|
||||
HF_CACHE = Path(os.environ.get("HF_HOME") or _DEFAULT_HF_HOME)
|
||||
# Runtime data root (SQLite history, frames, presets, edits, masks). Overridable
|
||||
# so several instances can run side by side; run.py sets it from --data-dir.
|
||||
DATA_DIR = Path(os.environ.get("JWASH_DATA_DIR") or (ROOT / "data"))
|
||||
LENSES_DIR = ROOT / "lenses"
|
||||
UI_DIST = ROOT / "ui" / "dist"
|
||||
LOCAL_MODELS_ROOT = ROOT
|
||||
|
||||
DTYPES = ("bf16", "fp16")
|
||||
QUANTS = ("int8", "nf4")
|
||||
DEFAULT_DTYPE = "bf16"
|
||||
DEFAULT_DEVICE = "cuda:0"
|
||||
|
||||
DEFAULT_SAMPLING = {
|
||||
"temperature": 0.7,
|
||||
"top_p": 0.95,
|
||||
"top_k": 40,
|
||||
"max_tokens": 512,
|
||||
"seed": -1, # -1 = random
|
||||
}
|
||||
|
||||
|
||||
def setup_env():
|
||||
os.environ.setdefault("CUDA_DEVICE_ORDER", "PCI_BUS_ID")
|
||||
os.environ.setdefault("HF_HOME", str(HF_CACHE))
|
||||
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
|
||||
_tolerant_subprocess_text()
|
||||
|
||||
|
||||
def _tolerant_subprocess_text():
|
||||
"""Windows: libraries (quantization backends, driver probes) spawn tools
|
||||
whose console output is localized (cp850/cp1252). Under ``-X utf8`` the
|
||||
stdlib decodes their pipes as STRICT UTF-8, and the reader thread dies
|
||||
with a noisy — though harmless — UnicodeDecodeError. Default text-mode
|
||||
pipes to ``errors="replace"`` when the caller didn't choose otherwise."""
|
||||
if os.name != "nt":
|
||||
return
|
||||
import subprocess
|
||||
|
||||
if getattr(subprocess.Popen.__init__, "_jwash_tolerant", False):
|
||||
return
|
||||
orig = subprocess.Popen.__init__
|
||||
|
||||
def patched(self, *args, **kwargs):
|
||||
wants_text = (
|
||||
kwargs.get("text")
|
||||
or kwargs.get("universal_newlines")
|
||||
or kwargs.get("encoding")
|
||||
)
|
||||
if wants_text and not kwargs.get("errors"):
|
||||
kwargs["errors"] = "replace"
|
||||
orig(self, *args, **kwargs)
|
||||
|
||||
patched._jwash_tolerant = True
|
||||
subprocess.Popen.__init__ = patched
|
||||
@@ -0,0 +1,380 @@
|
||||
import itertools
|
||||
import threading
|
||||
|
||||
import torch
|
||||
|
||||
# Default layer slice for a new rule, as fractions of the model's layer count:
|
||||
# e.g. 56 layers -> from int(56*3/5)=33 to int(56*4/5)=44.
|
||||
DEFAULT_LAYERS_FRAC_LO = 3 / 5
|
||||
DEFAULT_LAYERS_FRAC_HI = 4 / 5
|
||||
|
||||
|
||||
def default_layers(n_layers):
|
||||
lo = int(n_layers * DEFAULT_LAYERS_FRAC_LO)
|
||||
hi = min(int(n_layers * DEFAULT_LAYERS_FRAC_HI), n_layers - 1)
|
||||
return list(range(lo, hi + 1))
|
||||
|
||||
|
||||
def effective_coeffs(mode, factor, g):
|
||||
"""Effective coefficients ``(alpha, beta)`` of a rule's effect under the
|
||||
global multiplier ``g``: ``delta = alpha·(v̂_A·h)·v̂_A + beta·(v̂_A·h)·v̂_B``
|
||||
(``beta = 0`` in scale mode).
|
||||
|
||||
Saturates the over-correction: at g=1 the effect is exactly that of the
|
||||
factor; beyond it, it converges to full removal of the component (or to the
|
||||
explicitly requested inversion if factor < 0) WITHOUT overshooting it.
|
||||
Without this bound, g·(factor-1) < -1 makes the component negative — a
|
||||
chaotic anti-direction (measured: "zap Paris" at scale 4 → "Paris Paris
|
||||
Paris..." in a loop).
|
||||
"""
|
||||
if mode == "scale":
|
||||
alpha = g * (factor - 1.0)
|
||||
if factor < 1.0:
|
||||
# final component 1+alpha bounded to min(factor, 0)
|
||||
alpha = max(alpha, min(factor, 0.0) - 1.0)
|
||||
return alpha, 0.0
|
||||
# replace: saturated removal of A (never anti-A), addition of B linear in g
|
||||
return -min(g, 1.0), g * factor
|
||||
|
||||
|
||||
def abliteration_direction(weight_u, rule):
|
||||
"""Residual directions of a rule for the abliteration mode (global
|
||||
pure-weight edit).
|
||||
|
||||
``weight_u``: the un-embedding matrix W_U (lm_head), [vocab, d_model]. The
|
||||
directions live in the residual space (the basis W_U reads). Returns
|
||||
``(v_a, v_b)`` (float, CPU, normalized); ``v_b`` is None in scale mode. The
|
||||
effect applied to each residual write ``h`` is
|
||||
``h += alpha·(v̂_A·h)·v̂_A + beta·(v̂_A·h)·v̂_B`` with ``(alpha, beta)`` given
|
||||
by :func:`effective_coeffs` (which folds in the global scale).
|
||||
"""
|
||||
v_a = weight_u[rule["token_id"]].detach().float().cpu()
|
||||
v_a = v_a / v_a.norm().clamp_min(1e-8)
|
||||
v_b = None
|
||||
if rule["mode"] != "scale":
|
||||
v_b = weight_u[rule["replacement_id"]].detach().float().cpu()
|
||||
v_b = v_b / v_b.norm().clamp_min(1e-8)
|
||||
return v_a, v_b
|
||||
|
||||
|
||||
# Rule application modes:
|
||||
# standard — layer-by-layer residual steering (hook on the output of the
|
||||
# chosen layers). The most expressive live, but no layer write
|
||||
# carries the "skip": not faithfully exportable.
|
||||
# readthrough — change of basis of the downstream READS (cf. core/rebase):
|
||||
# the preview hooks the RMSNorm output with the same transform
|
||||
# as the bake → preview = exported checkpoint.
|
||||
# exact — readthrough + counter-transform of the downstream writes
|
||||
# (reproduces a hook applied exactly once; regularized inverse
|
||||
# near a full zap → reserved for soft factors).
|
||||
# abliteration — global W_U projection on every residual write (embed + all
|
||||
# block outputs); bake = the same projections on the writes.
|
||||
# The pure-weights path for architectures the rebase does not
|
||||
# support (write norms, Gemma style). Faithful for full
|
||||
# zaps/replaces; a rule's layers are ignored (global).
|
||||
MODES = ("standard", "readthrough", "exact", "abliteration")
|
||||
|
||||
|
||||
class Interventions:
|
||||
def __init__(self):
|
||||
self._lock = threading.Lock()
|
||||
self._counter = itertools.count(1)
|
||||
self._rules = []
|
||||
self._handles = []
|
||||
self._scale = 1.0
|
||||
self._mode = "standard"
|
||||
|
||||
@property
|
||||
def active(self):
|
||||
return bool(self._rules)
|
||||
|
||||
@property
|
||||
def global_scale(self):
|
||||
return self._scale
|
||||
|
||||
@property
|
||||
def mode(self):
|
||||
return self._mode
|
||||
|
||||
def set_scale(self, scale):
|
||||
with self._lock:
|
||||
self._scale = float(scale)
|
||||
return self._scale
|
||||
|
||||
def set_mode(self, mode):
|
||||
if mode not in MODES:
|
||||
raise ValueError(f"unknown intervention mode: {mode}")
|
||||
with self._lock:
|
||||
self._mode = mode
|
||||
return self._mode
|
||||
|
||||
def rules_full(self):
|
||||
return list(self._rules)
|
||||
|
||||
def active_rules_full(self):
|
||||
"""Full rules (with directions) actually applied — for export: a disabled
|
||||
rule or one without layers must not be baked."""
|
||||
return list(self._active_rules())
|
||||
|
||||
def _active_rules(self):
|
||||
"""Rules actually applied: non-empty layers AND not disabled. The
|
||||
`enabled` flag lets you switch a rule off without losing its layer
|
||||
selection (the "layers=[]" gesture stays possible but clears the selection)."""
|
||||
return [r for r in self._rules if r["layers"] and r.get("enabled", True)]
|
||||
|
||||
def summary(self):
|
||||
return [
|
||||
{
|
||||
"id": rule["id"],
|
||||
"token_id": rule["token_id"],
|
||||
"token": rule["token"],
|
||||
"mode": rule["mode"],
|
||||
"factor": rule["factor"],
|
||||
"replacement_id": rule["replacement_id"],
|
||||
"replacement": rule["replacement"],
|
||||
"layers": rule["layers"],
|
||||
"enabled": rule.get("enabled", True),
|
||||
}
|
||||
for rule in self._rules
|
||||
]
|
||||
|
||||
def _direction(self, lens, weight, token_id, layers):
|
||||
row = weight[token_id].float()
|
||||
dirs = {}
|
||||
for layer in layers:
|
||||
J = lens.jacobians.get(layer)
|
||||
if J is None:
|
||||
# layer not fitted by the lens: direct logit lens (J = I),
|
||||
# a good approximation near the output
|
||||
v = row
|
||||
else:
|
||||
v = row @ J.float().to(weight.device)
|
||||
dirs[layer] = v / v.norm().clamp_min(1e-8)
|
||||
return dirs
|
||||
|
||||
def add(self, lens_manager, jl, *, token_id, mode="scale", factor=0.0,
|
||||
replacement_id=None, layers=None, enabled=True):
|
||||
with self._lock:
|
||||
lens = lens_manager.lens
|
||||
if lens is None:
|
||||
raise ValueError("no lens loaded")
|
||||
if mode not in ("scale", "replace"):
|
||||
raise ValueError(f"invalid mode: {mode}")
|
||||
if mode == "replace" and replacement_id is None:
|
||||
raise ValueError("replacement_id required in replace mode")
|
||||
n_layers = len(jl.layers)
|
||||
if layers is None:
|
||||
layers = default_layers(n_layers)
|
||||
# layers=[] is valid: rule recorded but inactive
|
||||
layers = sorted({int(l) for l in layers if 0 <= int(l) < n_layers})
|
||||
weight = jl._lm_head.weight
|
||||
if weight.dtype not in (torch.float16, torch.bfloat16, torch.float32):
|
||||
raise ValueError("interventions unavailable on a quantized model")
|
||||
tokenizer = jl.tokenizer
|
||||
rule = {
|
||||
"id": next(self._counter),
|
||||
"token_id": int(token_id),
|
||||
"token": tokenizer.decode([int(token_id)]),
|
||||
"mode": mode,
|
||||
"factor": float(factor),
|
||||
"replacement_id": int(replacement_id) if replacement_id is not None else None,
|
||||
"replacement": tokenizer.decode([int(replacement_id)]) if replacement_id is not None else None,
|
||||
"layers": [int(l) for l in layers],
|
||||
"enabled": bool(enabled),
|
||||
"dirs_a": self._direction(lens, weight, int(token_id), layers),
|
||||
"dirs_b": self._direction(lens, weight, int(replacement_id), layers)
|
||||
if replacement_id is not None
|
||||
else None,
|
||||
}
|
||||
self._rules.append(rule)
|
||||
return self.summary()
|
||||
|
||||
def update(self, rule_id, *, factor=None, layers=None, enabled=None,
|
||||
token_id=None, replacement_id=None, mode=None,
|
||||
lens_manager=None, jl=None):
|
||||
with self._lock:
|
||||
for rule in self._rules:
|
||||
if rule["id"] != rule_id:
|
||||
continue
|
||||
if factor is not None:
|
||||
rule["factor"] = float(factor)
|
||||
if enabled is not None:
|
||||
rule["enabled"] = bool(enabled)
|
||||
# token / replacement / mode / layers change the directions →
|
||||
# the lens and model are required to re-resolve them
|
||||
needs_dirs = any(x is not None for x in (layers, token_id, replacement_id, mode))
|
||||
if not needs_dirs:
|
||||
return self.summary()
|
||||
if lens_manager is None or jl is None:
|
||||
raise ValueError("model and lens required to edit the rule")
|
||||
lens = lens_manager.lens
|
||||
if lens is None:
|
||||
raise ValueError("no lens loaded")
|
||||
tokenizer = jl.tokenizer
|
||||
if mode is not None:
|
||||
if mode not in ("scale", "replace"):
|
||||
raise ValueError(f"invalid mode: {mode}")
|
||||
rule["mode"] = mode
|
||||
if token_id is not None:
|
||||
rule["token_id"] = int(token_id)
|
||||
rule["token"] = tokenizer.decode([int(token_id)])
|
||||
if replacement_id is not None:
|
||||
rule["replacement_id"] = int(replacement_id)
|
||||
rule["replacement"] = tokenizer.decode([int(replacement_id)])
|
||||
if rule["mode"] == "scale":
|
||||
rule["replacement_id"] = None
|
||||
rule["replacement"] = None
|
||||
elif rule["replacement_id"] is None:
|
||||
raise ValueError("replacement_id required in replace mode")
|
||||
if layers is not None:
|
||||
n_layers = len(jl.layers)
|
||||
# new_layers=[] is valid: rule kept but inactive
|
||||
rule["layers"] = sorted({int(l) for l in layers if 0 <= int(l) < n_layers})
|
||||
weight = jl._lm_head.weight
|
||||
rule["dirs_a"] = self._direction(lens, weight, rule["token_id"], rule["layers"])
|
||||
rule["dirs_b"] = (
|
||||
self._direction(lens, weight, rule["replacement_id"], rule["layers"])
|
||||
if rule["replacement_id"] is not None
|
||||
else None
|
||||
)
|
||||
return self.summary()
|
||||
raise ValueError(f"unknown rule {rule_id}")
|
||||
|
||||
def remove(self, rule_id=None):
|
||||
with self._lock:
|
||||
self.detach()
|
||||
if rule_id is None:
|
||||
self._rules = []
|
||||
else:
|
||||
self._rules = [r for r in self._rules if r["id"] != rule_id]
|
||||
return self.summary()
|
||||
|
||||
def attach(self, jl):
|
||||
if not self._rules:
|
||||
return
|
||||
if self._mode == "abliteration":
|
||||
self._attach_abliteration(jl)
|
||||
return
|
||||
if self._mode in ("readthrough", "exact"):
|
||||
self._attach_rebase(jl, exact=self._mode == "exact")
|
||||
return
|
||||
by_layer = {}
|
||||
for rule in self._active_rules():
|
||||
for layer in rule["layers"]:
|
||||
by_layer.setdefault(layer, []).append(rule)
|
||||
|
||||
def make_hook(layer, rules):
|
||||
def hook(module, inputs, output):
|
||||
h = output[0] if isinstance(output, tuple) else output
|
||||
g = self._scale
|
||||
for rule in rules:
|
||||
alpha, beta = effective_coeffs(rule["mode"], rule["factor"], g)
|
||||
vA = rule["dirs_a"][layer].to(h.device, h.dtype)
|
||||
coef = (h * vA).sum(-1, keepdim=True)
|
||||
h = h + alpha * coef * vA
|
||||
if beta:
|
||||
vB = rule["dirs_b"][layer].to(h.device, h.dtype)
|
||||
h = h + beta * coef * vB
|
||||
if isinstance(output, tuple):
|
||||
return (h,) + tuple(output[1:])
|
||||
return h
|
||||
|
||||
return hook
|
||||
|
||||
self._handles = [
|
||||
jl.layers[layer].register_forward_hook(make_hook(layer, rules))
|
||||
for layer, rules in by_layer.items()
|
||||
]
|
||||
|
||||
def _attach_abliteration(self, jl):
|
||||
# Abliteration-mode preview: the SAME projection on every residual write
|
||||
# (embed + each block's output), mirroring the pure-weight bake. A rule's
|
||||
# layers make no sense here (global projection), but layers=[] stays THE
|
||||
# "rule disabled" gesture: we honor it too.
|
||||
active = self._active_rules()
|
||||
if not active:
|
||||
return
|
||||
weight_u = jl._lm_head.weight
|
||||
dirs = [(abliteration_direction(weight_u, r), r) for r in active]
|
||||
|
||||
def apply(h):
|
||||
g = self._scale
|
||||
for (v_a, v_b), rule in dirs:
|
||||
alpha, beta = effective_coeffs(rule["mode"], rule["factor"], g)
|
||||
va = v_a.to(h.device, h.dtype)
|
||||
coef = (h * va).sum(-1, keepdim=True)
|
||||
h = h + alpha * coef * va
|
||||
if beta:
|
||||
h = h + beta * coef * v_b.to(h.device, h.dtype)
|
||||
return h
|
||||
|
||||
def emb_hook(module, inputs, output):
|
||||
return apply(output)
|
||||
|
||||
def blk_hook(module, inputs, output):
|
||||
h = output[0] if isinstance(output, tuple) else output
|
||||
h = apply(h)
|
||||
return (h,) + tuple(output[1:]) if isinstance(output, tuple) else h
|
||||
|
||||
self._handles = [jl._embed_tokens.register_forward_hook(emb_hook)]
|
||||
self._handles += [blk.register_forward_hook(blk_hook) for blk in jl.layers]
|
||||
|
||||
def _attach_rebase(self, jl, exact):
|
||||
# readthrough/exact preview: the SAME transform as the bake (core/rebase),
|
||||
# applied by hooks on the OUTPUT of the reading RMSNorms (and, in exact
|
||||
# mode, on the downstream writes) — the preview and the exported
|
||||
# checkpoint differ only by rounding.
|
||||
from core import rebase # local import (rebase imports effective_coeffs from here)
|
||||
|
||||
active = self._active_rules()
|
||||
if not active:
|
||||
return
|
||||
n_layers = len(jl.layers)
|
||||
cums = rebase.cumulative(active, self._scale, n_layers)
|
||||
if not cums:
|
||||
return
|
||||
|
||||
def read_hook_for(norm, U, V):
|
||||
Ug, Vg = rebase.gamma_pair(norm, U, V)
|
||||
weight = norm.weight
|
||||
Ug = Ug.to(weight.device, weight.dtype)
|
||||
Vg = Vg.to(weight.device, weight.dtype)
|
||||
|
||||
def hook(module, inputs, output):
|
||||
return output + (output @ Vg) @ Ug.T
|
||||
|
||||
return hook
|
||||
|
||||
def write_hook_for(module, U_inv, V):
|
||||
weight = module.weight
|
||||
U_inv = U_inv.to(weight.device, weight.dtype)
|
||||
V = V.to(weight.device, weight.dtype)
|
||||
|
||||
def hook(module, inputs, output):
|
||||
return output - (output @ V) @ U_inv.T
|
||||
|
||||
return hook
|
||||
|
||||
handles = []
|
||||
for m in sorted(k for k in cums if k < n_layers):
|
||||
U, V = cums[m]
|
||||
block = jl.layers[m]
|
||||
norms = {}
|
||||
for _suffix, _module, norm in rebase.iter_reads(block):
|
||||
norms[id(norm)] = norm
|
||||
for norm in norms.values():
|
||||
handles.append(norm.register_forward_hook(read_hook_for(norm, U, V)))
|
||||
if exact:
|
||||
U_inv, Vw, _regularized = rebase.inverse_uv(U, V)
|
||||
for _suffix, module in rebase.iter_writes(block):
|
||||
handles.append(module.register_forward_hook(write_hook_for(module, U_inv, Vw)))
|
||||
U, V = cums[n_layers]
|
||||
handles.append(
|
||||
jl._final_norm.register_forward_hook(read_hook_for(jl._final_norm, U, V))
|
||||
)
|
||||
self._handles = handles
|
||||
|
||||
def detach(self):
|
||||
for handle in self._handles:
|
||||
handle.remove()
|
||||
self._handles = []
|
||||
+621
@@ -0,0 +1,621 @@
|
||||
import json
|
||||
import re
|
||||
import shutil
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
from safetensors import safe_open
|
||||
from safetensors.torch import save_file
|
||||
|
||||
import config
|
||||
from core import rebase
|
||||
from core.ablation import abliteration_direction, effective_coeffs
|
||||
|
||||
EDITS_DIR = config.DATA_DIR / "edits"
|
||||
PRESETS_DIR = config.DATA_DIR / "presets"
|
||||
|
||||
# Residual writes edited by the global abliteration (embed aside)
|
||||
TARGET_SUFFIXES = ("self_attn.o_proj", "mlp.down_proj")
|
||||
|
||||
|
||||
def _now():
|
||||
return datetime.now(timezone.utc).isoformat(timespec="seconds")
|
||||
|
||||
|
||||
def list_presets():
|
||||
PRESETS_DIR.mkdir(parents=True, exist_ok=True)
|
||||
out = []
|
||||
for path in sorted(PRESETS_DIR.glob("*.json")):
|
||||
data = json.loads(path.read_text(encoding="utf-8"))
|
||||
out.append({"name": path.stem, "n_rules": len(data.get("rules", [])), "model_id": data.get("model_id")})
|
||||
return out
|
||||
|
||||
|
||||
def save_preset(name, rules, model_id, scale=1.0):
|
||||
PRESETS_DIR.mkdir(parents=True, exist_ok=True)
|
||||
payload = {"model_id": model_id, "saved_at": _now(), "scale": scale, "rules": rules}
|
||||
(PRESETS_DIR / f"{name}.json").write_text(
|
||||
json.dumps(payload, ensure_ascii=False, indent=1), encoding="utf-8"
|
||||
)
|
||||
return payload
|
||||
|
||||
|
||||
def load_preset(name):
|
||||
path = PRESETS_DIR / f"{name}.json"
|
||||
if not path.exists():
|
||||
raise ValueError(f"unknown preset {name}")
|
||||
return json.loads(path.read_text(encoding="utf-8"))
|
||||
|
||||
|
||||
def delete_preset(name):
|
||||
(PRESETS_DIR / f"{name}.json").unlink(missing_ok=True)
|
||||
|
||||
|
||||
def compute_abliteration(rules, jl, scale=1.0):
|
||||
"""Global pure-weight edit reproducing the abliteration-mode preview.
|
||||
|
||||
Applies to EVERY residual write (embed_tokens + o_proj/down_proj of every
|
||||
layer) the same transform as the abliteration-mode hooks: for each rule,
|
||||
``out += scale·(v̂_A·out)·w`` (applied sequentially, like the hooks). Since
|
||||
the residual is the sum of all these writes, the direction is
|
||||
removed/redirected across the whole residual — hence the fidelity (~0.97
|
||||
cosine on the logits). This is the pure-weights path for architectures the
|
||||
rebase does not support (write norms, Gemma style).
|
||||
|
||||
Returns ``(tensors, info)``:
|
||||
- ``tensors``: {param_name: W_new (cpu, float32)}
|
||||
- ``info``: {tied, embed_key, lm_head_key, path, delta_max, lowrank}
|
||||
where ``lowrank`` = {param_name: (B [out, r], A [r, in])} — the SAME edit
|
||||
as per-rule rank-1 factors (delta = B·A), exact, for the LoRA export.
|
||||
For the embed, delta = (B·A)ᵀ (PEFT lookup convention).
|
||||
"""
|
||||
# layers=[] = disabled rule, in this mode too (consistent with the preview)
|
||||
rules = [r for r in rules if r["layers"]]
|
||||
if not rules:
|
||||
raise ValueError("no active rule (all have 0 layers): nothing to export")
|
||||
path = jl.layout.path
|
||||
weight_u = jl._lm_head.weight
|
||||
# (v_a, w_eff) per rule, with w_eff = alpha·v̂_A + beta·v̂_B: the SAME effective
|
||||
# coefficients (saturation included) as the preview hooks
|
||||
pairs = []
|
||||
for r in rules:
|
||||
v_a, v_b = abliteration_direction(weight_u, r)
|
||||
alpha, beta = effective_coeffs(r["mode"], r["factor"], scale)
|
||||
w_eff = alpha * v_a
|
||||
if beta:
|
||||
w_eff = w_eff + beta * v_b
|
||||
pairs.append((v_a, w_eff))
|
||||
|
||||
# bake on CPU: the float32 matrices (embed ~1.5 GB) don't fit alongside the
|
||||
# model on the GPU (OOM measured on 12 GB with a 4B loaded)
|
||||
def apply_cols(W): # [d_model, d_in]: residual output = rows
|
||||
cur, us, rows = W, [], []
|
||||
for v_a, w in pairs:
|
||||
row = v_a @ cur # composed over the previous rules
|
||||
us.append(w)
|
||||
rows.append(row)
|
||||
cur = cur + torch.outer(w, row)
|
||||
return cur, torch.stack(us, dim=1), torch.stack(rows, dim=0)
|
||||
|
||||
def apply_rows(E): # [vocab, d_model]: each ROW is a residual vector
|
||||
cur, us, rows = E, [], []
|
||||
for v_a, w in pairs:
|
||||
col = cur @ v_a # [vocab]
|
||||
us.append(w)
|
||||
rows.append(col)
|
||||
cur = cur + torch.outer(col, w)
|
||||
return cur, torch.stack(us, dim=1), torch.stack(rows, dim=0)
|
||||
|
||||
tensors = {}
|
||||
lowrank = {}
|
||||
delta_max = 0.0
|
||||
|
||||
embed_key = f"{path}.{jl.layout.embed}.weight"
|
||||
E = jl._embed_tokens.weight.detach().float().cpu()
|
||||
E_new, B, A = apply_rows(E)
|
||||
delta_max = max(delta_max, (E_new - E).abs().max().item())
|
||||
tensors[embed_key] = E_new
|
||||
lowrank[embed_key] = (B, A) # delta_embed = (B·A)ᵀ = summed outer(A_k, B_k)
|
||||
|
||||
skipped_writes = 0
|
||||
for i, block in enumerate(jl.layers):
|
||||
for suffix in TARGET_SUFFIXES:
|
||||
module = block
|
||||
for part in suffix.split("."):
|
||||
module = getattr(module, part, None)
|
||||
if module is None:
|
||||
break
|
||||
if module is None: # e.g. linear-attention blocks (no self_attn)
|
||||
skipped_writes += 1
|
||||
continue
|
||||
W = module.weight.detach().float().cpu()
|
||||
W_new, B, A = apply_cols(W)
|
||||
delta_max = max(delta_max, (W_new - W).abs().max().item())
|
||||
name = f"{path}.layers.{i}.{suffix}.weight"
|
||||
tensors[name] = W_new
|
||||
lowrank[name] = (B, A)
|
||||
|
||||
tied = jl._lm_head.weight.data_ptr() == jl._embed_tokens.weight.data_ptr()
|
||||
info = {
|
||||
"tied": tied,
|
||||
"embed_key": embed_key,
|
||||
"lm_head_key": f"{jl.layout.lm_head}.weight",
|
||||
"path": path,
|
||||
"delta_max": delta_max,
|
||||
"lowrank": lowrank,
|
||||
"skipped_writes": skipped_writes,
|
||||
}
|
||||
return tensors, info
|
||||
|
||||
|
||||
def _abliteration_warnings(rules):
|
||||
warns = []
|
||||
for r in rules:
|
||||
if r["mode"] == "scale" and r["factor"] > 1.0:
|
||||
warns.append(
|
||||
f"\"{(r['token'] or '').strip()}\" ×{r['factor']}: amplifying (factor > 1) "
|
||||
"is approximate in pure weights (the hook composes over the layers)"
|
||||
)
|
||||
return warns
|
||||
|
||||
|
||||
def export_abliteration(rules, jl, model_meta, *, fmt, name, source_dir=None, scale=1.0):
|
||||
"""Pure-weight export (global abliteration). Formats: ``full`` (full
|
||||
checkpoint), ``layers`` (safetensors of only the modified matrices) and
|
||||
``lora`` (exact PEFT adapter, rank = n_rules; embed omitted if embeddings
|
||||
are tied). Unties ``lm_head`` (full/layers) if the model has tied embeddings,
|
||||
to preserve the original un-embedding."""
|
||||
rules = [r for r in rules if r["layers"]] # layers=[] = disabled rule
|
||||
if not rules:
|
||||
raise ValueError("no active intervention to export")
|
||||
if fmt not in ("full", "layers", "lora"):
|
||||
raise ValueError(f"unknown format for abliteration: {fmt}")
|
||||
|
||||
tensors, info = compute_abliteration(rules, jl, scale=scale)
|
||||
if info["delta_max"] < 1e-8:
|
||||
raise ValueError(
|
||||
"the bake changes no weight (neutral factors, scale=0 or null "
|
||||
"directions) — the export would be identical to the original model"
|
||||
)
|
||||
out_dir = EDITS_DIR / name
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
dtype = torch.bfloat16 if model_meta.get("dtype") == "bf16" else torch.float16
|
||||
lm_head_key = info["lm_head_key"]
|
||||
summary = [
|
||||
{k: r[k] for k in ("token_id", "token", "mode", "factor", "replacement_id", "replacement")}
|
||||
for r in rules
|
||||
]
|
||||
meta = {
|
||||
"name": name,
|
||||
"format": fmt,
|
||||
"method": "abliteration-global",
|
||||
"model_id": model_meta.get("model_id"),
|
||||
"model_revision": model_meta.get("revision"),
|
||||
"dtype": model_meta.get("dtype"),
|
||||
"global_scale": scale,
|
||||
"untied_lm_head": info["tied"] and fmt in ("full", "layers"),
|
||||
"rules": summary,
|
||||
"modified_params_count": len(tensors) + (1 if info["tied"] else 0),
|
||||
"warnings": _abliteration_warnings(rules) + (
|
||||
[f"{info['skipped_writes']} residual write(s) without o_proj/down_proj "
|
||||
"(hybrid architecture) left untouched — the bake is partial there; "
|
||||
"prefer read projection when the architecture supports it"]
|
||||
if info["skipped_writes"] else []
|
||||
),
|
||||
"note": (
|
||||
"global abliteration: the token's direction is removed/redirected in "
|
||||
"every residual write (embed + o_proj/down_proj of all layers). "
|
||||
"Reproduces the abliteration-mode preview (~0.97 cosine on the logits). "
|
||||
"Pure weights: a standard safetensors checkpoint."
|
||||
),
|
||||
"created_at": _now(),
|
||||
}
|
||||
|
||||
if fmt == "layers":
|
||||
out = {k: v.to(dtype) for k, v in tensors.items()}
|
||||
if info["tied"]:
|
||||
# original un-embedding (unedited embed) to write separately
|
||||
out[lm_head_key] = jl._embed_tokens.weight.detach().to(dtype).cpu()
|
||||
save_file(out, str(out_dir / "modified_layers.safetensors"))
|
||||
|
||||
elif fmt == "lora":
|
||||
# The abliteration delta is EXACTLY rank-n_rules per matrix (delta = B·A),
|
||||
# so the LoRA is exact — except the embed of a tied-embeddings model: PEFT
|
||||
# can't untie lm_head, and editing the embed would corrupt the shared
|
||||
# un-embedding → we omit it (reduced fidelity).
|
||||
include_embed = not info["tied"]
|
||||
if not include_embed:
|
||||
meta["warnings"] = meta["warnings"] + [
|
||||
"tied embeddings: the embed is not included in the LoRA (PEFT "
|
||||
"cannot untie lm_head) — prefer \"full checkpoint\" for maximum "
|
||||
"fidelity"
|
||||
]
|
||||
out = {}
|
||||
target_modules = set()
|
||||
for pname, (B, A) in info["lowrank"].items():
|
||||
base = pname.removesuffix(".weight")
|
||||
if pname == info["embed_key"]:
|
||||
if not include_embed:
|
||||
continue
|
||||
target_modules.add(base.rsplit(".", 1)[-1])
|
||||
# PEFT Embedding convention: delta_lookup = (B·A)ᵀ,
|
||||
# A = lora_embedding_A [r, vocab], B = lora_embedding_B [d_model, r]
|
||||
out[f"base_model.model.{base}.lora_embedding_A"] = A.contiguous()
|
||||
out[f"base_model.model.{base}.lora_embedding_B"] = B.contiguous()
|
||||
else:
|
||||
target_modules.add(base.rsplit(".", 1)[-1])
|
||||
out[f"base_model.model.{base}.lora_A.weight"] = A.contiguous()
|
||||
out[f"base_model.model.{base}.lora_B.weight"] = B.contiguous()
|
||||
rank = len(rules)
|
||||
save_file(out, str(out_dir / "adapter_model.safetensors"))
|
||||
adapter_config = {
|
||||
"peft_type": "LORA",
|
||||
"base_model_name_or_path": model_meta.get("model_id"),
|
||||
"r": rank,
|
||||
"lora_alpha": rank,
|
||||
"lora_dropout": 0.0,
|
||||
"target_modules": sorted(target_modules),
|
||||
"bias": "none",
|
||||
"fan_in_fan_out": False,
|
||||
"task_type": "CAUSAL_LM",
|
||||
}
|
||||
(out_dir / "adapter_config.json").write_text(
|
||||
json.dumps(adapter_config, indent=1), encoding="utf-8"
|
||||
)
|
||||
|
||||
elif fmt == "full":
|
||||
if source_dir is None or not Path(source_dir).is_dir():
|
||||
raise ValueError("full checkpoint: model source folder not found")
|
||||
source_dir = Path(source_dir)
|
||||
shards = sorted(source_dir.glob("*.safetensors"))
|
||||
if not shards:
|
||||
raise ValueError("full checkpoint: no safetensors in the source")
|
||||
|
||||
lm_head_value = None # original un-embedding (if tied) = original embed from disk
|
||||
embed_shard_name = None
|
||||
seen = set()
|
||||
for shard in shards:
|
||||
ino = shard.stat().st_ino
|
||||
if ino in seen:
|
||||
continue
|
||||
seen.add(ino)
|
||||
out = {}
|
||||
with safe_open(str(shard), framework="pt") as f:
|
||||
keys = list(f.keys())
|
||||
for key in keys:
|
||||
original = f.get_tensor(key)
|
||||
if info["tied"] and key == info["embed_key"]:
|
||||
lm_head_value = original.clone() # BEFORE editing
|
||||
embed_shard_name = shard.name
|
||||
out[key] = tensors[key].to(original.dtype) if key in tensors else original
|
||||
# if this shard already carries lm_head (untied model), don't touch it
|
||||
save_file(out, str(out_dir / shard.name))
|
||||
|
||||
# untie: add lm_head.weight (= original embed) into the embed's shard
|
||||
if info["tied"]:
|
||||
if lm_head_value is None:
|
||||
raise ValueError("cannot untie: embed not found in the source")
|
||||
target_shard = out_dir / embed_shard_name
|
||||
with safe_open(str(target_shard), framework="pt") as f:
|
||||
merged = {k: f.get_tensor(k) for k in f.keys()}
|
||||
merged[lm_head_key] = lm_head_value
|
||||
save_file(merged, str(target_shard))
|
||||
|
||||
# config.json: copy, force tie_word_embeddings=False if untied
|
||||
cfg_path = source_dir / "config.json"
|
||||
if cfg_path.exists():
|
||||
cfg = json.loads(cfg_path.read_text(encoding="utf-8"))
|
||||
if info["tied"]:
|
||||
cfg["tie_word_embeddings"] = False
|
||||
(out_dir / "config.json").write_text(
|
||||
json.dumps(cfg, ensure_ascii=False, indent=1), encoding="utf-8"
|
||||
)
|
||||
# other tokenizer/config files (json, merges.txt, tokenizer.model…):
|
||||
# copy as-is, then fix the index if present
|
||||
for pattern in ("*.json", "*.txt", "*.model", "*.tiktoken", "*.jinja"):
|
||||
for extra in source_dir.glob(pattern):
|
||||
if extra.name == "config.json":
|
||||
continue
|
||||
shutil.copy2(extra, out_dir / extra.name)
|
||||
index_path = out_dir / "model.safetensors.index.json"
|
||||
if info["tied"] and index_path.exists():
|
||||
index = json.loads(index_path.read_text(encoding="utf-8"))
|
||||
wm = index.setdefault("weight_map", {})
|
||||
wm[lm_head_key] = embed_shard_name
|
||||
if "metadata" in index and "total_size" in index["metadata"]:
|
||||
index["metadata"]["total_size"] += lm_head_value.numel() * lm_head_value.element_size()
|
||||
index_path.write_text(json.dumps(index, indent=1), encoding="utf-8")
|
||||
|
||||
(out_dir / "edit_meta.json").write_text(
|
||||
json.dumps(meta, ensure_ascii=False, indent=1), encoding="utf-8"
|
||||
)
|
||||
return {"out_dir": str(out_dir), **meta}
|
||||
|
||||
|
||||
def _disk_mapper(mem_embed_key, disk_keys):
|
||||
"""Memory keys (instantiated model's layout) → disk checkpoint keys.
|
||||
|
||||
transformers renames on load: e.g. Qwen3.5 is instantiated as ForCausalLM
|
||||
("model.layers.*" in memory) but saved in ConditionalGeneration format
|
||||
("model.language_model.layers.*"). Without this mapping, a "full" export
|
||||
would copy the source verbatim without transforming anything. We anchor the
|
||||
disk prefix on the embed, whose suffix is unique in the checkpoint."""
|
||||
if mem_embed_key in disk_keys:
|
||||
return lambda key: key
|
||||
suffix = "." + ".".join(mem_embed_key.rsplit(".", 2)[-2:]) # ".embed_tokens.weight"
|
||||
candidates = [k for k in disk_keys if k.endswith(suffix)]
|
||||
if len(candidates) != 1:
|
||||
raise ValueError(
|
||||
f"checkpoint prefix undecidable: {mem_embed_key} absent from the source "
|
||||
f"and {len(candidates)} key(s) end with {suffix}"
|
||||
)
|
||||
mem_prefix = mem_embed_key.removesuffix(suffix)
|
||||
disk_prefix = candidates[0].removesuffix(suffix)
|
||||
|
||||
def to_disk(key):
|
||||
if key == mem_prefix or key.startswith(mem_prefix + "."):
|
||||
return disk_prefix + key[len(mem_prefix):]
|
||||
return key
|
||||
|
||||
return to_disk
|
||||
|
||||
|
||||
def export_rebase(rules, jl, model_meta, *, fmt, name, source_dir=None, scale=1.0, exact=False):
|
||||
"""Pure-weight export by change of basis of the reads (cf. core/rebase).
|
||||
|
||||
``readthrough`` (exact=False): the downstream read matrices + lm_head.
|
||||
``exact``: adds the counter-transform of the downstream writes.
|
||||
Formats: ``full`` (checkpoint), ``layers`` (safetensors of the modified
|
||||
matrices) and ``lora`` (PEFT adapter = the exact low-rank diff between the
|
||||
baked weights and the originals; the lm_head delta is applied at forward
|
||||
time, so tied embeddings need no untying). The bake is done streaming, one
|
||||
float32 CPU matrix at a time. Tied-embeddings model (full/layers): the embed
|
||||
stays INTACT, it's lm_head (untied) that receives the final read transform."""
|
||||
method = "rebase-exact" if exact else "rebase-readthrough"
|
||||
if fmt not in ("full", "layers", "lora"):
|
||||
raise ValueError(f"unknown format for {method}: {fmt}")
|
||||
transforms, info = rebase.build_plan(rules, jl, scale, exact=exact)
|
||||
lm_head_key = info["lm_head_key"]
|
||||
|
||||
delta_max = 0.0
|
||||
applied = set()
|
||||
|
||||
def bake(key, tensor):
|
||||
nonlocal delta_max
|
||||
W = tensor.detach().to("cpu", torch.float32)
|
||||
W_new, _B, _A = rebase.apply_transform(transforms[key], W)
|
||||
delta_max = max(delta_max, (W_new - W).abs().max().item())
|
||||
applied.add(key)
|
||||
return W_new
|
||||
|
||||
out_dir = EDITS_DIR / name
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
dtype = torch.bfloat16 if model_meta.get("dtype") == "bf16" else torch.float16
|
||||
warnings = []
|
||||
if exact and info["regularized_layers"]:
|
||||
warnings.append(
|
||||
"regularized inverse (full zap ⇒ singular transform) on layers "
|
||||
f"{info['regularized_layers']} — the effect there equals readthrough; "
|
||||
"prefer readthrough mode for full removals"
|
||||
)
|
||||
if fmt == "lora" and info["tied"]:
|
||||
warnings.append(
|
||||
"tied embeddings: use the adapter at runtime (PEFT applies the "
|
||||
"lm_head delta at forward time, leaving the shared embed intact); "
|
||||
"merging it into the base weights (merge_and_unload) would write "
|
||||
"that delta into the embed too — export a full checkpoint if you "
|
||||
"need merged weights"
|
||||
)
|
||||
|
||||
def source_weight(state, key):
|
||||
source = state.get(key)
|
||||
if source is None and key == lm_head_key and info["tied"]:
|
||||
source = state[info["embed_key"]] # tied: the un-embedding IS the embed
|
||||
if source is None:
|
||||
raise ValueError(
|
||||
f"parameter {key} not found in the loaded model — "
|
||||
"unexpected layout, export cancelled"
|
||||
)
|
||||
return source
|
||||
|
||||
if fmt == "layers":
|
||||
state = jl._hf_model.state_dict()
|
||||
to_disk = lambda key: key # noqa: E731 — refined if the source is available
|
||||
if source_dir is not None and Path(source_dir).is_dir():
|
||||
disk_keys = set()
|
||||
for shard in Path(source_dir).glob("*.safetensors"):
|
||||
with safe_open(str(shard), framework="pt") as f:
|
||||
disk_keys.update(f.keys())
|
||||
if disk_keys:
|
||||
to_disk = _disk_mapper(info["embed_key"], disk_keys)
|
||||
tensors = {}
|
||||
for key in transforms:
|
||||
tensors[to_disk(key)] = bake(key, source_weight(state, key)).to(dtype)
|
||||
save_file(tensors, str(out_dir / "modified_layers.safetensors"))
|
||||
|
||||
elif fmt == "lora":
|
||||
# The rebase delta is low-rank by construction (delta = B·A exactly, cf.
|
||||
# rebase.apply_transform): the adapter is the exact diff between the
|
||||
# baked weights and the originals, not an approximation. lm_head: PEFT
|
||||
# adds the delta at forward time without writing to the (possibly tied)
|
||||
# weight, so the un-embedding is effectively untied while the embed
|
||||
# stays intact. Module names follow the model as instantiated by
|
||||
# AutoModelForCausalLM (the same loading path as the UI).
|
||||
state = jl._hf_model.state_dict()
|
||||
factors = {}
|
||||
max_rank = 0
|
||||
for key in transforms:
|
||||
W = source_weight(state, key).detach().to("cpu", torch.float32)
|
||||
W_new, B, A = rebase.apply_transform(transforms[key], W)
|
||||
delta_max = max(delta_max, (W_new - W).abs().max().item())
|
||||
applied.add(key)
|
||||
factors[key] = (B, A)
|
||||
max_rank = max(max_rank, B.shape[1])
|
||||
tensors = {}
|
||||
module_paths = []
|
||||
for key, (B, A) in factors.items():
|
||||
base = key.removesuffix(".weight")
|
||||
module_paths.append(base)
|
||||
if B.shape[1] < max_rank: # pad so a single config `r` fits every module
|
||||
pad = max_rank - B.shape[1]
|
||||
B = torch.cat([B, torch.zeros(B.shape[0], pad)], dim=1)
|
||||
A = torch.cat([A, torch.zeros(pad, A.shape[1])], dim=0)
|
||||
tensors[f"base_model.model.{base}.lora_A.weight"] = A.contiguous()
|
||||
tensors[f"base_model.model.{base}.lora_B.weight"] = B.contiguous()
|
||||
save_file(tensors, str(out_dir / "adapter_model.safetensors"))
|
||||
# target_modules as an anchored regex over the modules actually edited:
|
||||
# a plain suffix list would wrap the same projections in EVERY layer and
|
||||
# leave benign but alarming "missing adapter keys" warnings at load time
|
||||
target_regex = "(.*\\.)?(" + "|".join(re.escape(p) for p in sorted(module_paths)) + ")"
|
||||
adapter_config = {
|
||||
"peft_type": "LORA",
|
||||
"base_model_name_or_path": model_meta.get("model_id"),
|
||||
"r": max_rank,
|
||||
"lora_alpha": max_rank, # scaling alpha/r = 1: B·A is the raw delta
|
||||
"lora_dropout": 0.0,
|
||||
"target_modules": target_regex,
|
||||
"bias": "none",
|
||||
"fan_in_fan_out": False,
|
||||
"task_type": "CAUSAL_LM",
|
||||
}
|
||||
(out_dir / "adapter_config.json").write_text(
|
||||
json.dumps(adapter_config, indent=1), encoding="utf-8"
|
||||
)
|
||||
|
||||
elif fmt == "full":
|
||||
if source_dir is None or not Path(source_dir).is_dir():
|
||||
raise ValueError("full checkpoint: model source folder not found")
|
||||
source_dir = Path(source_dir)
|
||||
shards = sorted(source_dir.glob("*.safetensors"))
|
||||
if not shards:
|
||||
raise ValueError("full checkpoint: no safetensors in the source")
|
||||
|
||||
disk_keys = set()
|
||||
seen = set()
|
||||
for shard in shards:
|
||||
ino = shard.stat().st_ino
|
||||
if ino in seen:
|
||||
continue
|
||||
seen.add(ino)
|
||||
with safe_open(str(shard), framework="pt") as f:
|
||||
disk_keys.update(f.keys())
|
||||
to_disk = _disk_mapper(info["embed_key"], disk_keys)
|
||||
transforms = {to_disk(k): fn for k, fn in transforms.items()}
|
||||
lm_head_key = to_disk(lm_head_key)
|
||||
embed_key = to_disk(info["embed_key"])
|
||||
|
||||
lm_head_written = False
|
||||
embed_shard_name = None
|
||||
seen = set()
|
||||
for shard in shards:
|
||||
ino = shard.stat().st_ino
|
||||
if ino in seen:
|
||||
continue
|
||||
seen.add(ino)
|
||||
out = {}
|
||||
with safe_open(str(shard), framework="pt") as f:
|
||||
for key in f.keys():
|
||||
original = f.get_tensor(key)
|
||||
if key in transforms:
|
||||
out[key] = bake(key, original).to(original.dtype)
|
||||
if key == lm_head_key:
|
||||
lm_head_written = True
|
||||
else:
|
||||
out[key] = original
|
||||
if key == embed_key:
|
||||
embed_shard_name = shard.name
|
||||
save_file(out, str(out_dir / shard.name))
|
||||
del out
|
||||
|
||||
# untie: the transformed un-embedding becomes a separate lm_head, baked
|
||||
# from the original embed (which stays intact)
|
||||
if info["tied"] and not lm_head_written:
|
||||
if embed_shard_name is None:
|
||||
raise ValueError("cannot untie: embed not found in the source")
|
||||
target_shard = out_dir / embed_shard_name
|
||||
with safe_open(str(target_shard), framework="pt") as f:
|
||||
merged = {k: f.get_tensor(k) for k in f.keys()}
|
||||
embed_original = merged[embed_key]
|
||||
lm_head_value = bake(lm_head_key, embed_original).to(embed_original.dtype)
|
||||
merged[lm_head_key] = lm_head_value
|
||||
save_file(merged, str(out_dir / embed_shard_name))
|
||||
del merged
|
||||
|
||||
missing = set(transforms) - applied
|
||||
if missing:
|
||||
sample = sorted(missing)[:3]
|
||||
raise ValueError(
|
||||
f"{len(missing)} parameter(s) to transform absent from the source "
|
||||
f"checkpoint (e.g. {sample}) — unexpected key names, export cancelled "
|
||||
"(the written checkpoint would be partially original)"
|
||||
)
|
||||
|
||||
cfg_path = source_dir / "config.json"
|
||||
if cfg_path.exists():
|
||||
cfg = json.loads(cfg_path.read_text(encoding="utf-8"))
|
||||
if info["tied"]:
|
||||
cfg["tie_word_embeddings"] = False
|
||||
text_cfg = cfg.get("text_config")
|
||||
if isinstance(text_cfg, dict) and "tie_word_embeddings" in text_cfg:
|
||||
text_cfg["tie_word_embeddings"] = False
|
||||
(out_dir / "config.json").write_text(
|
||||
json.dumps(cfg, ensure_ascii=False, indent=1), encoding="utf-8"
|
||||
)
|
||||
for pattern in ("*.json", "*.txt", "*.model", "*.tiktoken", "*.jinja"):
|
||||
for extra in source_dir.glob(pattern):
|
||||
if extra.name == "config.json":
|
||||
continue
|
||||
shutil.copy2(extra, out_dir / extra.name)
|
||||
index_path = out_dir / "model.safetensors.index.json"
|
||||
if info["tied"] and not lm_head_written and index_path.exists():
|
||||
index = json.loads(index_path.read_text(encoding="utf-8"))
|
||||
wm = index.setdefault("weight_map", {})
|
||||
wm[lm_head_key] = embed_shard_name
|
||||
if "metadata" in index and "total_size" in index["metadata"]:
|
||||
index["metadata"]["total_size"] += (
|
||||
lm_head_value.numel() * lm_head_value.element_size()
|
||||
)
|
||||
index_path.write_text(json.dumps(index, indent=1), encoding="utf-8")
|
||||
|
||||
if delta_max < 1e-8:
|
||||
shutil.rmtree(out_dir, ignore_errors=True)
|
||||
raise ValueError(
|
||||
"the bake changes no weight (null directions?) — the export would be "
|
||||
"identical to the original model, folder deleted"
|
||||
)
|
||||
|
||||
summary = [
|
||||
{k: r[k] for k in ("token_id", "token", "mode", "factor", "replacement_id", "replacement", "layers")}
|
||||
for r in rules if r["layers"]
|
||||
]
|
||||
meta = {
|
||||
"name": name,
|
||||
"format": fmt,
|
||||
"method": method,
|
||||
"model_id": model_meta.get("model_id"),
|
||||
"model_revision": model_meta.get("revision"),
|
||||
"dtype": model_meta.get("dtype"),
|
||||
"global_scale": scale,
|
||||
# lora: no physical untying — the lm_head delta lives in the adapter
|
||||
"untied_lm_head": info["tied"] and fmt != "lora",
|
||||
"rules": summary,
|
||||
"layers_span": info["layers_span"],
|
||||
"rank": info["rank_final"],
|
||||
"modified_params_count": len(transforms),
|
||||
"delta_max": delta_max,
|
||||
"min_gamma": info["min_gamma"],
|
||||
"warnings": warnings,
|
||||
"note": (
|
||||
"change of basis of the reads: every matrix that READS the residual "
|
||||
"downstream of the hooked layers (q/k/v, in_proj*, gate/up + lm_head) sees "
|
||||
"the residual transformed by the same J-space directions as the live preview"
|
||||
+ (" ; downstream writes counter-transformed (exact mode)" if exact else "")
|
||||
+ ". Pure weights: a standard safetensors checkpoint."
|
||||
),
|
||||
"created_at": _now(),
|
||||
}
|
||||
(out_dir / "edit_meta.json").write_text(
|
||||
json.dumps(meta, ensure_ascii=False, indent=1), encoding="utf-8"
|
||||
)
|
||||
return {"out_dir": str(out_dir), **meta}
|
||||
+374
@@ -0,0 +1,374 @@
|
||||
import hashlib
|
||||
import json
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from datetime import datetime, timezone
|
||||
|
||||
from jlens.lens import JacobianLens
|
||||
|
||||
import config
|
||||
from core.gpus import gpu_stats
|
||||
|
||||
FITS_DIR = config.DATA_DIR / "fits"
|
||||
WORKER = config.ROOT / "scripts" / "fit_worker.py"
|
||||
|
||||
# Fit corpora. "mixed" = both, equal parts (rounded to the nearest prompt).
|
||||
DATASET_WIKITEXT = "Salesforce/wikitext-103-raw-v1"
|
||||
DATASET_HARMLESS = "heretic-org/Semantic-Harmless"
|
||||
FIT_DATASETS = (DATASET_WIKITEXT, DATASET_HARMLESS, "mixed")
|
||||
|
||||
|
||||
def _load_corpus(dataset, n, skip=0):
|
||||
"""``n`` prompts from ``dataset``, skipping the first ``skip`` picks
|
||||
(continue-from: the new prompts must not overlap the base lens's).
|
||||
|
||||
wikitext keeps the historical behavior (first records ≥600 chars, streamed).
|
||||
Semantic-Harmless is a small instruct set (~416 one-line prompts): we draw a
|
||||
seeded random sample — sample(skip+n) then drop the head, so a continued fit
|
||||
extends the same sequence — and PACK the picks into ~350-char sequences
|
||||
(median prompt ≈ 10 tokens, and jlens skips the first 16 positions of every
|
||||
sequence as attention sinks: unpacked, almost every pick would be dropped as
|
||||
"too short"). ``n``/``skip`` count SOURCE prompts, not packs. "mixed" takes
|
||||
equal parts of both (n odd: the extra prompt goes to wikitext) and shuffles
|
||||
the union so multi-GPU slices stay mixed."""
|
||||
if dataset == "mixed":
|
||||
import random
|
||||
|
||||
n_wiki = (n + 1) // 2
|
||||
s_wiki = (skip + 1) // 2
|
||||
prompts = _load_corpus(DATASET_WIKITEXT, n_wiki, s_wiki)
|
||||
prompts += _load_corpus(DATASET_HARMLESS, n - n_wiki, skip // 2)
|
||||
random.Random(1729).shuffle(prompts)
|
||||
return prompts
|
||||
if dataset == DATASET_HARMLESS:
|
||||
import random
|
||||
|
||||
from datasets import load_dataset
|
||||
|
||||
texts = [r["text"] for r in load_dataset(DATASET_HARMLESS, split="train")]
|
||||
if skip + n > len(texts):
|
||||
raise ValueError(
|
||||
f"{DATASET_HARMLESS} has {len(texts)} prompts, "
|
||||
f"{skip + n} requested (continue included) — lower n_prompts"
|
||||
)
|
||||
picks = random.Random(1729).sample(texts, skip + n)[skip:]
|
||||
packs, cur = [], ""
|
||||
for text in picks:
|
||||
cur = f"{cur}\n\n{text}" if cur else text
|
||||
if len(cur) >= 350:
|
||||
packs.append(cur)
|
||||
cur = ""
|
||||
if cur:
|
||||
# a lone sub-16-token tail would be skipped by jlens anyway: fold it
|
||||
# into the previous pack instead of losing it
|
||||
if packs and len(cur) < 120:
|
||||
packs[-1] += "\n\n" + cur
|
||||
else:
|
||||
packs.append(cur)
|
||||
return packs
|
||||
from jlens.examples import load_wikitext_prompts
|
||||
|
||||
# load skip + n then keep the tail: the new prompts don't overlap
|
||||
# those of the base lens
|
||||
return load_wikitext_prompts(skip + n)[skip:]
|
||||
|
||||
def _default_dim_batch(device):
|
||||
"""Default dim_batch scaled to the device's VRAM.
|
||||
Measured on a 4B bf16 fit: 8 fits in 16 GB, 4 in 12 GB."""
|
||||
try:
|
||||
total = gpu_stats()[int(device.split(":")[1])]["vram_total"]
|
||||
return 8 if total >= 15 * 2**30 else 4
|
||||
except Exception:
|
||||
return 4
|
||||
|
||||
|
||||
def _now():
|
||||
return datetime.now(timezone.utc).isoformat(timespec="seconds")
|
||||
|
||||
|
||||
class FitManager:
|
||||
def __init__(self):
|
||||
self._lock = threading.Lock()
|
||||
self._procs = []
|
||||
self.state = {"state": "idle"}
|
||||
self.on_progress = None
|
||||
|
||||
def _emit(self):
|
||||
if self.on_progress:
|
||||
self.on_progress(dict(self.state))
|
||||
|
||||
def start(self, *, model_id, source, n_prompts=100, dtype="bf16", quant=None,
|
||||
devices=("cuda:0",), name=None, dim_batch=None,
|
||||
max_seq_len=128, source_layers=None, model_revision=None,
|
||||
continue_from=None, dataset=DATASET_WIKITEXT):
|
||||
with self._lock:
|
||||
if self.state.get("state") == "running":
|
||||
raise ValueError("a fitting is already in progress")
|
||||
if not devices:
|
||||
raise ValueError("at least one device required")
|
||||
if dataset not in FIT_DATASETS:
|
||||
raise ValueError(f"unknown dataset: {dataset} (choices: {', '.join(FIT_DATASETS)})")
|
||||
skip_prompts = 0
|
||||
base_lens = None
|
||||
if continue_from:
|
||||
base_lens = JacobianLens.load(continue_from)
|
||||
# new prompts: skip those already seen by the base lens
|
||||
skip_prompts = base_lens.n_prompts
|
||||
if source_layers is None:
|
||||
source_layers = list(base_lens.source_layers)
|
||||
if name is None:
|
||||
base = model_id.split("/")[-1]
|
||||
if dataset == "mixed":
|
||||
base += "_mixed"
|
||||
elif dataset == DATASET_HARMLESS:
|
||||
base += "_harmless"
|
||||
total = n_prompts + skip_prompts
|
||||
name = f"{base}_n{total}" if continue_from else f"{base}_n{n_prompts}"
|
||||
params = {
|
||||
"model_id": model_id,
|
||||
"source": source,
|
||||
"model_revision": model_revision,
|
||||
"dtype": dtype,
|
||||
"quant": quant,
|
||||
"n_prompts": n_prompts,
|
||||
"dataset": dataset,
|
||||
"devices": list(devices),
|
||||
"dim_batch": dim_batch,
|
||||
"max_seq_len": max_seq_len,
|
||||
"source_layers": source_layers,
|
||||
"continue_from": continue_from,
|
||||
"skip_prompts": skip_prompts,
|
||||
}
|
||||
self.state = {
|
||||
"state": "running",
|
||||
"name": name,
|
||||
"phase": "corpus",
|
||||
"total": n_prompts,
|
||||
"done": 0,
|
||||
"workers": [],
|
||||
"eta_seconds": None,
|
||||
"started_at": _now(),
|
||||
"params": params,
|
||||
"error": None,
|
||||
}
|
||||
self._procs = []
|
||||
threading.Thread(target=self._run, args=(name, params), daemon=True).start()
|
||||
return dict(self.state)
|
||||
|
||||
def stop(self):
|
||||
with self._lock:
|
||||
for proc in self._procs:
|
||||
if proc.poll() is None:
|
||||
proc.terminate()
|
||||
if self.state.get("state") == "running":
|
||||
self.state["state"] = "stopping"
|
||||
self._emit()
|
||||
return dict(self.state)
|
||||
|
||||
def _run(self, name, params):
|
||||
try:
|
||||
job_dir = FITS_DIR / name
|
||||
job_dir.mkdir(parents=True, exist_ok=True)
|
||||
corpus_path = job_dir / "corpus.json"
|
||||
if corpus_path.exists():
|
||||
prompts = json.loads(corpus_path.read_text(encoding="utf-8"))
|
||||
else:
|
||||
prompts = _load_corpus(
|
||||
params.get("dataset", DATASET_WIKITEXT),
|
||||
params["n_prompts"],
|
||||
params.get("skip_prompts", 0),
|
||||
)
|
||||
corpus_path.write_text(
|
||||
json.dumps(prompts, ensure_ascii=False), encoding="utf-8"
|
||||
)
|
||||
devices = params["devices"]
|
||||
n = len(prompts)
|
||||
if len(devices) == 2:
|
||||
cut = int(n * 0.65)
|
||||
slices = [prompts[:cut], prompts[cut:]]
|
||||
else:
|
||||
slices = [prompts]
|
||||
|
||||
self.state.update(phase="fitting", total=n)
|
||||
workers = []
|
||||
started = time.perf_counter()
|
||||
for i, (device, chunk) in enumerate(zip(devices, slices)):
|
||||
slice_path = job_dir / f"slice{i}.json"
|
||||
if not slice_path.exists():
|
||||
slice_path.write_text(json.dumps(chunk, ensure_ascii=False), encoding="utf-8")
|
||||
dim_batch = params["dim_batch"] or _default_dim_batch(device)
|
||||
cmd = [
|
||||
sys.executable, "-X", "utf8", str(WORKER),
|
||||
"--model", params["source"],
|
||||
"--device", device,
|
||||
"--dtype", params["dtype"],
|
||||
"--prompts", str(slice_path),
|
||||
"--checkpoint", str(job_dir / f"ckpt{i}.pt"),
|
||||
"--out", str(job_dir / f"lens{i}.pt"),
|
||||
"--dim-batch", str(dim_batch),
|
||||
"--max-seq-len", str(params["max_seq_len"]),
|
||||
]
|
||||
if params["quant"]:
|
||||
cmd += ["--quant", params["quant"]]
|
||||
if params["source_layers"]:
|
||||
cmd += ["--source-layers", json.dumps(params["source_layers"])]
|
||||
proc = subprocess.Popen(
|
||||
cmd,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
text=True,
|
||||
encoding="utf-8",
|
||||
cwd=str(config.ROOT),
|
||||
)
|
||||
self._procs.append(proc)
|
||||
worker_state = {
|
||||
"device": device,
|
||||
"done": 0,
|
||||
"total": len(chunk),
|
||||
"dim_batch": dim_batch,
|
||||
"state": "loading",
|
||||
"elapsed": 0.0,
|
||||
# [done, elapsed] of the last 10 updates: the ETA follows the
|
||||
# RECENT pace (throughput can degrade mid-fit, e.g. VRAM
|
||||
# saturated — a global average would then freeze the ETA)
|
||||
"hist": [],
|
||||
}
|
||||
workers.append(worker_state)
|
||||
threading.Thread(
|
||||
target=self._read_worker, args=(proc, worker_state, started), daemon=True
|
||||
).start()
|
||||
self.state["workers"] = workers
|
||||
self._emit()
|
||||
|
||||
stderr_tails = [""] * len(self._procs)
|
||||
|
||||
def drain_err(index, proc):
|
||||
data = proc.stderr.read()
|
||||
stderr_tails[index] = (data or "")[-2000:]
|
||||
|
||||
drainers = [
|
||||
threading.Thread(target=drain_err, args=(i, p), daemon=True)
|
||||
for i, p in enumerate(self._procs)
|
||||
]
|
||||
for t in drainers:
|
||||
t.start()
|
||||
for proc in self._procs:
|
||||
proc.wait()
|
||||
for t in drainers:
|
||||
t.join()
|
||||
failed = [i for i, p in enumerate(self._procs) if p.returncode != 0]
|
||||
if self.state.get("state") == "stopping":
|
||||
self.state.update(state="stopped")
|
||||
self._emit()
|
||||
return
|
||||
if failed:
|
||||
detail = " | ".join(stderr_tails[i].strip().splitlines()[-1] if stderr_tails[i].strip() else "?" for i in failed)
|
||||
raise RuntimeError(f"worker(s) {failed} failed: {detail}")
|
||||
|
||||
self.state.update(phase="merge")
|
||||
self._emit()
|
||||
partials = [
|
||||
JacobianLens.load(str(job_dir / f"lens{i}.pt"))
|
||||
for i in range(len(slices))
|
||||
]
|
||||
merged = JacobianLens.merge(partials) if len(partials) > 1 else partials[0]
|
||||
if params.get("continue_from"):
|
||||
base_lens = JacobianLens.load(params["continue_from"])
|
||||
if base_lens.source_layers != merged.source_layers:
|
||||
raise RuntimeError(
|
||||
"cannot continue: the source layers differ from the base lens "
|
||||
f"({base_lens.source_layers[0]}..{base_lens.source_layers[-1]} vs "
|
||||
f"{merged.source_layers[0]}..{merged.source_layers[-1]})"
|
||||
)
|
||||
# weighted average by n_prompts = equivalent to a fit over the union
|
||||
merged = JacobianLens.merge([base_lens, merged])
|
||||
out_dir = config.LENSES_DIR / name
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
lens_path = out_dir / "lens.pt"
|
||||
merged.save(str(lens_path))
|
||||
meta = {
|
||||
"name": name,
|
||||
"model_id": params["model_id"],
|
||||
"model_revision": params["model_revision"],
|
||||
"model_source": params["source"],
|
||||
"d_model": merged.d_model,
|
||||
"source_layers": [merged.source_layers[0], merged.source_layers[-1]],
|
||||
"dtype": params["dtype"],
|
||||
"quant": params["quant"],
|
||||
"n_prompts": merged.n_prompts,
|
||||
"corpus": (
|
||||
f"mixed: {DATASET_WIKITEXT} + {DATASET_HARMLESS} (equal parts)"
|
||||
if params.get("dataset") == "mixed"
|
||||
else params.get("dataset", DATASET_WIKITEXT)
|
||||
),
|
||||
"max_seq_len": params["max_seq_len"],
|
||||
"devices": params["devices"],
|
||||
"continued_from": params.get("continue_from"),
|
||||
"config_hash": hashlib.sha1(
|
||||
json.dumps(params, sort_keys=True).encode()
|
||||
).hexdigest()[:16],
|
||||
"created_at": _now(),
|
||||
"fit_seconds": round(time.perf_counter() - started, 1),
|
||||
}
|
||||
(out_dir / "meta.json").write_text(
|
||||
json.dumps(meta, ensure_ascii=False, indent=1), encoding="utf-8"
|
||||
)
|
||||
self.state.update(
|
||||
state="done",
|
||||
phase="done",
|
||||
lens_path=str(lens_path),
|
||||
meta=meta,
|
||||
eta_seconds=0,
|
||||
)
|
||||
self._emit()
|
||||
except Exception as exc:
|
||||
self.state.update(state="error", error=str(exc))
|
||||
self._emit()
|
||||
|
||||
def _read_worker(self, proc, worker_state, started):
|
||||
for line in proc.stdout:
|
||||
line = line.strip()
|
||||
if not line:
|
||||
continue
|
||||
try:
|
||||
event = json.loads(line)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
if event["event"] == "loading":
|
||||
worker_state["state"] = "loading"
|
||||
elif event["event"] in ("progress", "resume"):
|
||||
worker_state["state"] = "fitting"
|
||||
worker_state["done"] = event["done"]
|
||||
worker_state["total"] = event["total"]
|
||||
worker_state["elapsed"] = round(time.perf_counter() - started, 1)
|
||||
hist = worker_state.setdefault("hist", [])
|
||||
hist.append([worker_state["done"], worker_state["elapsed"]])
|
||||
del hist[:-10]
|
||||
elif event["event"] == "done":
|
||||
worker_state["state"] = "done"
|
||||
worker_state["done"] = worker_state["total"]
|
||||
self._refresh_totals(started)
|
||||
self._emit()
|
||||
|
||||
def _refresh_totals(self, started):
|
||||
workers = self.state.get("workers", [])
|
||||
self.state["done"] = sum(w["done"] for w in workers)
|
||||
etas = []
|
||||
for w in workers:
|
||||
if not (w["done"] > 0 and w["elapsed"] > 0 and w["done"] < w["total"]):
|
||||
continue
|
||||
hist = w.get("hist") or []
|
||||
if len(hist) >= 2 and hist[-1][1] > hist[0][1] and hist[-1][0] > hist[0][0]:
|
||||
# pace over the last 10 updates (sliding window)
|
||||
rate = (hist[-1][0] - hist[0][0]) / (hist[-1][1] - hist[0][1])
|
||||
else:
|
||||
rate = w["done"] / w["elapsed"]
|
||||
etas.append((w["total"] - w["done"]) / rate)
|
||||
# multi-GPU: the fit ETA = the slowest worker
|
||||
self.state["eta_seconds"] = round(max(etas), 0) if etas else None
|
||||
self.state["vram"] = [
|
||||
{"index": g["index"], "used_gb": round(g["vram_used"] / 2**30, 1)}
|
||||
for g in gpu_stats()
|
||||
]
|
||||
@@ -0,0 +1,32 @@
|
||||
import pynvml
|
||||
|
||||
_initialized = False
|
||||
|
||||
|
||||
def _ensure_init():
|
||||
global _initialized
|
||||
if not _initialized:
|
||||
pynvml.nvmlInit()
|
||||
_initialized = True
|
||||
|
||||
|
||||
def gpu_stats():
|
||||
_ensure_init()
|
||||
stats = []
|
||||
for index in range(pynvml.nvmlDeviceGetCount()):
|
||||
handle = pynvml.nvmlDeviceGetHandleByIndex(index)
|
||||
memory = pynvml.nvmlDeviceGetMemoryInfo(handle)
|
||||
util = pynvml.nvmlDeviceGetUtilizationRates(handle)
|
||||
name = pynvml.nvmlDeviceGetName(handle)
|
||||
if isinstance(name, bytes):
|
||||
name = name.decode()
|
||||
stats.append(
|
||||
{
|
||||
"index": index,
|
||||
"name": name,
|
||||
"vram_total": memory.total,
|
||||
"vram_used": memory.used,
|
||||
"util_pct": util.gpu,
|
||||
}
|
||||
)
|
||||
return stats
|
||||
@@ -0,0 +1,353 @@
|
||||
import hashlib
|
||||
import itertools
|
||||
import json
|
||||
import threading
|
||||
from collections import OrderedDict
|
||||
|
||||
import torch
|
||||
from jlens.lens import JacobianLens
|
||||
|
||||
import config
|
||||
|
||||
GEN_STORE_MAX = 4
|
||||
|
||||
MASKS_DIR = config.DATA_DIR / "masks"
|
||||
|
||||
# Last range of layers captured per lens: {lens key: [layers]}.
|
||||
# Avoids re-entering the range on every reload (user request).
|
||||
LENS_PREFS_PATH = config.DATA_DIR / "lens_prefs.json"
|
||||
|
||||
|
||||
def _load_lens_prefs():
|
||||
try:
|
||||
return json.loads(LENS_PREFS_PATH.read_text(encoding="utf-8"))
|
||||
except Exception:
|
||||
return {}
|
||||
|
||||
|
||||
def _save_lens_pref(key, layers):
|
||||
prefs = _load_lens_prefs()
|
||||
prefs[key] = [int(l) for l in layers]
|
||||
LENS_PREFS_PATH.parent.mkdir(parents=True, exist_ok=True)
|
||||
LENS_PREFS_PATH.write_text(json.dumps(prefs, indent=1), encoding="utf-8")
|
||||
|
||||
|
||||
def _lens_pref_key(source):
|
||||
if source.get("path"):
|
||||
return f"path:{source['path']}"
|
||||
return f"hub:{source['repo_id']}:{source['filename']}@{source.get('revision') or 'main'}"
|
||||
|
||||
|
||||
class ActivationCatcher:
|
||||
def __init__(self, layers, indices):
|
||||
self.acts = {}
|
||||
self._handles = [
|
||||
layers[i].register_forward_hook(self._make(i)) for i in indices
|
||||
]
|
||||
|
||||
def _make(self, index):
|
||||
def hook(module, inputs, output):
|
||||
tensor = output[0] if isinstance(output, tuple) else output
|
||||
self.acts[index] = tensor.detach()
|
||||
|
||||
return hook
|
||||
|
||||
def close(self):
|
||||
for handle in self._handles:
|
||||
handle.remove()
|
||||
self._handles = []
|
||||
|
||||
|
||||
def _vocab_fingerprint(tokenizer):
|
||||
payload = json.dumps(sorted(tokenizer.get_vocab().items()), ensure_ascii=False)
|
||||
return hashlib.sha1(payload.encode()).hexdigest()[:16]
|
||||
|
||||
|
||||
def _wordlike(raw):
|
||||
s = raw.strip()
|
||||
if len(s) < 1 or "<|" in s or (s.startswith("<") and s.endswith(">")):
|
||||
return False
|
||||
if s.isascii():
|
||||
return (
|
||||
raw.startswith(" ")
|
||||
and len(s) > 2
|
||||
and s[0].isalpha()
|
||||
and all(c.isalpha() or c in "'-" for c in s)
|
||||
)
|
||||
return all(ch.isalnum() for ch in s)
|
||||
|
||||
|
||||
def display_token_mask(tokenizer, vocab_size):
|
||||
MASKS_DIR.mkdir(parents=True, exist_ok=True)
|
||||
path = MASKS_DIR / f"{_vocab_fingerprint(tokenizer)}_{vocab_size}.pt"
|
||||
if path.exists():
|
||||
return torch.load(path, weights_only=True)
|
||||
mask = torch.zeros(vocab_size, dtype=torch.bool)
|
||||
n_decodable = min(vocab_size, len(tokenizer))
|
||||
decoded = tokenizer.batch_decode(
|
||||
[[tid] for tid in range(n_decodable)], clean_up_tokenization_spaces=False
|
||||
)
|
||||
for tid, raw in enumerate(decoded):
|
||||
mask[tid] = _wordlike(raw)
|
||||
torch.save(mask, path)
|
||||
return mask
|
||||
|
||||
|
||||
class LensManager:
|
||||
def __init__(self):
|
||||
self._lock = threading.Lock()
|
||||
self.lens = None
|
||||
self.meta = None
|
||||
self.layers = []
|
||||
self.k = 8
|
||||
self.mask = None
|
||||
self._J = None
|
||||
self._tok_strs = {}
|
||||
self.gen_store = OrderedDict()
|
||||
self._gen_counter = itertools.count(1)
|
||||
self._pref_key = None
|
||||
|
||||
def load(self, model_manager, *, repo_id=None, filename="lens.pt", revision=None,
|
||||
path=None, layers=None, k=8):
|
||||
with self._lock:
|
||||
if model_manager.hf_model is None:
|
||||
raise ValueError("load a model first")
|
||||
if path:
|
||||
lens = JacobianLens.from_pretrained(path)
|
||||
source = {"path": path, "repo_id": None, "filename": None, "revision": None}
|
||||
else:
|
||||
lens = JacobianLens.from_pretrained(
|
||||
repo_id, filename=filename, revision=revision
|
||||
)
|
||||
source = {"path": None, "repo_id": repo_id, "filename": filename, "revision": revision}
|
||||
|
||||
model_meta = model_manager.meta
|
||||
if lens.d_model != model_meta["d_model"]:
|
||||
raise ValueError(
|
||||
f"lens d_model ({lens.d_model}) != model ({model_meta['d_model']})"
|
||||
)
|
||||
n_layers = model_meta["n_layers"]
|
||||
fitted = lens.source_layers
|
||||
if fitted[-1] >= n_layers:
|
||||
raise ValueError(
|
||||
f"the lens covers layer {fitted[-1]}, outside a model with {n_layers} layers"
|
||||
)
|
||||
pref_key = _lens_pref_key(source)
|
||||
if layers:
|
||||
tapped = sorted(set(layers) & set(fitted))
|
||||
if not tapped:
|
||||
raise ValueError(
|
||||
f"no requested layer is fitted (fitted: {fitted[0]}..{fitted[-1]})"
|
||||
)
|
||||
else:
|
||||
# last range used for THIS lens, otherwise all the fitted layers
|
||||
# (= the selection made when the fit was created; max range for a
|
||||
# downloaded lens)
|
||||
saved = _load_lens_prefs().get(pref_key)
|
||||
tapped = (sorted(set(saved) & set(fitted)) if saved else None) or list(fitted)
|
||||
_save_lens_pref(pref_key, tapped)
|
||||
self._pref_key = pref_key
|
||||
|
||||
device = model_manager.jl.input_device
|
||||
stacked = torch.stack([lens.jacobians[l].float() for l in tapped]).to(device)
|
||||
tokenizer = model_manager.tokenizer
|
||||
vocab_size = model_manager.hf_model.get_output_embeddings().weight.shape[0]
|
||||
mask = display_token_mask(tokenizer, vocab_size).to(device)
|
||||
|
||||
warnings = []
|
||||
if model_meta.get("quant"):
|
||||
warnings.append(
|
||||
f"model loaded in {model_meta['quant']}: the lens was probably "
|
||||
"fitted on the unquantized weights, the readouts may drift"
|
||||
)
|
||||
if model_meta["model_id"].startswith("local/"):
|
||||
warnings.append(
|
||||
"local model: cannot verify that the lens matches these exact weights"
|
||||
)
|
||||
|
||||
self.lens = lens
|
||||
self.layers = tapped
|
||||
self.k = int(k)
|
||||
self.mask = mask
|
||||
self._J = stacked
|
||||
self._tok_strs = {}
|
||||
self.meta = {
|
||||
**source,
|
||||
"model_id": model_meta["model_id"],
|
||||
"model_revision": model_meta.get("revision"),
|
||||
"d_model": lens.d_model,
|
||||
"n_prompts": lens.n_prompts,
|
||||
"fitted_layers": [int(fitted[0]), int(fitted[-1])],
|
||||
"fitted_layers_all": [int(l) for l in fitted],
|
||||
"tapped_layers": [int(l) for l in tapped],
|
||||
"k": self.k,
|
||||
"warnings": warnings,
|
||||
}
|
||||
return self.meta
|
||||
|
||||
def set_layers(self, model_manager, layers, k=None):
|
||||
with self._lock:
|
||||
if self.lens is None:
|
||||
raise ValueError("no lens loaded")
|
||||
fitted = self.lens.source_layers
|
||||
tapped = sorted(set(layers) & set(fitted))
|
||||
if not tapped:
|
||||
raise ValueError(
|
||||
f"no requested layer is fitted (fitted: {fitted[0]}..{fitted[-1]})"
|
||||
)
|
||||
device = model_manager.jl.input_device
|
||||
self.layers = tapped
|
||||
self._J = torch.stack(
|
||||
[self.lens.jacobians[l].float() for l in tapped]
|
||||
).to(device)
|
||||
if k:
|
||||
self.k = int(k)
|
||||
self.meta = dict(self.meta, tapped_layers=[int(l) for l in tapped], k=self.k)
|
||||
if getattr(self, "_pref_key", None):
|
||||
_save_lens_pref(self._pref_key, tapped)
|
||||
return self.meta
|
||||
|
||||
def unload(self):
|
||||
with self._lock:
|
||||
self.lens = None
|
||||
self.meta = None
|
||||
self.layers = []
|
||||
self.mask = None
|
||||
self._J = None
|
||||
self._tok_strs = {}
|
||||
self.gen_store.clear()
|
||||
torch.cuda.empty_cache()
|
||||
return {"unloaded": True}
|
||||
|
||||
def start_gen(self):
|
||||
gen_id = next(self._gen_counter)
|
||||
self.gen_store[gen_id] = {
|
||||
"layers": list(self.layers),
|
||||
"residuals": {l: [] for l in self.layers},
|
||||
"positions": [],
|
||||
"token_ids": [],
|
||||
"phases": [],
|
||||
}
|
||||
while len(self.gen_store) > GEN_STORE_MAX:
|
||||
self.gen_store.popitem(last=False)
|
||||
return gen_id
|
||||
|
||||
@torch.no_grad()
|
||||
def pin_ranks(self, gen_id, token_ids, jl, chunk=32):
|
||||
store = self.gen_store.get(gen_id)
|
||||
if store is None:
|
||||
raise ValueError("unknown generation (residual store expired)")
|
||||
layers = store["layers"]
|
||||
device = self._J.device
|
||||
tids = torch.tensor(token_ids, dtype=torch.long, device=device)
|
||||
pins = {
|
||||
int(t): {"ranks": [], "p": []} for t in token_ids
|
||||
}
|
||||
for layer in layers:
|
||||
residuals = torch.cat(store["residuals"][layer]).to(device).float()
|
||||
J = self.lens.jacobians[layer].float().to(device)
|
||||
layer_ranks = {int(t): [] for t in token_ids}
|
||||
layer_p = {int(t): [] for t in token_ids}
|
||||
for start in range(0, residuals.shape[0], chunk):
|
||||
h = residuals[start : start + chunk]
|
||||
logits = jl.unembed(h @ J.T).float()
|
||||
probs = torch.softmax(logits, -1)
|
||||
sel = logits[:, tids]
|
||||
rank = (logits.unsqueeze(-1) > sel.unsqueeze(1)).sum(1)
|
||||
p_sel = probs[:, tids]
|
||||
rank_l, p_l = rank.tolist(), p_sel.tolist()
|
||||
for ti, t in enumerate(token_ids):
|
||||
layer_ranks[int(t)].extend(row[ti] for row in rank_l)
|
||||
layer_p[int(t)].extend(round(row[ti], 6) for row in p_l)
|
||||
for t in token_ids:
|
||||
pins[int(t)]["ranks"].append(layer_ranks[int(t)])
|
||||
pins[int(t)]["p"].append(layer_p[int(t)])
|
||||
return {
|
||||
"gen_id": gen_id,
|
||||
"layers": [int(l) for l in layers],
|
||||
"positions": store["positions"],
|
||||
"phases": store["phases"],
|
||||
"tokens": self._strs(jl.tokenizer, store["token_ids"]),
|
||||
"pins": pins,
|
||||
}
|
||||
|
||||
def _strs(self, tokenizer, ids):
|
||||
out = []
|
||||
for tid in ids:
|
||||
s = self._tok_strs.get(tid)
|
||||
if s is None:
|
||||
s = tokenizer.decode([tid], clean_up_tokenization_spaces=False)
|
||||
self._tok_strs[tid] = s
|
||||
out.append(s)
|
||||
return out
|
||||
|
||||
@torch.no_grad()
|
||||
def compute_frames(self, acts, positions, phase, jl, token_ids, gen_id=None,
|
||||
abs_positions=None, chunk=None):
|
||||
tokenizer = jl.tokenizer
|
||||
if chunk is None:
|
||||
chunk = max(1, 96 // max(1, len(self.layers)))
|
||||
if abs_positions is None:
|
||||
abs_positions = positions
|
||||
frames = [
|
||||
{
|
||||
"type": "frame",
|
||||
"phase": phase,
|
||||
"pos": int(pos),
|
||||
"token_id": int(tid),
|
||||
"tok": self._strs(tokenizer, [tid])[0],
|
||||
"gen": gen_id,
|
||||
"layers": {},
|
||||
}
|
||||
for pos, tid in zip(abs_positions, token_ids)
|
||||
]
|
||||
store = self.gen_store.get(gen_id) if gen_id is not None else None
|
||||
if store is not None:
|
||||
store["positions"].extend(int(p) for p in abs_positions)
|
||||
store["token_ids"].extend(int(t) for t in token_ids)
|
||||
store["phases"].extend(phase for _ in abs_positions)
|
||||
device = self._J.device
|
||||
for start in range(0, len(positions), chunk):
|
||||
batch_positions = positions[start : start + chunk]
|
||||
gathered = []
|
||||
for layer in self.layers:
|
||||
full = acts[layer][0]
|
||||
gathered.append(full[list(batch_positions)].float().to(device))
|
||||
h = torch.stack(gathered)
|
||||
if store is not None:
|
||||
for li, layer in enumerate(self.layers):
|
||||
store["residuals"][layer].append(h[li].half().cpu())
|
||||
# L2 norm of the residual per layer/position ("Activations" view)
|
||||
h_norms = h.norm(dim=-1).tolist()
|
||||
transported = torch.einsum("lij,lpj->lpi", self._J, h)
|
||||
logits = jl.unembed(transported).float()
|
||||
lse = logits.logsumexp(-1, keepdim=True)
|
||||
raw_v, raw_ids = logits.topk(self.k)
|
||||
raw_p = (raw_v - lse).exp()
|
||||
m_v, m_ids = logits.masked_fill(~self.mask, float("-inf")).topk(self.k)
|
||||
m_p = (m_v - lse).exp()
|
||||
sel = logits.gather(-1, m_ids)
|
||||
# rank of each top-k token in the full distribution. We loop over k
|
||||
# rather than materializing a boolean [L, P, k, V] (≈760 MB at k=32 /
|
||||
# 32 layers → OOM): each iteration only touches [L, P, V].
|
||||
m_rank = torch.empty_like(m_ids)
|
||||
for ki in range(m_ids.shape[-1]):
|
||||
m_rank[..., ki] = (logits > sel[..., ki : ki + 1]).sum(-1)
|
||||
del logits
|
||||
raw_ids_l, raw_p_l = raw_ids.tolist(), raw_p.tolist()
|
||||
m_ids_l, m_p_l, m_rank_l = m_ids.tolist(), m_p.tolist(), m_rank.tolist()
|
||||
for li, layer in enumerate(self.layers):
|
||||
for pi in range(len(batch_positions)):
|
||||
ids = raw_ids_l[li][pi]
|
||||
mids = m_ids_l[li][pi]
|
||||
frames[start + pi]["layers"][str(layer)] = {
|
||||
"ids": ids,
|
||||
"p": [round(v, 5) for v in raw_p_l[li][pi]],
|
||||
"strs": self._strs(tokenizer, ids),
|
||||
"m_ids": mids,
|
||||
"m_p": [round(v, 5) for v in m_p_l[li][pi]],
|
||||
"m_rank": m_rank_l[li][pi],
|
||||
"m_strs": self._strs(tokenizer, mids),
|
||||
"h_norm": round(h_norms[li][pi], 2),
|
||||
}
|
||||
return frames
|
||||
@@ -0,0 +1,761 @@
|
||||
import gc
|
||||
import json
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
import transformers
|
||||
from huggingface_hub import scan_cache_dir, try_to_load_from_cache
|
||||
|
||||
import config
|
||||
import jlens
|
||||
from core.lens_manager import ActivationCatcher
|
||||
|
||||
SKIP_LOCAL_DIRS = {"vendor", "ui", "data", "hf_cache", "lenses", "core", "api", "scripts"}
|
||||
|
||||
# Many "base" models (e.g. non-Instruct Llama-3.2-1B) ship no chat_template. The
|
||||
# right one is their instruct sibling's, which shares the same tokenizer: so we
|
||||
# look it up on the Hub before any fallback.
|
||||
INSTRUCT_SIBLING_SUFFIXES = ("-Instruct", "-instruct", "-it", "-Chat", "-chat")
|
||||
|
||||
# End-of-turn markers per model family; added to the stop tokens when they appear
|
||||
# in the applied template (useful when a base model is given an instruct template:
|
||||
# it must stop on <|eot_id|>, <|im_end|>, <end_of_turn>, etc.)
|
||||
TURN_END_MARKERS = ("<|eot_id|>", "<|im_end|>", "<end_of_turn>", "<|end|>", "<|endoftext|>")
|
||||
|
||||
# Last-resort fallback when no template can be found (offline, no reachable
|
||||
# sibling): a readable "User:/Assistant:" format a completion model can continue.
|
||||
FALLBACK_CHAT_TEMPLATE = (
|
||||
"{% for message in messages %}"
|
||||
"{% if message['role'] == 'system' %}{{ message['content'] + '\n\n' }}"
|
||||
"{% elif message['role'] == 'user' %}{{ 'User: ' + message['content'] + '\n' }}"
|
||||
"{% elif message['role'] == 'assistant' %}{{ 'Assistant: ' + message['content'] + '\n' }}"
|
||||
"{% endif %}{% endfor %}"
|
||||
"{% if add_generation_prompt %}{{ 'Assistant:' }}{% endif %}"
|
||||
)
|
||||
|
||||
|
||||
def _extract_template(chat_template):
|
||||
"""chat_template may be a string or a list [{name, template}] (multi-template)."""
|
||||
if isinstance(chat_template, str):
|
||||
return chat_template
|
||||
if isinstance(chat_template, list):
|
||||
for entry in chat_template:
|
||||
if isinstance(entry, dict) and entry.get("name") == "default":
|
||||
return entry.get("template")
|
||||
if chat_template and isinstance(chat_template[0], dict):
|
||||
return chat_template[0].get("template")
|
||||
return None
|
||||
|
||||
|
||||
def _read_hub_template(repo, token, revision=None):
|
||||
"""Read a chat_template from a Hub repo: chat_template.jinja (raw) then the
|
||||
chat_template key of tokenizer_config.json / chat_template.json."""
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
try:
|
||||
path = hf_hub_download(repo, "chat_template.jinja", token=token, revision=revision)
|
||||
text = Path(path).read_text(encoding="utf-8").strip()
|
||||
if text:
|
||||
return text
|
||||
except Exception:
|
||||
pass
|
||||
for fname in ("tokenizer_config.json", "chat_template.json"):
|
||||
try:
|
||||
path = hf_hub_download(repo, fname, token=token, revision=revision)
|
||||
data = json.loads(Path(path).read_text(encoding="utf-8"))
|
||||
except Exception:
|
||||
continue
|
||||
tmpl = _extract_template(data.get("chat_template") if isinstance(data, dict) else None)
|
||||
if tmpl:
|
||||
return tmpl
|
||||
return None
|
||||
|
||||
|
||||
def fetch_chat_template(model_id, token, revision=None):
|
||||
"""Look up the real chat_template on the Hub: first the model's own repo, then
|
||||
its instruct siblings (shared tokenizer). Returns (template, source_repo) or
|
||||
(None, None). Skips local models/paths (no Hub repo)."""
|
||||
if "/" not in model_id or model_id.startswith("local/") or os.path.isabs(model_id):
|
||||
return None, None
|
||||
candidates = [(model_id, revision)]
|
||||
for suffix in INSTRUCT_SIBLING_SUFFIXES:
|
||||
if not model_id.endswith(suffix):
|
||||
candidates.append((model_id + suffix, None)) # sibling revision unknown
|
||||
for repo, rev in candidates:
|
||||
try:
|
||||
tmpl = _read_hub_template(repo, token, revision=rev)
|
||||
except Exception:
|
||||
tmpl = None
|
||||
if tmpl:
|
||||
return tmpl, repo
|
||||
return None, None
|
||||
|
||||
|
||||
def _config_n_layers(config_path):
|
||||
try:
|
||||
cfg = json.loads(Path(config_path).read_text(encoding="utf-8"))
|
||||
except Exception:
|
||||
return None
|
||||
tc = cfg.get("text_config", cfg)
|
||||
return tc.get("num_hidden_layers") or cfg.get("num_hidden_layers")
|
||||
|
||||
|
||||
def _config_dtype(config_path):
|
||||
try:
|
||||
cfg = json.loads(Path(config_path).read_text(encoding="utf-8"))
|
||||
except Exception:
|
||||
return None
|
||||
return cfg.get("torch_dtype") or cfg.get("text_config", {}).get("torch_dtype")
|
||||
|
||||
|
||||
# Model folders registered by hand (Browse): a list of absolute paths kept in
|
||||
# the data dir. Registering never copies or moves anything; unregistering only
|
||||
# forgets the entry, the files stay untouched.
|
||||
REGISTERED_PATH = config.DATA_DIR / "registered_models.json"
|
||||
|
||||
|
||||
def _read_registered():
|
||||
try:
|
||||
entries = json.loads(REGISTERED_PATH.read_text(encoding="utf-8"))
|
||||
return [str(e) for e in entries if isinstance(e, str)]
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
|
||||
def _write_registered(entries):
|
||||
REGISTERED_PATH.parent.mkdir(parents=True, exist_ok=True)
|
||||
REGISTERED_PATH.write_text(
|
||||
json.dumps(entries, ensure_ascii=False, indent=1), encoding="utf-8"
|
||||
)
|
||||
|
||||
|
||||
def register_model_dir(path):
|
||||
p = Path(path).expanduser().resolve()
|
||||
if not _dir_is_model(p):
|
||||
raise ValueError(f"not a model folder (config.json + weights required): {p}")
|
||||
entries = _read_registered()
|
||||
if str(p) not in entries:
|
||||
entries.append(str(p))
|
||||
_write_registered(entries)
|
||||
return {"registered": str(p)}
|
||||
|
||||
|
||||
def unregister_model_dir(path):
|
||||
wanted = str(Path(path).expanduser().resolve())
|
||||
entries = _read_registered()
|
||||
kept = [e for e in entries if e != path and str(Path(e)) != wanted]
|
||||
if len(kept) == len(entries):
|
||||
raise ValueError(f"not a registered entry: {path}")
|
||||
_write_registered(kept)
|
||||
return {"unregistered": path}
|
||||
|
||||
|
||||
def _registered_models():
|
||||
out = []
|
||||
for entry in _read_registered():
|
||||
path = Path(entry)
|
||||
missing = not _dir_is_model(path)
|
||||
stats = [] if missing else [f.stat() for f in path.glob("*.safetensors")]
|
||||
unique = {(s.st_ino, s.st_size): s.st_size for s in stats}
|
||||
out.append(
|
||||
{
|
||||
# the absolute path IS the id: resolve_source passes it through
|
||||
"id": entry,
|
||||
"source": "registered",
|
||||
"path": entry,
|
||||
"missing": missing,
|
||||
"size_bytes": sum(unique.values()),
|
||||
"n_layers": None if missing else _config_n_layers(path / "config.json"),
|
||||
"dtype": None if missing else _config_dtype(path / "config.json"),
|
||||
}
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
def _local_models():
|
||||
found = []
|
||||
for child in sorted(config.LOCAL_MODELS_ROOT.iterdir()):
|
||||
if not child.is_dir() or child.name in SKIP_LOCAL_DIRS:
|
||||
continue
|
||||
if not (child / "config.json").exists():
|
||||
continue
|
||||
stats = [f.stat() for f in child.glob("*.safetensors")]
|
||||
if not stats:
|
||||
continue
|
||||
unique = {(s.st_ino, s.st_size): s.st_size for s in stats}
|
||||
found.append(
|
||||
{
|
||||
"id": f"local/{child.name}",
|
||||
"source": "local",
|
||||
"path": str(child),
|
||||
"size_bytes": sum(unique.values()),
|
||||
"n_layers": _config_n_layers(child / "config.json"),
|
||||
"dtype": _config_dtype(child / "config.json"),
|
||||
}
|
||||
)
|
||||
return found
|
||||
|
||||
|
||||
def _cached_models():
|
||||
hub = config.HF_CACHE / "hub"
|
||||
if not hub.exists():
|
||||
return []
|
||||
out = []
|
||||
for repo in scan_cache_dir(hub).repos:
|
||||
if repo.repo_type != "model":
|
||||
continue
|
||||
config_path = None
|
||||
for rev in repo.revisions:
|
||||
for f in rev.files:
|
||||
if f.file_name == "config.json":
|
||||
config_path = f.file_path
|
||||
if config_path is None:
|
||||
continue
|
||||
out.append(
|
||||
{
|
||||
"id": repo.repo_id,
|
||||
"source": "hf-cache",
|
||||
"size_bytes": repo.size_on_disk,
|
||||
"n_layers": _config_n_layers(config_path),
|
||||
"dtype": _config_dtype(config_path),
|
||||
"path": str(Path(config_path).parent),
|
||||
}
|
||||
)
|
||||
return sorted(out, key=lambda r: r["id"])
|
||||
|
||||
|
||||
def _dir_is_model(path):
|
||||
return (path / "config.json").exists() and (
|
||||
any(path.glob("*.safetensors")) or any(path.glob("*.bin"))
|
||||
)
|
||||
|
||||
|
||||
def browse_dir(path=None):
|
||||
"""Minimal file browser: subfolders + loadable model folders.
|
||||
Empty path -> list drive letters (Windows)."""
|
||||
import string
|
||||
|
||||
if not path:
|
||||
drives = []
|
||||
for letter in string.ascii_uppercase:
|
||||
root = Path(f"{letter}:/")
|
||||
if root.exists():
|
||||
drives.append({"name": f"{letter}:", "path": str(root)})
|
||||
return {"path": "", "parent": None, "dirs": drives, "models": []}
|
||||
|
||||
base = Path(path)
|
||||
if not base.is_dir():
|
||||
raise ValueError(f"folder not found: {path}")
|
||||
dirs = []
|
||||
try:
|
||||
children = sorted(base.iterdir(), key=lambda p: p.name.lower())
|
||||
except PermissionError:
|
||||
children = []
|
||||
for child in children:
|
||||
try:
|
||||
if child.is_dir():
|
||||
dirs.append({
|
||||
"name": child.name,
|
||||
"path": str(child),
|
||||
"is_model": _dir_is_model(child),
|
||||
})
|
||||
except OSError:
|
||||
continue
|
||||
return {
|
||||
"path": str(base),
|
||||
"parent": str(base.parent) if base.parent != base else None,
|
||||
"dirs": dirs,
|
||||
"is_model": _dir_is_model(base),
|
||||
}
|
||||
|
||||
|
||||
def delete_model(model_id):
|
||||
"""Delete a model: local folder or HF cache repo.
|
||||
Refuses anything outside the managed roots (guards against arbitrary paths)."""
|
||||
import shutil
|
||||
|
||||
if model_id.startswith("local/"):
|
||||
name = model_id.removeprefix("local/")
|
||||
if name in SKIP_LOCAL_DIRS or "/" in name or "\\" in name or ".." in name:
|
||||
raise ValueError("protected folder or invalid name")
|
||||
path = (config.LOCAL_MODELS_ROOT / name).resolve()
|
||||
root = config.LOCAL_MODELS_ROOT.resolve()
|
||||
if root not in path.parents or not (path / "config.json").exists():
|
||||
raise ValueError(f"unmanaged path: {path}")
|
||||
shutil.rmtree(path)
|
||||
return {"deleted": str(path), "freed_bytes": None}
|
||||
|
||||
# otherwise: a Hugging Face cache repo (delete all of its revisions)
|
||||
hub = config.HF_CACHE / "hub"
|
||||
if not hub.exists():
|
||||
raise ValueError(f"unknown model: {model_id}")
|
||||
info = scan_cache_dir(hub)
|
||||
hashes, freed = [], 0
|
||||
for repo in info.repos:
|
||||
if repo.repo_id == model_id and repo.repo_type == "model":
|
||||
hashes = [rev.commit_hash for rev in repo.revisions]
|
||||
freed = repo.size_on_disk
|
||||
break
|
||||
if not hashes:
|
||||
raise ValueError(f"unknown model in cache: {model_id}")
|
||||
info.delete_revisions(*hashes).execute()
|
||||
return {"deleted": model_id, "freed_bytes": freed}
|
||||
|
||||
|
||||
def convert_to_bf16(src_dir, out_dir=None):
|
||||
"""Rewrite an fp32 model's safetensors as bf16 into a sibling local folder.
|
||||
Leaves the source untouched. Returns the new local id."""
|
||||
from safetensors import safe_open
|
||||
from safetensors.torch import save_file
|
||||
|
||||
src = Path(src_dir)
|
||||
if not src.is_dir():
|
||||
raise ValueError(f"source not found: {src_dir}")
|
||||
shards = sorted(src.glob("*.safetensors"))
|
||||
if not shards:
|
||||
raise ValueError("no safetensors in the source")
|
||||
out = Path(out_dir) if out_dir else (config.LOCAL_MODELS_ROOT / f"{src.name}-bf16")
|
||||
name = out.name
|
||||
out.mkdir(parents=True, exist_ok=True)
|
||||
for shard in shards:
|
||||
tensors = {}
|
||||
with safe_open(str(shard), framework="pt") as f:
|
||||
metadata = f.metadata()
|
||||
for key in f.keys():
|
||||
t = f.get_tensor(key)
|
||||
if t.dtype == torch.float32:
|
||||
t = t.to(torch.bfloat16)
|
||||
tensors[key] = t
|
||||
save_file(tensors, str(out / shard.name), metadata=metadata)
|
||||
for extra in src.iterdir():
|
||||
if extra.suffix in (".json", ".txt", ".model") or extra.name.startswith("tokenizer"):
|
||||
data = extra.read_bytes()
|
||||
if extra.name == "config.json":
|
||||
cfg = json.loads(data)
|
||||
cfg["torch_dtype"] = "bfloat16"
|
||||
if "text_config" in cfg and isinstance(cfg["text_config"], dict):
|
||||
cfg["text_config"]["torch_dtype"] = "bfloat16"
|
||||
(out / extra.name).write_text(json.dumps(cfg, indent=2), encoding="utf-8")
|
||||
else:
|
||||
(out / extra.name).write_bytes(data)
|
||||
return {"id": f"local/{name}", "path": str(out)}
|
||||
|
||||
|
||||
def _torch_allocated():
|
||||
return {
|
||||
f"cuda:{i}": torch.cuda.memory_allocated(i)
|
||||
for i in range(torch.cuda.device_count())
|
||||
}
|
||||
|
||||
|
||||
def _torch_reserved():
|
||||
return {
|
||||
f"cuda:{i}": torch.cuda.memory_reserved(i)
|
||||
for i in range(torch.cuda.device_count())
|
||||
}
|
||||
|
||||
|
||||
def _free_cuda():
|
||||
"""Hand the caching allocator's blocks back to the driver (gc then empty_cache).
|
||||
Call this on EVERY error/unload path: without it, allocations from an OOM load
|
||||
or from an aborted generation's KV cache stay reserved and pile up until the
|
||||
server restarts. We loop per device with a sync: pending frees must be visible
|
||||
before empty_cache can hand the segments back."""
|
||||
for _ in range(2):
|
||||
gc.collect()
|
||||
if not torch.cuda.is_available():
|
||||
return
|
||||
# cuBLAS keeps a persistent workspace (~8 MB) per device; under
|
||||
# expandable_segments:True (see config.setup_env) that single live allocation
|
||||
# pins the WHOLE segment (~8 GB) → empty_cache returns nothing after unload. So
|
||||
# we explicitly clear the cuBLAS workspaces first.
|
||||
try:
|
||||
torch._C._cuda_clearCublasWorkspaces()
|
||||
except Exception:
|
||||
pass
|
||||
for i in range(torch.cuda.device_count()):
|
||||
with torch.cuda.device(i):
|
||||
torch.cuda.synchronize()
|
||||
torch.cuda.empty_cache()
|
||||
try:
|
||||
torch.cuda.ipc_collect()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def _input_device(hf_model):
|
||||
return hf_model.get_input_embeddings().weight.device
|
||||
|
||||
|
||||
def resolve_source(model_id):
|
||||
if model_id.startswith("local/"):
|
||||
return str(config.LOCAL_MODELS_ROOT / model_id.removeprefix("local/"))
|
||||
return model_id
|
||||
|
||||
|
||||
def resolve_local_dir(model_id):
|
||||
source = resolve_source(model_id)
|
||||
path = Path(source)
|
||||
if path.is_dir():
|
||||
return str(path)
|
||||
cached = try_to_load_from_cache(source, "config.json")
|
||||
if isinstance(cached, str):
|
||||
return str(Path(cached).parent)
|
||||
return None
|
||||
|
||||
|
||||
def _resolve_revision(source):
|
||||
if Path(source).exists():
|
||||
return None
|
||||
cached = try_to_load_from_cache(source, "config.json")
|
||||
if isinstance(cached, str):
|
||||
parts = Path(cached).parts
|
||||
if "snapshots" in parts:
|
||||
return parts[parts.index("snapshots") + 1]
|
||||
return None
|
||||
|
||||
|
||||
def _sample(logits, temperature, top_p, top_k, generator=None):
|
||||
if temperature <= 0:
|
||||
return int(logits.argmax())
|
||||
probs = torch.softmax(logits / temperature, -1)
|
||||
if top_k > 0:
|
||||
kth = probs.topk(top_k).values[-1]
|
||||
probs = probs.masked_fill(probs < kth, 0.0)
|
||||
if 0 < top_p < 1:
|
||||
sorted_probs, sorted_idx = probs.sort(descending=True)
|
||||
keep = sorted_probs.cumsum(-1) - sorted_probs < top_p
|
||||
sorted_probs = sorted_probs * keep
|
||||
probs = torch.zeros_like(probs).scatter_(0, sorted_idx, sorted_probs)
|
||||
return int(torch.multinomial(probs / probs.sum(), 1, generator=generator))
|
||||
|
||||
|
||||
class ModelManager:
|
||||
def __init__(self):
|
||||
self._lock = threading.Lock()
|
||||
self.hf_model = None
|
||||
self.tokenizer = None
|
||||
self.jl = None
|
||||
self.meta = None
|
||||
self.busy = None
|
||||
|
||||
def list_models(self):
|
||||
return _local_models() + _registered_models() + _cached_models()
|
||||
|
||||
def load(self, model_id, dtype, quant, device):
|
||||
with self._lock:
|
||||
self._unload_locked()
|
||||
self.busy = "loading"
|
||||
hf_model = tokenizer = None
|
||||
try:
|
||||
torch_dtype = torch.bfloat16 if dtype == "bf16" else torch.float16
|
||||
source = resolve_source(model_id)
|
||||
kwargs = {"dtype": torch_dtype}
|
||||
if quant == "int8":
|
||||
kwargs["quantization_config"] = transformers.BitsAndBytesConfig(
|
||||
load_in_8bit=True
|
||||
)
|
||||
elif 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,
|
||||
)
|
||||
kwargs["device_map"] = "auto" if device == "auto" else {"": device}
|
||||
# model already present (HF cache or local folder) → load WITHOUT network:
|
||||
# otherwise from_pretrained queries the Hub and fails offline, even if cached.
|
||||
offline_ok = resolve_local_dir(model_id) is not None
|
||||
if offline_ok:
|
||||
kwargs["local_files_only"] = True
|
||||
tok_kwargs = {"local_files_only": True} if offline_ok else {}
|
||||
started = time.perf_counter()
|
||||
hf_model = transformers.AutoModelForCausalLM.from_pretrained(source, **kwargs)
|
||||
tokenizer = transformers.AutoTokenizer.from_pretrained(source, **tok_kwargs)
|
||||
# "base" models with no chat template: we fetch the real template from
|
||||
# the Hub (instruct sibling with shared tokenizer), generic as a last resort
|
||||
chat_template_source = None
|
||||
if not getattr(tokenizer, "chat_template", None):
|
||||
token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
|
||||
fetched, src = fetch_chat_template(
|
||||
model_id, token, revision=_resolve_revision(source)
|
||||
)
|
||||
if fetched:
|
||||
tokenizer.chat_template = fetched
|
||||
chat_template_source = src
|
||||
else:
|
||||
tokenizer.chat_template = FALLBACK_CHAT_TEMPLATE
|
||||
chat_template_source = "generic"
|
||||
chat_template_fallback = chat_template_source == "generic"
|
||||
hf_model.eval()
|
||||
text_config = hf_model.config.get_text_config()
|
||||
self.hf_model = hf_model
|
||||
self.tokenizer = tokenizer
|
||||
self.jl = jlens.from_hf(hf_model, tokenizer)
|
||||
# Read-projection support: write-norm architectures (Gemma
|
||||
# style) can't take the reads change of basis — the UI falls
|
||||
# back to the global abliteration for pure-weights edits.
|
||||
from core import rebase
|
||||
try:
|
||||
for block in self.jl.layers:
|
||||
rebase.check_block_supported(block)
|
||||
rebase_supported = True
|
||||
except ValueError:
|
||||
rebase_supported = False
|
||||
self.meta = {
|
||||
"model_id": model_id,
|
||||
"revision": _resolve_revision(source),
|
||||
"dtype": dtype,
|
||||
"quant": quant,
|
||||
"device": device,
|
||||
"n_layers": text_config.num_hidden_layers,
|
||||
"d_model": text_config.hidden_size,
|
||||
"rebase_supported": rebase_supported,
|
||||
"chat_template_source": chat_template_source,
|
||||
"chat_template_fallback": chat_template_fallback,
|
||||
"load_seconds": round(time.perf_counter() - started, 1),
|
||||
}
|
||||
return self.meta
|
||||
except Exception:
|
||||
# failure (often OOM): drop any partial allocation and return the
|
||||
# reserved blocks, otherwise they linger until the server restarts
|
||||
self.hf_model = self.tokenizer = self.jl = self.meta = None
|
||||
hf_model = None
|
||||
tokenizer = None
|
||||
_free_cuda()
|
||||
raise
|
||||
finally:
|
||||
self.busy = None
|
||||
|
||||
def unload(self):
|
||||
with self._lock:
|
||||
return self._unload_locked()
|
||||
|
||||
def _unload_locked(self):
|
||||
if self.hf_model is None:
|
||||
return {"unloaded": False, "vram_allocated": _torch_allocated()}
|
||||
before = _torch_allocated()
|
||||
self.hf_model = None
|
||||
self.tokenizer = None
|
||||
self.jl = None
|
||||
self.meta = None
|
||||
_free_cuda()
|
||||
return {
|
||||
"unloaded": True,
|
||||
"vram_allocated_before": before,
|
||||
"vram_allocated_after": _torch_allocated(),
|
||||
"vram_reserved_after": _torch_reserved(),
|
||||
}
|
||||
|
||||
@torch.no_grad()
|
||||
def generate(self, messages, sampling, stop_event, emit, lens=None, ablator=None,
|
||||
continue_final=False):
|
||||
"""``continue_final=True``: the last message is an assistant reply to
|
||||
EXTEND — the template leaves its turn open instead of starting a new
|
||||
one, and the model picks up where it stopped."""
|
||||
hf_model, tokenizer = self.hf_model, self.tokenizer
|
||||
self.busy = "generating"
|
||||
reader = None
|
||||
ok = False
|
||||
try:
|
||||
if ablator is not None:
|
||||
ablator.attach(self.jl)
|
||||
is_gpt_oss = "gpt-oss" in (self.meta or {}).get("model_id", "").lower()
|
||||
template_kwargs = {}
|
||||
if is_gpt_oss:
|
||||
# harmony format: the system slot always carries an identity —
|
||||
# "You are ChatGPT, a large language model trained by OpenAI."
|
||||
# unless model_identity overrides it — while a user "system"
|
||||
# message is APPENDED as a developer message. We make the
|
||||
# user's system prompt BE the identity (no OpenAI default, no
|
||||
# duplicated developer copy); with no system prompt, a neutral
|
||||
# identity replaces the default.
|
||||
sys_prompts = [m["content"] for m in messages if m["role"] == "system"]
|
||||
identity = (sys_prompts[0] or "").strip() if sys_prompts else ""
|
||||
template_kwargs["model_identity"] = identity or "You are a helpful assistant."
|
||||
if sys_prompts:
|
||||
messages = [m for m in messages if m["role"] != "system"]
|
||||
encoded = tokenizer.apply_chat_template(
|
||||
messages,
|
||||
add_generation_prompt=not continue_final,
|
||||
continue_final_message=continue_final,
|
||||
return_tensors="pt",
|
||||
enable_thinking=False,
|
||||
**template_kwargs,
|
||||
)
|
||||
input_ids = encoded if isinstance(encoded, torch.Tensor) else encoded["input_ids"]
|
||||
input_ids = input_ids.to(_input_device(hf_model))
|
||||
|
||||
# gpt-oss: enable_thinking does not apply to the harmony template.
|
||||
# We prime the "final" channel directly to skip the CoT ("analysis"
|
||||
# channel) → direct answer, no chain of thought.
|
||||
# (Not when continuing: the final message is already mid-channel.)
|
||||
if is_gpt_oss and not continue_final:
|
||||
final_prefix = torch.tensor(
|
||||
[tokenizer.encode("<|channel|>final<|message|>", add_special_tokens=False)],
|
||||
device=input_ids.device, dtype=input_ids.dtype,
|
||||
)
|
||||
input_ids = torch.cat([input_ids, final_prefix], dim=1)
|
||||
|
||||
read_from = 0
|
||||
gen_id = None
|
||||
if lens is not None and lens.lens is not None:
|
||||
reader = ActivationCatcher(self.jl.layers, lens.layers)
|
||||
gen_id = lens.start_gen()
|
||||
if len(messages) > 1 and any(m["role"] != "system" for m in messages[:-1]):
|
||||
prev = tokenizer.apply_chat_template(
|
||||
messages[:-1],
|
||||
add_generation_prompt=False,
|
||||
return_tensors="pt",
|
||||
enable_thinking=False,
|
||||
**template_kwargs,
|
||||
)
|
||||
prev_ids = prev if isinstance(prev, torch.Tensor) else prev["input_ids"]
|
||||
read_from = min(prev_ids.shape[1], input_ids.shape[1] - 1)
|
||||
|
||||
temperature = float(sampling.get("temperature", config.DEFAULT_SAMPLING["temperature"]))
|
||||
top_p = float(sampling.get("top_p", config.DEFAULT_SAMPLING["top_p"]))
|
||||
top_k = int(sampling.get("top_k", config.DEFAULT_SAMPLING["top_k"]))
|
||||
max_tokens = int(sampling.get("max_tokens", config.DEFAULT_SAMPLING["max_tokens"]))
|
||||
seed = int(sampling.get("seed", config.DEFAULT_SAMPLING["seed"]))
|
||||
# base model with a generic template: the model has no notion of dialogue
|
||||
# turns, so we cut as soon as it reopens one (User: or a new Assistant:)
|
||||
stop_seqs = (
|
||||
["\nUser:", "\nAssistant:"]
|
||||
if (self.meta or {}).get("chat_template_fallback")
|
||||
else []
|
||||
)
|
||||
|
||||
out = hf_model(input_ids=input_ids, use_cache=True)
|
||||
cache = out.past_key_values
|
||||
logits = out.logits[:, -1]
|
||||
eos = hf_model.generation_config.eos_token_id
|
||||
eos_ids = set(eos) if isinstance(eos, list) else {eos}
|
||||
# if the template applies end-of-turn markers (e.g. an instruct template
|
||||
# placed on a base model), add them to the stop tokens.
|
||||
applied_template = getattr(tokenizer, "chat_template", "") or ""
|
||||
if isinstance(applied_template, str) and not is_gpt_oss:
|
||||
unk = tokenizer.unk_token_id
|
||||
for marker in TURN_END_MARKERS:
|
||||
if marker in applied_template:
|
||||
tid = tokenizer.convert_tokens_to_ids(marker)
|
||||
if isinstance(tid, int) and tid >= 0 and tid != unk:
|
||||
eos_ids.add(tid)
|
||||
if is_gpt_oss:
|
||||
# in harmony <|end|> separates MESSAGES (analysis → final), but our
|
||||
# prompt primes the final channel directly (and "continue" resumes
|
||||
# mid-final), so there is never a transition to protect: the first
|
||||
# <|end|> IS the end of the turn. The model often emits it instead
|
||||
# of <|return|>; without this stop it then replays a whole
|
||||
# "assistant analysis ..." turn in plain text up to max_tokens.
|
||||
tid = tokenizer.convert_tokens_to_ids("<|end|>")
|
||||
if isinstance(tid, int) and tid >= 0:
|
||||
eos_ids.add(tid)
|
||||
|
||||
# seed >= 0: reproducible sampling; -1 = random
|
||||
generator = None
|
||||
if seed >= 0:
|
||||
generator = torch.Generator(device=logits.device).manual_seed(seed)
|
||||
|
||||
if reader is not None:
|
||||
positions = list(range(read_from, input_ids.shape[1]))
|
||||
reading_frames = lens.compute_frames(
|
||||
reader.acts,
|
||||
positions,
|
||||
"reading",
|
||||
self.jl,
|
||||
input_ids[0, read_from:].tolist(),
|
||||
gen_id=gen_id,
|
||||
)
|
||||
for frame in reading_frames:
|
||||
emit(frame)
|
||||
|
||||
reply_ids = []
|
||||
emitted = ""
|
||||
started = time.perf_counter()
|
||||
for _ in range(max_tokens):
|
||||
if stop_event.is_set():
|
||||
break
|
||||
next_id = _sample(logits[0].float(), temperature, top_p, top_k, generator)
|
||||
if next_id in eos_ids:
|
||||
break
|
||||
reply_ids.append(next_id)
|
||||
text = tokenizer.decode(reply_ids, skip_special_tokens=True)
|
||||
stop_hit = next((s for s in stop_seqs if s in text), None)
|
||||
if stop_hit:
|
||||
text = text[: text.index(stop_hit)]
|
||||
if not text.endswith("�") and len(text) > len(emitted):
|
||||
emit({"type": "token", "text": text[len(emitted):]})
|
||||
emitted = text
|
||||
if stop_hit:
|
||||
break
|
||||
out = hf_model(
|
||||
input_ids=torch.tensor([[next_id]], device=_input_device(hf_model)),
|
||||
past_key_values=cache,
|
||||
use_cache=True,
|
||||
)
|
||||
cache = out.past_key_values
|
||||
logits = out.logits[:, -1]
|
||||
if reader is not None:
|
||||
frame = lens.compute_frames(
|
||||
reader.acts,
|
||||
[-1],
|
||||
"thinking",
|
||||
self.jl,
|
||||
[next_id],
|
||||
gen_id=gen_id,
|
||||
abs_positions=[input_ids.shape[1] + len(reply_ids) - 1],
|
||||
)[0]
|
||||
emit(frame)
|
||||
|
||||
elapsed = time.perf_counter() - started
|
||||
text = tokenizer.decode(reply_ids, skip_special_tokens=True)
|
||||
for s in stop_seqs:
|
||||
if s in text:
|
||||
text = text[: text.index(s)]
|
||||
break
|
||||
emit(
|
||||
{
|
||||
"type": "done",
|
||||
"text": text,
|
||||
"gen_id": gen_id,
|
||||
"stopped": stop_event.is_set(),
|
||||
"stats": {
|
||||
"tokens": len(reply_ids),
|
||||
"seconds": round(elapsed, 2),
|
||||
"tok_per_s": round(len(reply_ids) / elapsed, 2) if reply_ids and elapsed > 0 else 0.0,
|
||||
},
|
||||
"meta": dict(
|
||||
self.meta or {},
|
||||
sampling=sampling,
|
||||
lens=dict(lens.meta) if lens is not None and lens.meta else None,
|
||||
interventions=ablator.summary() if ablator is not None else None,
|
||||
interventions_scale=ablator.global_scale if ablator is not None else None,
|
||||
),
|
||||
}
|
||||
)
|
||||
ok = True
|
||||
finally:
|
||||
if ablator is not None:
|
||||
ablator.detach()
|
||||
if reader is not None:
|
||||
reader.close()
|
||||
self.busy = None
|
||||
# aborted generation (OOM/error/hard stop): the KV cache and captured
|
||||
# activations are now dereferenced — return the blocks
|
||||
if not ok:
|
||||
_free_cuda()
|
||||
elif torch.cuda.is_available():
|
||||
# success path: when the device is nearly full (big model + long
|
||||
# KV cache), the freed cache fragments the reserve and the next
|
||||
# prefill hits costly allocator retries — generation gets slower
|
||||
# with every message. Hand segments back once the reserve crosses
|
||||
# 92 % of the device; a no-op (no sync, no gc) below that.
|
||||
for i in range(torch.cuda.device_count()):
|
||||
total = torch.cuda.get_device_properties(i).total_memory
|
||||
if torch.cuda.memory_reserved(i) > 0.92 * total:
|
||||
with torch.cuda.device(i):
|
||||
torch.cuda.empty_cache()
|
||||
@@ -0,0 +1,115 @@
|
||||
import re
|
||||
import threading
|
||||
|
||||
import torch
|
||||
|
||||
from core.lens_manager import MASKS_DIR, _vocab_fingerprint
|
||||
|
||||
# ── EVALUATION TOGGLE ────────────────────────────────────────────────────────
|
||||
# Nearest tokens (the "translation" of a non-latin token to readable neighbors):
|
||||
# True = keep only ENGLISH words (pure ASCII, no accents) as targets
|
||||
# False = any readable latin script (accents included: fr/de/es…)
|
||||
# Set to True by default; flip it to compare.
|
||||
ENGLISH_ONLY = True
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
# "translation" targets: readable tokens (2+ letter word, apostrophe/hyphen
|
||||
# allowed) so the neighbors are interpretable
|
||||
_LATIN_RE = re.compile(r"^[ A-Za-zÀ-ɏ'\-]+$")
|
||||
_LATIN_LETTERS_RE = re.compile(r"[A-Za-zÀ-ɏ]{2}")
|
||||
# english variant: pure ASCII (excludes café, über, naïve… → filters out the
|
||||
# other latin-script languages)
|
||||
_ENGLISH_RE = re.compile(r"^[ A-Za-z'\-]+$")
|
||||
_ENGLISH_LETTERS_RE = re.compile(r"[A-Za-z]{2}")
|
||||
|
||||
|
||||
def _latin_target_mask(tokenizer, vocab_size):
|
||||
MASKS_DIR.mkdir(parents=True, exist_ok=True)
|
||||
# distinct cache per mode (otherwise a "latin" mask would serve in english mode)
|
||||
tag = "english" if ENGLISH_ONLY else "latin"
|
||||
word_re = _ENGLISH_RE if ENGLISH_ONLY else _LATIN_RE
|
||||
letters_re = _ENGLISH_LETTERS_RE if ENGLISH_ONLY else _LATIN_LETTERS_RE
|
||||
path = MASKS_DIR / f"{_vocab_fingerprint(tokenizer)}_{vocab_size}_{tag}.pt"
|
||||
if path.exists():
|
||||
return torch.load(path, weights_only=True)
|
||||
mask = torch.zeros(vocab_size, dtype=torch.bool)
|
||||
n_decodable = min(vocab_size, len(tokenizer))
|
||||
decoded = tokenizer.batch_decode(
|
||||
[[tid] for tid in range(n_decodable)], clean_up_tokenization_spaces=False
|
||||
)
|
||||
for tid, raw in enumerate(decoded):
|
||||
s = raw.strip()
|
||||
mask[tid] = bool(
|
||||
len(s) >= 2 and word_re.match(s) and letters_re.search(s)
|
||||
)
|
||||
torch.save(mask, path)
|
||||
return mask
|
||||
|
||||
|
||||
class TokenNeighbors:
|
||||
"""Approximate local translation: latin tokens whose output direction (row of
|
||||
W_U) is closest in cosine to a non-latin token most often carry the same
|
||||
meaning (答案 → ' answer')."""
|
||||
|
||||
def __init__(self):
|
||||
self._lock = threading.Lock()
|
||||
self._key = None
|
||||
self._mask = None
|
||||
self._norms = None
|
||||
self._cache = {}
|
||||
|
||||
def _prepare(self, jl, tokenizer, model_key):
|
||||
if self._key == model_key and self._norms is not None:
|
||||
return
|
||||
weight = jl._lm_head.weight
|
||||
if weight.dtype not in (torch.float16, torch.bfloat16, torch.float32):
|
||||
raise ValueError("neighbors unavailable on a quantized model")
|
||||
vocab_size = weight.shape[0]
|
||||
self._mask = _latin_target_mask(tokenizer, vocab_size).to(weight.device)
|
||||
norms = torch.empty(vocab_size, dtype=torch.float32, device=weight.device)
|
||||
with torch.no_grad():
|
||||
for start in range(0, vocab_size, 8192):
|
||||
chunk = weight[start:start + 8192].float()
|
||||
norms[start:start + 8192] = chunk.norm(dim=1)
|
||||
self._norms = norms.clamp_min(1e-8)
|
||||
self._cache = {}
|
||||
self._key = model_key
|
||||
|
||||
def lookup(self, jl, tokenizer, model_key, token_ids, k=3):
|
||||
with self._lock:
|
||||
self._prepare(jl, tokenizer, model_key)
|
||||
weight = jl._lm_head.weight
|
||||
out = {}
|
||||
for tid in token_ids:
|
||||
tid = int(tid)
|
||||
if tid < 0 or tid >= weight.shape[0]:
|
||||
out[tid] = []
|
||||
continue
|
||||
if tid in self._cache:
|
||||
out[tid] = self._cache[tid]
|
||||
continue
|
||||
with torch.no_grad():
|
||||
v = weight[tid]
|
||||
sims = (weight @ v).float() / (self._norms * self._norms[tid])
|
||||
sims[~self._mask] = float("-inf")
|
||||
sims[tid] = float("-inf")
|
||||
top = torch.topk(sims, min(k, int(self._mask.sum())))
|
||||
entries = [
|
||||
{
|
||||
"id": int(i),
|
||||
"str": tokenizer.decode([int(i)]),
|
||||
"sim": round(float(s), 3),
|
||||
}
|
||||
for s, i in zip(top.values.tolist(), top.indices.tolist())
|
||||
if s != float("-inf")
|
||||
]
|
||||
self._cache[tid] = entries
|
||||
out[tid] = entries
|
||||
return out
|
||||
|
||||
def reset(self):
|
||||
with self._lock:
|
||||
self._key = None
|
||||
self._mask = None
|
||||
self._norms = None
|
||||
self._cache = {}
|
||||
+300
@@ -0,0 +1,300 @@
|
||||
"""Change of basis of the residual: faithful pure-weight bake of the steering.
|
||||
|
||||
The standard hook applies ``h ← M_l·h`` at the output of each hooked layer, with
|
||||
``M_l = Π_rules (I + w·v̂ᵀ)`` (rank-1 per rule, the layer's J-space directions).
|
||||
This transformed residual is then READ by everything downstream through matrices:
|
||||
each sub-block reads ``W·(γ ⊙ h/rms(h))`` via its RMSNorm, and lm_head reads via
|
||||
the final norm. So we realize the transform in the downstream READS instead of the
|
||||
writes (the "skip" escapes no one in reading, whereas no matrix carries it in
|
||||
writing — the cause of the ~1.5 % of the per-layer bake):
|
||||
|
||||
read of layer m: W ← W·Γ·C_m·Γ⁻¹ (Γ = diag(γ) of the read RMSNorm)
|
||||
lm_head: W ← W·Γ_f·C_fin·Γ_f⁻¹
|
||||
write of layer m: W ← C_m⁻¹·W ("exact" mode only)
|
||||
|
||||
where ``C_m = M_{m-1}···M_{l0}`` composes the hooks strictly upstream of m.
|
||||
``C = I + U·Vᵀ`` stays low-rank end to end (one column per rule and per hooked
|
||||
layer), so each matrix receives a rank-r update.
|
||||
|
||||
Two variants:
|
||||
- "readthrough": reads only. For saturated zaps/replaces (M idempotent), this
|
||||
equals the hook applied over a range extended to the last layer, with a slight
|
||||
bias toward MORE effect (the range's intermediate writes are projected too).
|
||||
No inversion: robust in bf16 and to GGUF quantization.
|
||||
- "exact": adds the counter-transform of the writes to reproduce a hook applied
|
||||
ONCE at the chosen point. C⁻¹ blows up near a full zap (1 + v̂ᵀw → 0):
|
||||
reserved for soft factors, regularized inverse.
|
||||
|
||||
Assumed approximation (the only one): the rms in the RMSNorm denominator stays
|
||||
that of the untransformed residual — a per-position scalar error, second-order
|
||||
when the modified component is small compared to ‖h‖. Same assumption as all of
|
||||
the weight-orthogonalization literature.
|
||||
|
||||
The live preview mode (core/ablation) applies the SAME transform via hooks on the
|
||||
RMSNorm output: the preview and the exported checkpoint differ only by rounding.
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
from core.ablation import effective_coeffs
|
||||
|
||||
# division by γ: channels with γ=0 are dead (never read via this norm), the clamp
|
||||
# is exact there; between 0 and EPS the error is bounded and negligible
|
||||
GAMMA_EPS = 1e-6
|
||||
|
||||
# regularization threshold of the inverse (exact mode): below it, the
|
||||
# counter-transform amplifies the downstream writes (×1/σ), which makes the RMS
|
||||
# error first-order and destroys bf16 precision then GGUF quantization. 0.2 bounds
|
||||
# the amplification to ×5; a saturated replace (α = −1 as soon as scale ≥ 1) is
|
||||
# ALWAYS in this regime → prefer readthrough.
|
||||
INV_COND_EPS = 0.2
|
||||
|
||||
# Residual reads per sub-block: {module suffix: suffix of the read RMSNorm}.
|
||||
# Covers Llama/Qwen/Mistral (self_attn+mlp) and Qwen3.5/Qwen3-Next
|
||||
# (linear_attn GatedDeltaNet). conv1d/q_norm/k_norm operate AFTER these
|
||||
# projections: they see the transformed residual without us touching them.
|
||||
READS = {
|
||||
"self_attn.q_proj": "input_layernorm",
|
||||
"self_attn.k_proj": "input_layernorm",
|
||||
"self_attn.v_proj": "input_layernorm",
|
||||
"linear_attn.in_proj_qkv": "input_layernorm",
|
||||
"linear_attn.in_proj_z": "input_layernorm",
|
||||
"linear_attn.in_proj_b": "input_layernorm",
|
||||
"linear_attn.in_proj_a": "input_layernorm",
|
||||
"mlp.gate_proj": "input_layernorm", # replaced if post_attention is present
|
||||
"mlp.up_proj": "input_layernorm",
|
||||
}
|
||||
# most archs read the MLP via post_attention_layernorm
|
||||
READS_POST = {"mlp.gate_proj", "mlp.up_proj"}
|
||||
|
||||
# Writes into the residual (exact mode only)
|
||||
WRITES = ("self_attn.o_proj", "linear_attn.out_proj", "mlp.down_proj")
|
||||
|
||||
# archs where post_attention_layernorm normalizes the attention WRITE (not the
|
||||
# MLP read): the read transform would be wrong there
|
||||
_UNSUPPORTED_MARKERS = ("pre_feedforward_layernorm", "post_feedforward_layernorm")
|
||||
|
||||
|
||||
def _submodule(block, dotted):
|
||||
module = block
|
||||
for part in dotted.split("."):
|
||||
module = getattr(module, part, None)
|
||||
if module is None:
|
||||
return None
|
||||
return module
|
||||
|
||||
|
||||
def check_block_supported(block):
|
||||
for marker in _UNSUPPORTED_MARKERS:
|
||||
if getattr(block, marker, None) is not None:
|
||||
raise ValueError(
|
||||
"architecture not supported by the readthrough/exact modes: "
|
||||
f"the layer has {marker} (write norm, Gemma style) — "
|
||||
"the read transform would be incorrect there"
|
||||
)
|
||||
|
||||
|
||||
def iter_reads(block):
|
||||
"""Yields ``(suffix, module, norm)`` for each residual read."""
|
||||
check_block_supported(block)
|
||||
for suffix, norm_name in READS.items():
|
||||
module = _submodule(block, suffix)
|
||||
if module is None:
|
||||
continue
|
||||
if suffix in READS_POST and getattr(block, "post_attention_layernorm", None) is not None:
|
||||
norm_name = "post_attention_layernorm"
|
||||
norm = getattr(block, norm_name, None)
|
||||
if norm is None or not hasattr(norm, "weight"):
|
||||
raise ValueError(f"RMSNorm {norm_name} not found for {suffix}")
|
||||
yield suffix, module, norm
|
||||
|
||||
|
||||
def iter_writes(block):
|
||||
for suffix in WRITES:
|
||||
module = _submodule(block, suffix)
|
||||
if module is not None:
|
||||
yield suffix, module
|
||||
|
||||
|
||||
def rule_factors(rules, scale):
|
||||
"""Rank-1 factors ``{layer: [(w, v̂), ...]}`` float32 CPU, in the standard
|
||||
hook's application order (increasing layers, rules in order).
|
||||
``w = α·v̂_A + β·v̂_B`` with the effective coefficients (saturation included).
|
||||
Returns an empty dict if all coefficients are neutral."""
|
||||
by_layer = {}
|
||||
for rule in rules:
|
||||
alpha, beta = effective_coeffs(rule["mode"], rule["factor"], scale)
|
||||
if alpha == 0.0 and not beta:
|
||||
continue
|
||||
for layer in rule["layers"]:
|
||||
v_a = rule["dirs_a"][layer].detach().float().cpu()
|
||||
w = alpha * v_a
|
||||
if beta:
|
||||
w = w + beta * rule["dirs_b"][layer].detach().float().cpu()
|
||||
if w.norm() < 1e-8:
|
||||
continue # null W_U row → empty direction, nothing to apply
|
||||
by_layer.setdefault(int(layer), []).append((w, v_a))
|
||||
return by_layer
|
||||
|
||||
|
||||
def _compose_left(U, V, w, v):
|
||||
"""``(I + w·vᵀ)·(I + U·Vᵀ)`` → new ``(U, V)`` (one more column)."""
|
||||
if U is None:
|
||||
return w.unsqueeze(1), v.unsqueeze(1)
|
||||
v_new = v + V @ (U.T @ v)
|
||||
return torch.cat([U, w.unsqueeze(1)], dim=1), torch.cat([V, v_new.unsqueeze(1)], dim=1)
|
||||
|
||||
|
||||
def compress_uv(U, V, tol=1e-5):
|
||||
"""Recompacts ``C − I = U·Vᵀ`` via QR + truncated SVD.
|
||||
|
||||
Essential, not cosmetic: a token's directions across layers are nearly
|
||||
collinear, so naive composition inflates the columns (multiplicative cross
|
||||
terms) and the result only holds through cancellation between large numbers —
|
||||
invisible in float32, destructive in bf16 (live preview → random tokens,
|
||||
measured). After compression V is orthonormal and U carries the true singular
|
||||
values (~O(1)): stable in bf16 and rank reduced to the effective rank."""
|
||||
Qu, Ru = torch.linalg.qr(U)
|
||||
Qv, Rv = torch.linalg.qr(V)
|
||||
Us, S, Vh = torch.linalg.svd(Ru @ Rv.T)
|
||||
keep = S > tol * S.max().clamp_min(1e-12)
|
||||
return Qu @ (Us[:, keep] * S[keep]), Qv @ Vh.T[:, keep]
|
||||
|
||||
|
||||
def cumulative(rules, scale, n_layers):
|
||||
"""Cumulative transforms ``{m: (U, V)}`` for each read point:
|
||||
m = layer (its reads see ``C_m`` = hooks of layers < m);
|
||||
the ``n_layers`` key is the final norm / lm_head point.
|
||||
Returns ``{}`` if no factor is active. The (U, V) of consecutive layers with
|
||||
no intermediate hook share their tensors (never mutated)."""
|
||||
factors = rule_factors(rules, scale)
|
||||
if not factors:
|
||||
return {}
|
||||
l_min = min(factors)
|
||||
out = {}
|
||||
U = V = None
|
||||
for layer in range(l_min, n_layers):
|
||||
if factors.get(layer):
|
||||
for w, v in factors[layer]:
|
||||
U, V = _compose_left(U, V, w, v)
|
||||
U, V = compress_uv(U, V)
|
||||
if U is not None:
|
||||
out[layer + 1] = (U, V)
|
||||
return out
|
||||
|
||||
|
||||
def effective_gamma(norm):
|
||||
"""MEASURED effective γ: ``norm(1⃗) = γ_eff`` since rms(1⃗) = 1.
|
||||
|
||||
Do NOT read ``norm.weight`` directly: Qwen3.5 (like Gemma) uses a
|
||||
zero-centered RMSNorm where γ = 1 + weight — dividing by ``weight`` (~0, of
|
||||
arbitrary sign) made the transform chaotic (live preview → random tokens,
|
||||
measured). The functional measurement covers both styles."""
|
||||
weight = norm.weight
|
||||
with torch.no_grad():
|
||||
ones = torch.ones(1, weight.shape[-1], device=weight.device, dtype=torch.float32)
|
||||
return norm(ones).detach().flatten().float().cpu()
|
||||
|
||||
|
||||
def gamma_pair(norm, U, V):
|
||||
"""``(γ⊙U, V/γ)`` float32 CPU for the read via this RMSNorm:
|
||||
``W·Γ·C·Γ⁻¹ = W + (W·(γ⊙U))·(V/γ)ᵀ``."""
|
||||
gamma = effective_gamma(norm)
|
||||
safe = torch.where(gamma.abs() < GAMMA_EPS, torch.full_like(gamma, GAMMA_EPS), gamma)
|
||||
return gamma.unsqueeze(1) * U, V / safe.unsqueeze(1)
|
||||
|
||||
|
||||
def apply_read(W, Ug, Vg):
|
||||
"""``W ← W·(I + Ug·Vgᵀ)``; returns ``(W_new, B, A)`` with delta = B·A."""
|
||||
B = W @ Ug # [out, r]
|
||||
return W + B @ Vg.T, B, Vg.T.contiguous()
|
||||
|
||||
|
||||
def inverse_uv(U, V):
|
||||
"""``C⁻¹ = I − U_inv·Vᵀ`` (Woodbury: ``U_inv = U·(I_r + VᵀU)⁻¹``).
|
||||
Returns ``(U_inv, V, regularized)``; near a full zap the small matrix is
|
||||
singular → thresholded pseudo-inverse (the local effect ≈ readthrough)."""
|
||||
r = U.shape[1]
|
||||
small = torch.eye(r) + V.T @ U
|
||||
svals = torch.linalg.svdvals(small)
|
||||
regularized = bool(svals.min() < INV_COND_EPS * max(1.0, float(svals.max())))
|
||||
if regularized:
|
||||
inv = torch.linalg.pinv(small, rtol=INV_COND_EPS)
|
||||
else:
|
||||
inv = torch.linalg.inv(small)
|
||||
return U @ inv, V, regularized
|
||||
|
||||
|
||||
def apply_write(W, U_inv, V):
|
||||
"""``W ← (I − U_inv·Vᵀ)·W``; returns ``(W_new, B, A)`` with delta = B·A."""
|
||||
A = V.T @ W # [r, in]
|
||||
return W - U_inv @ A, (-U_inv).contiguous(), A
|
||||
|
||||
|
||||
def apply_transform(entry, W):
|
||||
"""Applies a plan entry to a float32 weight.
|
||||
|
||||
Returns ``(W_new, B, A)`` where ``B·A`` is the EXACT delta ``W_new − W``:
|
||||
the rebase update is low-rank by construction, which is what makes the LoRA
|
||||
export exact rather than an approximation."""
|
||||
kind, X, Y = entry
|
||||
if kind == "read":
|
||||
return apply_read(W, X, Y)
|
||||
return apply_write(W, X, Y)
|
||||
|
||||
|
||||
def build_plan(rules, jl, scale, exact=False):
|
||||
"""Bake plan: ``{param_name: entry}`` with ``entry = ("read", Ug, Vg)`` or
|
||||
``("write", U_inv, V)`` — apply with :func:`apply_transform` — plus the
|
||||
diagnostic metadata.
|
||||
|
||||
The names follow the model's layout (``{path}.layers.{m}.{suffix}.weight``,
|
||||
``{lm_head}.weight``); the guard matching them against the checkpoint keys is
|
||||
done by the export."""
|
||||
active = [r for r in rules if r["layers"]]
|
||||
if not active:
|
||||
raise ValueError("no active rule (all have 0 layers): nothing to export")
|
||||
n_layers = len(jl.layers)
|
||||
cums = cumulative(active, scale, n_layers)
|
||||
if not cums:
|
||||
raise ValueError(
|
||||
"all coefficients neutral (factors at 1 and/or scale=0): "
|
||||
"the bake would change no weight"
|
||||
)
|
||||
path = jl.layout.path
|
||||
transforms = {}
|
||||
regularized_layers = []
|
||||
min_gamma = None
|
||||
|
||||
for m in sorted(k for k in cums if k < n_layers):
|
||||
U, V = cums[m]
|
||||
block = jl.layers[m]
|
||||
for suffix, _module, norm in iter_reads(block):
|
||||
Ug, Vg = gamma_pair(norm, U, V)
|
||||
g_min = effective_gamma(norm).abs().min().item()
|
||||
min_gamma = g_min if min_gamma is None else min(min_gamma, g_min)
|
||||
transforms[f"{path}.layers.{m}.{suffix}.weight"] = ("read", Ug, Vg)
|
||||
if exact:
|
||||
U_inv, Vw, regularized = inverse_uv(U, V)
|
||||
if regularized:
|
||||
regularized_layers.append(m)
|
||||
for suffix, _module in iter_writes(block):
|
||||
transforms[f"{path}.layers.{m}.{suffix}.weight"] = ("write", U_inv, Vw)
|
||||
|
||||
U, V = cums[n_layers]
|
||||
Ug, Vg = gamma_pair(jl._final_norm, U, V)
|
||||
lm_head_key = f"{jl.layout.lm_head}.weight"
|
||||
transforms[lm_head_key] = ("read", Ug, Vg)
|
||||
|
||||
tied = jl._lm_head.weight.data_ptr() == jl._embed_tokens.weight.data_ptr()
|
||||
info = {
|
||||
"tied": tied,
|
||||
"lm_head_key": lm_head_key,
|
||||
"embed_key": f"{path}.{jl.layout.embed}.weight",
|
||||
"path": path,
|
||||
"rank_final": cums[n_layers][0].shape[1],
|
||||
"layers_span": [min(cums), n_layers - 1],
|
||||
"regularized_layers": regularized_layers,
|
||||
"min_gamma": min_gamma,
|
||||
}
|
||||
return transforms, info
|
||||
@@ -0,0 +1,253 @@
|
||||
import json
|
||||
import re
|
||||
import time
|
||||
|
||||
from huggingface_hub import HfApi
|
||||
|
||||
import config
|
||||
|
||||
HUB_SEED_REPO = "neuronpedia/jacobian-lens"
|
||||
HUB_CACHE_TTL = 600
|
||||
|
||||
_hub_cache = {"at": 0.0, "entries": None, "error": None}
|
||||
|
||||
|
||||
def local_lenses():
|
||||
out = []
|
||||
if not config.LENSES_DIR.exists():
|
||||
return out
|
||||
for entry in sorted(config.LENSES_DIR.iterdir()):
|
||||
lens_file = entry / "lens.pt"
|
||||
if not entry.is_dir() or not lens_file.exists():
|
||||
continue
|
||||
meta = {}
|
||||
meta_file = entry / "meta.json"
|
||||
if meta_file.exists():
|
||||
meta = json.loads(meta_file.read_text(encoding="utf-8"))
|
||||
out.append({"name": entry.name, "path": str(lens_file), "meta": meta})
|
||||
return out
|
||||
|
||||
|
||||
def _base_model_from(filename):
|
||||
stem = filename.rsplit("/", 1)[-1].removesuffix(".pt")
|
||||
match = re.match(r"(.+?)_jacobian_lens(?:_n\d+)?$", stem)
|
||||
return match.group(1) if match else None
|
||||
|
||||
|
||||
def _derived_model_id(base):
|
||||
if base is None:
|
||||
return None
|
||||
lowered = base.lower()
|
||||
if lowered.startswith("qwen"):
|
||||
return f"Qwen/{base}"
|
||||
if lowered.startswith("gemma"):
|
||||
return f"google/{base}"
|
||||
if lowered.startswith("llama"):
|
||||
return f"meta-llama/{base}"
|
||||
if lowered.startswith("gpt-oss"):
|
||||
return f"openai/{base}"
|
||||
if lowered == "gpt2":
|
||||
return "openai-community/gpt2"
|
||||
if lowered.startswith("pythia"):
|
||||
return f"EleutherAI/{base}"
|
||||
if lowered.startswith("olmo"):
|
||||
return f"allenai/{base}"
|
||||
return base
|
||||
|
||||
|
||||
def hub_lenses(force=False):
|
||||
now = time.time()
|
||||
if not force and _hub_cache["entries"] is not None and now - _hub_cache["at"] < HUB_CACHE_TTL:
|
||||
return _hub_cache["entries"]
|
||||
api = HfApi()
|
||||
repos = {HUB_SEED_REPO}
|
||||
try:
|
||||
for model in api.list_models(search="jacobian-lens", limit=50):
|
||||
repos.add(model.id)
|
||||
for model in api.list_models(filter="jacobian_lens", limit=50):
|
||||
repos.add(model.id)
|
||||
except Exception as exc:
|
||||
_hub_cache.update(error=f"Hub search unavailable: {exc}")
|
||||
entries = []
|
||||
for repo_id in sorted(repos):
|
||||
try:
|
||||
refs = api.list_repo_refs(repo_id)
|
||||
branches = [b.name for b in refs.branches] or ["main"]
|
||||
except Exception:
|
||||
continue
|
||||
for branch in branches:
|
||||
try:
|
||||
files = api.list_repo_files(repo_id, revision=branch)
|
||||
except Exception:
|
||||
continue
|
||||
for filename in files:
|
||||
if not filename.endswith(".pt"):
|
||||
continue
|
||||
base = _base_model_from(filename)
|
||||
entries.append(
|
||||
{
|
||||
"repo_id": repo_id,
|
||||
"revision": branch,
|
||||
"filename": filename,
|
||||
"base_model": base,
|
||||
"derived_model_id": _derived_model_id(base),
|
||||
"model_revision_verified": False,
|
||||
}
|
||||
)
|
||||
_hub_cache.update(at=now, entries=entries)
|
||||
return entries
|
||||
|
||||
|
||||
_base_cache = {}
|
||||
BASE_CACHE_TTL = 600
|
||||
|
||||
|
||||
def hub_base_model(model_id):
|
||||
"""Base model declared by the repo's model card (``base_model`` tags) —
|
||||
e.g. a finetune pointing at the checkpoint it was trained from. ``None``
|
||||
if unknown, offline, or not a Hub repo."""
|
||||
now = time.time()
|
||||
hit = _base_cache.get(model_id)
|
||||
if hit and now - hit["at"] < BASE_CACHE_TTL:
|
||||
return hit["base"]
|
||||
found = None
|
||||
if "/" in model_id and not model_id.startswith("local/"):
|
||||
try:
|
||||
info = HfApi().model_info(model_id)
|
||||
for tag in info.tags or []:
|
||||
if not tag.startswith("base_model:"):
|
||||
continue
|
||||
rest = tag[len("base_model:"):]
|
||||
if ":" in rest: # qualified form: finetune:X, adapter:X, quantized:X
|
||||
rest = rest.split(":", 1)[1]
|
||||
if rest and rest.lower() != model_id.lower():
|
||||
found = rest
|
||||
break
|
||||
except Exception:
|
||||
found = None
|
||||
_base_cache[model_id] = {"at": now, "base": found}
|
||||
return found
|
||||
|
||||
|
||||
def lenses_for_model(model_id, revision=None):
|
||||
matches_local = []
|
||||
for lens in local_lenses():
|
||||
meta = lens["meta"]
|
||||
if meta.get("model_id") != model_id:
|
||||
continue
|
||||
lens_rev = meta.get("model_revision")
|
||||
compatible = True
|
||||
reason = None
|
||||
if revision and lens_rev and lens_rev != revision:
|
||||
compatible = False
|
||||
reason = f"fit revision ({lens_rev[:12]}) != loaded model ({revision[:12]})"
|
||||
elif lens_rev is None and not model_id.startswith("local/"):
|
||||
reason = "fit revision unknown"
|
||||
matches_local.append(dict(lens, compatible=compatible, reason=reason))
|
||||
|
||||
base = model_id.split("/")[-1].lower()
|
||||
base_ref = hub_base_model(model_id) # e.g. "google/gemma-3-1b-it" for a finetune
|
||||
base_ref_name = base_ref.split("/")[-1].lower() if base_ref else None
|
||||
|
||||
def hub_entry(entry, via, reason=None, compatible=True):
|
||||
return dict(
|
||||
entry,
|
||||
via=via,
|
||||
compatible=compatible,
|
||||
reason=reason,
|
||||
cached=_lens_cached(entry["repo_id"], entry["filename"], entry["revision"]),
|
||||
)
|
||||
|
||||
# one entry per branch in hub_lenses → dedupe, main first
|
||||
entries = []
|
||||
seen = set()
|
||||
for entry in sorted(hub_lenses(), key=lambda e: e["revision"] != "main"):
|
||||
if entry["base_model"] is None or (entry["repo_id"], entry["filename"]) in seen:
|
||||
continue
|
||||
seen.add((entry["repo_id"], entry["filename"]))
|
||||
entries.append(entry)
|
||||
|
||||
matches_hub = []
|
||||
matched = set()
|
||||
prefix_hits = {}
|
||||
for entry in entries:
|
||||
key = (entry["repo_id"], entry["filename"])
|
||||
name = entry["base_model"].lower()
|
||||
derived = (entry["derived_model_id"] or "").lower()
|
||||
if name == base or derived == model_id.lower():
|
||||
# ⚠ only for a real problem (local merge); the fit revision not being
|
||||
# published is the normal state of Hub repos → discreet note
|
||||
reason = None
|
||||
if model_id.startswith("local/"):
|
||||
reason = "local model: a Hub lens fitted on the original checkpoint doesn't match a merge"
|
||||
matched.add(key)
|
||||
matches_hub.append(hub_entry(
|
||||
entry, "model", reason, compatible=not model_id.startswith("local/")))
|
||||
elif base_ref and (name == base_ref_name or derived == base_ref.lower()):
|
||||
matched.add(key)
|
||||
matches_hub.append(hub_entry(
|
||||
entry, "base-model",
|
||||
f"lens of the base model {base_ref} — fitted on the original "
|
||||
"weights, a finetune's readouts may drift slightly"))
|
||||
elif base_ref is None and name != base and len(name) >= 6 and base.startswith(name):
|
||||
prefix_hits.setdefault(len(name), []).append(entry)
|
||||
|
||||
# No card metadata: fall back to the longest name prefix (a finetune usually
|
||||
# keeps its base's name — "gemma-3-1b-it-toxicity" → "gemma-3-1b-it").
|
||||
if prefix_hits and not any(m["via"] == "model" for m in matches_hub):
|
||||
for entry in prefix_hits[max(prefix_hits)]:
|
||||
matched.add((entry["repo_id"], entry["filename"]))
|
||||
matches_hub.append(hub_entry(
|
||||
entry, "base-guess",
|
||||
f"the name suggests a finetune of {entry['base_model']} — fitted "
|
||||
"on the original weights, readouts may drift slightly"))
|
||||
|
||||
# Everything else stays reachable for cross-model loading (your own
|
||||
# architecture-compatible lens); d_model/layers are checked at load time.
|
||||
others = [
|
||||
hub_entry(entry, "other")
|
||||
for entry in entries
|
||||
if (entry["repo_id"], entry["filename"]) not in matched
|
||||
]
|
||||
return {
|
||||
"local": matches_local,
|
||||
"hub": matches_hub,
|
||||
"other": others,
|
||||
"base_model": base_ref,
|
||||
"hub_error": _hub_cache.get("error"),
|
||||
}
|
||||
|
||||
|
||||
def _lens_cached(repo_id, filename, revision=None):
|
||||
"""True if the lens file is already in the local HF cache (no download on load)."""
|
||||
from huggingface_hub import try_to_load_from_cache
|
||||
|
||||
try:
|
||||
result = try_to_load_from_cache(repo_id, filename, revision=revision)
|
||||
return isinstance(result, str)
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
||||
def resolve_lens(path=None, repo_id=None, filename=None):
|
||||
if path:
|
||||
for lens in local_lenses():
|
||||
if lens["path"] == path:
|
||||
meta = lens["meta"]
|
||||
return {
|
||||
"source": "local",
|
||||
"name": lens["name"],
|
||||
"required_model": meta.get("model_id"),
|
||||
"required_revision": meta.get("model_revision"),
|
||||
"meta": meta,
|
||||
}
|
||||
return {"source": "local", "required_model": None, "meta": {}, "warning": "meta.json missing: required model unknown"}
|
||||
base = _base_model_from(filename or "")
|
||||
return {
|
||||
"source": "hub",
|
||||
"repo_id": repo_id,
|
||||
"filename": filename,
|
||||
"required_model": _derived_model_id(base),
|
||||
"required_revision": None,
|
||||
"warning": "model derived from the filename; exact revision not published",
|
||||
}
|
||||
+378
@@ -0,0 +1,378 @@
|
||||
import json
|
||||
import sqlite3
|
||||
import threading
|
||||
from datetime import datetime, timezone
|
||||
|
||||
import msgpack
|
||||
import numpy as np
|
||||
|
||||
import config
|
||||
|
||||
DB_PATH = config.DATA_DIR / "jlens.db"
|
||||
FRAMES_DIR = config.DATA_DIR / "frames"
|
||||
|
||||
SCHEMA = """
|
||||
CREATE TABLE IF NOT EXISTS conversations (
|
||||
id INTEGER PRIMARY KEY,
|
||||
title TEXT NOT NULL DEFAULT '',
|
||||
tags TEXT NOT NULL DEFAULT '[]',
|
||||
created_at TEXT NOT NULL,
|
||||
updated_at TEXT NOT NULL
|
||||
);
|
||||
CREATE TABLE IF NOT EXISTS messages (
|
||||
id INTEGER PRIMARY KEY,
|
||||
conversation_id INTEGER NOT NULL REFERENCES conversations(id) ON DELETE CASCADE,
|
||||
parent_id INTEGER REFERENCES messages(id),
|
||||
role TEXT NOT NULL,
|
||||
content TEXT NOT NULL,
|
||||
meta TEXT,
|
||||
frames_file TEXT,
|
||||
created_at TEXT NOT NULL
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_messages_conv ON messages(conversation_id);
|
||||
CREATE VIRTUAL TABLE IF NOT EXISTS messages_fts USING fts5(
|
||||
content, content='messages', content_rowid='id'
|
||||
);
|
||||
CREATE TRIGGER IF NOT EXISTS messages_ai AFTER INSERT ON messages BEGIN
|
||||
INSERT INTO messages_fts(rowid, content) VALUES (new.id, new.content);
|
||||
END;
|
||||
CREATE TRIGGER IF NOT EXISTS messages_ad AFTER DELETE ON messages BEGIN
|
||||
INSERT INTO messages_fts(messages_fts, rowid, content)
|
||||
VALUES ('delete', old.id, old.content);
|
||||
END;
|
||||
CREATE TRIGGER IF NOT EXISTS messages_au AFTER UPDATE OF content ON messages BEGIN
|
||||
INSERT INTO messages_fts(messages_fts, rowid, content)
|
||||
VALUES ('delete', old.id, old.content);
|
||||
INSERT INTO messages_fts(rowid, content) VALUES (new.id, new.content);
|
||||
END;
|
||||
"""
|
||||
|
||||
|
||||
def _now():
|
||||
return datetime.now(timezone.utc).isoformat(timespec="seconds")
|
||||
|
||||
|
||||
class Store:
|
||||
def __init__(self):
|
||||
config.DATA_DIR.mkdir(parents=True, exist_ok=True)
|
||||
FRAMES_DIR.mkdir(parents=True, exist_ok=True)
|
||||
self._local = threading.local()
|
||||
conn = self._conn()
|
||||
conn.executescript(SCHEMA)
|
||||
conn.commit()
|
||||
|
||||
def _conn(self):
|
||||
conn = getattr(self._local, "conn", None)
|
||||
if conn is None:
|
||||
conn = sqlite3.connect(DB_PATH)
|
||||
conn.row_factory = sqlite3.Row
|
||||
conn.execute("PRAGMA journal_mode=WAL")
|
||||
conn.execute("PRAGMA foreign_keys=ON")
|
||||
self._local.conn = conn
|
||||
return conn
|
||||
|
||||
def create_conversation(self, title, tags=None):
|
||||
conn = self._conn()
|
||||
now = _now()
|
||||
cur = conn.execute(
|
||||
"INSERT INTO conversations (title, tags, created_at, updated_at) VALUES (?, ?, ?, ?)",
|
||||
(title, json.dumps(tags or []), now, now),
|
||||
)
|
||||
conn.commit()
|
||||
return cur.lastrowid
|
||||
|
||||
def update_conversation(self, conversation_id, title=None, tags=None):
|
||||
conn = self._conn()
|
||||
if title is not None:
|
||||
conn.execute(
|
||||
"UPDATE conversations SET title = ?, updated_at = ? WHERE id = ?",
|
||||
(title, _now(), conversation_id),
|
||||
)
|
||||
if tags is not None:
|
||||
conn.execute(
|
||||
"UPDATE conversations SET tags = ?, updated_at = ? WHERE id = ?",
|
||||
(json.dumps(tags), _now(), conversation_id),
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
def delete_conversation(self, conversation_id):
|
||||
conn = self._conn()
|
||||
rows = conn.execute(
|
||||
"SELECT frames_file FROM messages WHERE conversation_id = ? AND frames_file IS NOT NULL",
|
||||
(conversation_id,),
|
||||
).fetchall()
|
||||
conn.execute("DELETE FROM conversations WHERE id = ?", (conversation_id,))
|
||||
conn.commit()
|
||||
for row in rows:
|
||||
(FRAMES_DIR / row["frames_file"]).unlink(missing_ok=True)
|
||||
|
||||
def list_conversations(self, query=None, limit=200):
|
||||
conn = self._conn()
|
||||
if query:
|
||||
hits = conn.execute(
|
||||
"""
|
||||
SELECT messages_fts.rowid AS mid,
|
||||
snippet(messages_fts, 0, '[', ']', '…', 12) AS snip,
|
||||
rank
|
||||
FROM messages_fts
|
||||
WHERE messages_fts MATCH ?
|
||||
ORDER BY rank
|
||||
LIMIT 500
|
||||
""",
|
||||
(query,),
|
||||
).fetchall()
|
||||
best = {}
|
||||
for hit in hits:
|
||||
row = conn.execute(
|
||||
"SELECT conversation_id FROM messages WHERE id = ?", (hit["mid"],)
|
||||
).fetchone()
|
||||
if row and row["conversation_id"] not in best:
|
||||
best[row["conversation_id"]] = hit["snip"]
|
||||
rows = []
|
||||
for cid, snip in list(best.items())[:limit]:
|
||||
conv = conn.execute(
|
||||
"""
|
||||
SELECT c.id, c.title, c.tags, c.updated_at,
|
||||
(SELECT count(*) FROM messages m WHERE m.conversation_id = c.id) AS n_messages
|
||||
FROM conversations c WHERE c.id = ?
|
||||
""",
|
||||
(cid,),
|
||||
).fetchone()
|
||||
if conv:
|
||||
rows.append(dict(conv, snippet=snip))
|
||||
else:
|
||||
rows = conn.execute(
|
||||
"""
|
||||
SELECT c.id, c.title, c.tags, c.updated_at,
|
||||
count(m.id) AS n_messages, NULL AS snippet
|
||||
FROM conversations c
|
||||
LEFT JOIN messages m ON m.conversation_id = c.id
|
||||
GROUP BY c.id
|
||||
ORDER BY c.updated_at DESC
|
||||
LIMIT ?
|
||||
""",
|
||||
(limit,),
|
||||
).fetchall()
|
||||
return [
|
||||
dict(row, tags=json.loads(row["tags"]))
|
||||
for row in (dict(r) for r in rows)
|
||||
]
|
||||
|
||||
def get_conversation(self, conversation_id):
|
||||
conn = self._conn()
|
||||
conv = conn.execute(
|
||||
"SELECT * FROM conversations WHERE id = ?", (conversation_id,)
|
||||
).fetchone()
|
||||
if conv is None:
|
||||
raise ValueError(f"unknown conversation {conversation_id}")
|
||||
rows = conn.execute(
|
||||
"SELECT id, parent_id, role, content, meta, frames_file, created_at "
|
||||
"FROM messages WHERE conversation_id = ? ORDER BY id",
|
||||
(conversation_id,),
|
||||
).fetchall()
|
||||
messages = [
|
||||
{
|
||||
"id": row["id"],
|
||||
"parent_id": row["parent_id"],
|
||||
"role": row["role"],
|
||||
"content": row["content"],
|
||||
"meta": json.loads(row["meta"]) if row["meta"] else None,
|
||||
"has_frames": row["frames_file"] is not None,
|
||||
"created_at": row["created_at"],
|
||||
}
|
||||
for row in rows
|
||||
]
|
||||
return {
|
||||
"id": conv["id"],
|
||||
"title": conv["title"],
|
||||
"tags": json.loads(conv["tags"]),
|
||||
"created_at": conv["created_at"],
|
||||
"updated_at": conv["updated_at"],
|
||||
"messages": messages,
|
||||
}
|
||||
|
||||
def add_message(self, conversation_id, parent_id, role, content, meta=None):
|
||||
conn = self._conn()
|
||||
cur = conn.execute(
|
||||
"INSERT INTO messages (conversation_id, parent_id, role, content, meta, created_at) "
|
||||
"VALUES (?, ?, ?, ?, ?, ?)",
|
||||
(
|
||||
conversation_id,
|
||||
parent_id,
|
||||
role,
|
||||
content,
|
||||
json.dumps(meta, ensure_ascii=False) if meta else None,
|
||||
_now(),
|
||||
),
|
||||
)
|
||||
conn.execute(
|
||||
"UPDATE conversations SET updated_at = ? WHERE id = ?",
|
||||
(_now(), conversation_id),
|
||||
)
|
||||
conn.commit()
|
||||
return cur.lastrowid
|
||||
|
||||
def get_message(self, message_id):
|
||||
conn = self._conn()
|
||||
row = conn.execute(
|
||||
"SELECT id, conversation_id, parent_id, role, content, meta, frames_file "
|
||||
"FROM messages WHERE id = ?",
|
||||
(message_id,),
|
||||
).fetchone()
|
||||
if row is None:
|
||||
raise ValueError(f"unknown message {message_id}")
|
||||
return dict(row)
|
||||
|
||||
def update_message(self, message_id, content, meta=None):
|
||||
"""Rewrite a message's content (assistant edit, or a continuation
|
||||
appending to it). The FTS index follows via the update trigger."""
|
||||
conn = self._conn()
|
||||
row = conn.execute(
|
||||
"SELECT conversation_id FROM messages WHERE id = ?", (message_id,)
|
||||
).fetchone()
|
||||
if row is None:
|
||||
raise ValueError(f"unknown message {message_id}")
|
||||
if meta is not None:
|
||||
conn.execute(
|
||||
"UPDATE messages SET content = ?, meta = ? WHERE id = ?",
|
||||
(content, json.dumps(meta, ensure_ascii=False), message_id),
|
||||
)
|
||||
else:
|
||||
conn.execute(
|
||||
"UPDATE messages SET content = ? WHERE id = ?", (content, message_id)
|
||||
)
|
||||
conn.execute(
|
||||
"UPDATE conversations SET updated_at = ? WHERE id = ?",
|
||||
(_now(), row["conversation_id"]),
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
def path_to_root(self, message_id):
|
||||
conn = self._conn()
|
||||
path = []
|
||||
current = message_id
|
||||
while current is not None:
|
||||
row = conn.execute(
|
||||
"SELECT id, parent_id, role, content FROM messages WHERE id = ?",
|
||||
(current,),
|
||||
).fetchone()
|
||||
if row is None:
|
||||
break
|
||||
path.append({"role": row["role"], "content": row["content"]})
|
||||
current = row["parent_id"]
|
||||
path.reverse()
|
||||
return path
|
||||
|
||||
def save_frames(self, message_id, frames, layers, k):
|
||||
vocab = {}
|
||||
packed = []
|
||||
for frame in frames:
|
||||
vocab[frame["token_id"]] = frame["tok"]
|
||||
entry = {
|
||||
"pos": frame["pos"],
|
||||
"phase": frame["phase"],
|
||||
"token_id": frame["token_id"],
|
||||
"layers": {},
|
||||
}
|
||||
for layer, d in frame["layers"].items():
|
||||
for tid, s in zip(d["ids"], d["strs"]):
|
||||
vocab[tid] = s
|
||||
for tid, s in zip(d["m_ids"], d["m_strs"]):
|
||||
vocab[tid] = s
|
||||
entry["layers"][layer] = {
|
||||
"ids": np.asarray(d["ids"], np.int32).tobytes(),
|
||||
"p": np.asarray(d["p"], np.float16).tobytes(),
|
||||
"m_ids": np.asarray(d["m_ids"], np.int32).tobytes(),
|
||||
"m_p": np.asarray(d["m_p"], np.float16).tobytes(),
|
||||
"m_rank": np.asarray(d["m_rank"], np.int32).tobytes(),
|
||||
}
|
||||
packed.append(entry)
|
||||
blob = msgpack.packb(
|
||||
{
|
||||
"version": 1,
|
||||
"k": k,
|
||||
# generation id of the server-side residual store: lets pins
|
||||
# keep working after a page reload (same server session); a
|
||||
# restarted server simply reports the store as expired. Last
|
||||
# frame: after a continuation merge it's the freshest gen.
|
||||
"gen": frames[-1].get("gen") if frames else None,
|
||||
"layers": [int(l) for l in layers],
|
||||
"frames": packed,
|
||||
"vocab": {str(t): s for t, s in vocab.items()},
|
||||
}
|
||||
)
|
||||
filename = f"{message_id}.msgpack"
|
||||
(FRAMES_DIR / filename).write_bytes(blob)
|
||||
conn = self._conn()
|
||||
conn.execute(
|
||||
"UPDATE messages SET frames_file = ? WHERE id = ?", (filename, message_id)
|
||||
)
|
||||
conn.commit()
|
||||
return filename
|
||||
|
||||
def load_frames(self, message_id):
|
||||
conn = self._conn()
|
||||
row = conn.execute(
|
||||
"SELECT frames_file FROM messages WHERE id = ?", (message_id,)
|
||||
).fetchone()
|
||||
if row is None or row["frames_file"] is None:
|
||||
raise ValueError(f"no frames for message {message_id}")
|
||||
data = msgpack.unpackb((FRAMES_DIR / row["frames_file"]).read_bytes())
|
||||
vocab = data["vocab"]
|
||||
frames = []
|
||||
for entry in data["frames"]:
|
||||
frame = {
|
||||
"type": "frame",
|
||||
"phase": entry["phase"],
|
||||
"pos": entry["pos"],
|
||||
"token_id": entry["token_id"],
|
||||
"tok": vocab.get(str(entry["token_id"]), ""),
|
||||
"gen": data.get("gen"),
|
||||
"layers": {},
|
||||
}
|
||||
for layer, d in entry["layers"].items():
|
||||
ids = np.frombuffer(d["ids"], np.int32).tolist()
|
||||
m_ids = np.frombuffer(d["m_ids"], np.int32).tolist()
|
||||
frame["layers"][layer] = {
|
||||
"ids": ids,
|
||||
"p": [round(float(v), 5) for v in np.frombuffer(d["p"], np.float16)],
|
||||
"strs": [vocab.get(str(t), "") for t in ids],
|
||||
"m_ids": m_ids,
|
||||
"m_p": [round(float(v), 5) for v in np.frombuffer(d["m_p"], np.float16)],
|
||||
"m_rank": np.frombuffer(d["m_rank"], np.int32).tolist(),
|
||||
"m_strs": [vocab.get(str(t), "") for t in m_ids],
|
||||
}
|
||||
frames.append(frame)
|
||||
return {"k": data["k"], "layers": data["layers"], "frames": frames}
|
||||
|
||||
def export(self, conversation_id, fmt="json", include_frames=False):
|
||||
conv = self.get_conversation(conversation_id)
|
||||
if include_frames:
|
||||
for message in conv["messages"]:
|
||||
if message["has_frames"]:
|
||||
try:
|
||||
message["frames"] = self.load_frames(message["id"])
|
||||
except ValueError:
|
||||
pass
|
||||
if fmt == "json":
|
||||
return json.dumps(conv, ensure_ascii=False, indent=1), "application/json"
|
||||
lines = [f"# {conv['title']}", ""]
|
||||
if conv["tags"]:
|
||||
lines.append(f"tags: {', '.join(conv['tags'])}")
|
||||
lines.append("")
|
||||
for message in conv["messages"]:
|
||||
meta = message.get("meta") or {}
|
||||
head = f"**{message['role']}** (#{message['id']}"
|
||||
if message["parent_id"] is not None:
|
||||
head += f" ← #{message['parent_id']}"
|
||||
head += ")"
|
||||
if meta.get("model_id"):
|
||||
head += f" — {meta['model_id']} · {meta.get('quant') or meta.get('dtype')}"
|
||||
lines.append(head)
|
||||
lines.append("")
|
||||
lines.append(message["content"])
|
||||
lines.append("")
|
||||
if include_frames and message.get("frames"):
|
||||
lines.append(f"> {len(message['frames']['frames'])} lens frames (layers {message['frames']['layers']})")
|
||||
lines.append("")
|
||||
return "\n".join(lines), "text/markdown"
|
||||
@@ -0,0 +1,33 @@
|
||||
# ── J-Wash Python dependencies ───────────────────────────────────────────────
|
||||
# An NVIDIA GPU (CUDA) is required.
|
||||
#
|
||||
# 1. Install PyTorch FIRST, matching your CUDA version, from
|
||||
# https://pytorch.org/get-started/locally/ — for example:
|
||||
# pip install torch --index-url https://download.pytorch.org/whl/cu124
|
||||
# (Blackwell GPUs such as the RTX 50-series need a cu128+ build.)
|
||||
#
|
||||
# 2. The Jacobian-lens library ("jlens") is NOT on PyPI. Clone it and install it
|
||||
# editable (see the README for details):
|
||||
# git clone https://github.com/anthropics/jacobian-lens vendor/jacobian-lens
|
||||
# pip install -e vendor/jacobian-lens
|
||||
#
|
||||
# 3. Then install everything below:
|
||||
# pip install -r requirements.txt
|
||||
|
||||
transformers>=5.5
|
||||
accelerate>=1.0
|
||||
huggingface-hub>=0.26
|
||||
safetensors>=0.4
|
||||
|
||||
fastapi>=0.115
|
||||
uvicorn[standard]>=0.30
|
||||
websockets>=13
|
||||
pydantic>=2.0
|
||||
|
||||
numpy>=1.26
|
||||
msgpack>=1.0
|
||||
nvidia-ml-py>=12.0
|
||||
|
||||
# Optional — only needed for int8 / nf4 quantized loading (the "quant" option).
|
||||
# bitsandbytes wheels vary by platform/CUDA, so install it manually if you use it:
|
||||
# pip install bitsandbytes
|
||||
@@ -0,0 +1,56 @@
|
||||
"""Run the J-Wash server."""
|
||||
import argparse
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Run the J-Wash server.")
|
||||
parser.add_argument(
|
||||
"--hf-cache", metavar="PATH", default=None,
|
||||
help="Hugging Face cache directory. Default: the shared HF cache "
|
||||
"(HF_HOME, else ~/.cache/huggingface). Pass a path — e.g. ./hf_cache — "
|
||||
"to keep downloads isolated in a project-local cache.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--data-dir", metavar="PATH", default=None,
|
||||
help="Runtime data directory (history, frames, presets, edits). "
|
||||
"Default: ./data. Give each instance its own when running several "
|
||||
"servers side by side.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--port", type=int, default=None,
|
||||
help="HTTP port (default: 8381). Use a distinct port per instance.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
# Set the env overrides BEFORE importing config (which reads them).
|
||||
if args.data_dir:
|
||||
os.environ["JWASH_DATA_DIR"] = str(Path(args.data_dir).expanduser().resolve())
|
||||
if args.hf_cache:
|
||||
os.environ["HF_HOME"] = str(Path(args.hf_cache).expanduser().resolve())
|
||||
else:
|
||||
# HF cache chosen in the Options tab (data/settings.json); the
|
||||
# --hf-cache flag wins over it, the plain HF_HOME env loses to it.
|
||||
import json
|
||||
data_dir = Path(
|
||||
os.environ.get("JWASH_DATA_DIR") or Path(__file__).resolve().parent / "data"
|
||||
)
|
||||
try:
|
||||
saved = json.loads((data_dir / "settings.json").read_text(encoding="utf-8"))
|
||||
if saved.get("hf_cache"):
|
||||
os.environ["HF_HOME"] = str(Path(saved["hf_cache"]).expanduser().resolve())
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
import config
|
||||
|
||||
config.setup_env()
|
||||
|
||||
import uvicorn
|
||||
|
||||
uvicorn.run("api.app:app", host=config.HOST, port=args.port or config.PORT, log_level="info")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -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"]}
|
||||
]
|
||||
}
|
||||
@@ -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()
|
||||
+344
@@ -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()
|
||||
@@ -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")
|
||||
@@ -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()
|
||||
@@ -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())
|
||||
@@ -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()
|
||||
@@ -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)
|
||||
@@ -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())
|
||||
@@ -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())
|
||||
@@ -0,0 +1,6 @@
|
||||
@echo off
|
||||
REM Launch the J-Wash server on http://localhost:8381
|
||||
REM Activate your Python environment first (e.g. `conda activate jwash`),
|
||||
REM then run this script — or simply `python -X utf8 run.py`.
|
||||
cd /d %~dp0
|
||||
python -X utf8 run.py
|
||||
@@ -0,0 +1,13 @@
|
||||
# Stop the J-Wash servers: backend (8381) and the Vite dev server (5173).
|
||||
$ports = @(8381, 5173)
|
||||
foreach ($port in $ports) {
|
||||
$conns = Get-NetTCPConnection -LocalPort $port -State Listen -ErrorAction SilentlyContinue
|
||||
foreach ($conn in $conns) {
|
||||
$proc = Get-Process -Id $conn.OwningProcess -ErrorAction SilentlyContinue
|
||||
if ($proc) {
|
||||
Write-Host "port ${port}: stopping $($proc.ProcessName) (PID $($proc.Id))"
|
||||
Stop-Process -Id $proc.Id -Force -Confirm:$false
|
||||
}
|
||||
}
|
||||
if (-not $conns) { Write-Host "port ${port}: nothing to stop" }
|
||||
}
|
||||
@@ -0,0 +1,12 @@
|
||||
<!doctype html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<meta charset="UTF-8" />
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0" />
|
||||
<title>J-Wash</title>
|
||||
</head>
|
||||
<body>
|
||||
<div id="root"></div>
|
||||
<script type="module" src="/src/main.jsx"></script>
|
||||
</body>
|
||||
</html>
|
||||
Generated
+2019
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,22 @@
|
||||
{
|
||||
"name": "jlens-ui",
|
||||
"private": true,
|
||||
"version": "0.1.0",
|
||||
"type": "module",
|
||||
"scripts": {
|
||||
"dev": "vite",
|
||||
"build": "vite build",
|
||||
"preview": "vite preview"
|
||||
},
|
||||
"dependencies": {
|
||||
"d3": "^7.9.0",
|
||||
"dompurify": "^3.4.12",
|
||||
"marked": "^18.0.6",
|
||||
"react": "^18.3.1",
|
||||
"react-dom": "^18.3.1"
|
||||
},
|
||||
"devDependencies": {
|
||||
"@vitejs/plugin-react": "^4.3.4",
|
||||
"vite": "^5.4.11"
|
||||
}
|
||||
}
|
||||
+1730
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,82 @@
|
||||
import { useMemo } from 'react'
|
||||
import { fmtTok } from './tok'
|
||||
|
||||
const MAX_COLS = 200
|
||||
|
||||
function top1(frame, layer) {
|
||||
const d = frame.layers[layer]
|
||||
if (!d) return null
|
||||
return { str: fmtTok(d.m_strs[0]), rank: d.m_rank[0], alt: d.m_strs.slice(0, 3).map(fmtTok).join(', ') }
|
||||
}
|
||||
|
||||
export default function LensDiff({ framesA, framesB, labelA, labelB, onClose }) {
|
||||
const layers = useMemo(() => {
|
||||
if (!framesA.length || !framesB.length) return []
|
||||
const a = new Set(Object.keys(framesA[0].layers))
|
||||
return Object.keys(framesB[0].layers).filter((l) => a.has(l)).map(Number).sort((x, y) => x - y)
|
||||
}, [framesA, framesB])
|
||||
|
||||
const cols = Math.min(framesA.length, framesB.length, MAX_COLS)
|
||||
if (!cols || !layers.length) {
|
||||
return (
|
||||
<div className="lensview">
|
||||
<div className="lv-controls">
|
||||
diff not possible: incompatible layers or frames
|
||||
<button style={{ marginLeft: 'auto' }} onClick={onClose}>close</button>
|
||||
</div>
|
||||
</div>
|
||||
)
|
||||
}
|
||||
|
||||
let diffCount = 0
|
||||
const rows = layers.map((layer) => {
|
||||
const cells = []
|
||||
for (let i = 0; i < cols; i++) {
|
||||
const a = top1(framesA[i], String(layer))
|
||||
const b = top1(framesB[i], String(layer))
|
||||
const same = a && b && a.str === b.str
|
||||
if (!same) diffCount++
|
||||
cells.push({ a, b, same, tokA: framesA[i].tok, tokB: framesB[i].tok })
|
||||
}
|
||||
return { layer, cells }
|
||||
})
|
||||
|
||||
return (
|
||||
<div className="lensview">
|
||||
<div className="lv-controls">
|
||||
<span>diff: <b>A</b> = {labelA} · <b>B</b> = {labelB}</span>
|
||||
<span className="lv-sep" />
|
||||
<span>{diffCount} divergent cells / {cols * layers.length}</span>
|
||||
<button style={{ marginLeft: 'auto' }} onClick={onClose}>close</button>
|
||||
</div>
|
||||
<div className="lv-scroll">
|
||||
<table className="diff-table">
|
||||
<thead>
|
||||
<tr>
|
||||
<th></th>
|
||||
{Array.from({ length: cols }, (_, i) => (
|
||||
<th key={i} title={`A: ${framesA[i].tok} · B: ${framesB[i].tok}`}>
|
||||
{fmtTok(framesA[i].tok).slice(0, 6) || '·'}
|
||||
</th>
|
||||
))}
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
{rows.map(({ layer, cells }) => (
|
||||
<tr key={layer}>
|
||||
<td className="diff-lay">L{layer}</td>
|
||||
{cells.map((c, i) => (
|
||||
<td key={i} className={c.same ? 'diff-same' : 'diff-diff'}
|
||||
title={`A(${c.tokA}): ${c.a?.alt || '—'}\nB(${c.tokB}): ${c.b?.alt || '—'}`}>
|
||||
<div className="diff-a">{c.a?.str.slice(0, 7) || '—'}</div>
|
||||
<div className="diff-b">{c.b?.str.slice(0, 7) || '—'}</div>
|
||||
</td>
|
||||
))}
|
||||
</tr>
|
||||
))}
|
||||
</tbody>
|
||||
</table>
|
||||
</div>
|
||||
</div>
|
||||
)
|
||||
}
|
||||
@@ -0,0 +1,808 @@
|
||||
import { useEffect, useMemo, useRef, useState } from 'react'
|
||||
import { fmtTok } from './tok'
|
||||
|
||||
// Default layer slice for a new rule, as fractions of the model's layer count
|
||||
// (aligned with core/ablation.py): 56 layers -> 33 to 44.
|
||||
const DEFAULT_LAYERS_FRAC_LO = 3 / 5
|
||||
const DEFAULT_LAYERS_FRAC_HI = 4 / 5
|
||||
// default radius of the auto-selected slice around an edited token's "peak"
|
||||
// layer (band = peak ± radius) — adjustable in the Options tab. More inclusive
|
||||
// = more robust edit in readthrough.
|
||||
const AUTO_LAYER_RADIUS_DEFAULT = 2
|
||||
|
||||
async function jsonFetch(url, options) {
|
||||
const res = await fetch(url, options)
|
||||
const body = await res.json().catch(() => ({}))
|
||||
if (!res.ok) throw new Error(body.detail || res.statusText)
|
||||
return body
|
||||
}
|
||||
|
||||
const patchJson = (url, body) =>
|
||||
jsonFetch(url, { method: 'PATCH', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify(body) })
|
||||
|
||||
/* Layer selector: clickable cells, shift+click = range, shortcuts. */
|
||||
export function LayerPicker({ all, value, onChange, compact, defaults, fitted }) {
|
||||
const [anchor, setAnchor] = useState(null)
|
||||
const set = useMemo(() => new Set(value), [value])
|
||||
// "paint" drag: the state (on/off) is fixed by the first clicked layer, then
|
||||
// applied to the hovered layers as long as the button stays held.
|
||||
const dragRef = useRef(null) // { turnOn, sel } during the drag
|
||||
useEffect(() => {
|
||||
const up = () => { dragRef.current = null }
|
||||
window.addEventListener('mouseup', up)
|
||||
return () => window.removeEventListener('mouseup', up)
|
||||
}, [])
|
||||
if (!all.length) return null
|
||||
|
||||
function onCellDown(layer, ev) {
|
||||
ev.preventDefault() // prevents text selection during the drag
|
||||
if (ev.shiftKey && anchor != null) {
|
||||
const [lo, hi] = anchor < layer ? [anchor, layer] : [layer, anchor]
|
||||
const range = all.filter((l) => l >= lo && l <= hi)
|
||||
const turnOn = !set.has(layer)
|
||||
const next = new Set(set)
|
||||
range.forEach((l) => (turnOn ? next.add(l) : next.delete(l)))
|
||||
setAnchor(layer)
|
||||
onChange([...next].sort((a, b) => a - b))
|
||||
return
|
||||
}
|
||||
const turnOn = !set.has(layer)
|
||||
const sel = new Set(set)
|
||||
turnOn ? sel.add(layer) : sel.delete(layer)
|
||||
dragRef.current = { turnOn, sel }
|
||||
setAnchor(layer)
|
||||
onChange([...sel].sort((a, b) => a - b))
|
||||
}
|
||||
|
||||
function onCellEnter(layer) {
|
||||
const d = dragRef.current
|
||||
if (!d) return
|
||||
if (d.turnOn === d.sel.has(layer)) return // already in the desired state
|
||||
d.turnOn ? d.sel.add(layer) : d.sel.delete(layer)
|
||||
onChange([...d.sel].sort((a, b) => a - b))
|
||||
}
|
||||
|
||||
return (
|
||||
<div className={`layerpicker ${compact ? 'lp-compact' : ''}`}>
|
||||
<div className="lp-cells">
|
||||
{all.map((l) => (
|
||||
<span
|
||||
key={l}
|
||||
className={`lp-cell ${set.has(l) ? 'on' : ''} ${fitted && !fitted.has(l) ? 'lp-approx' : ''}`}
|
||||
title={`layer ${l} (click-drag = paint, shift+click = range)${fitted && !fitted.has(l) ? ' — outside the lens: direct logit lens (approx.)' : ''}`}
|
||||
onMouseDown={(e) => onCellDown(l, e)}
|
||||
onMouseEnter={() => onCellEnter(l)}
|
||||
>{l}</span>
|
||||
))}
|
||||
</div>
|
||||
{!compact && (
|
||||
<div className="lp-quick">
|
||||
{defaults?.length > 0 && <button onClick={() => onChange(defaults)}>default</button>}
|
||||
<button onClick={() => onChange([...all])}>all</button>
|
||||
<button onClick={() => onChange([])}>none</button>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
)
|
||||
}
|
||||
|
||||
/* Mini-bar: one segment per available layer, filled if the rule is active there. */
|
||||
function RuleLayerBar({ all, layers, onClick }) {
|
||||
const set = new Set(layers)
|
||||
return (
|
||||
<div
|
||||
className="rulebar"
|
||||
title={layers.length ? `layers ${layers.join(', ')} — click to edit` : 'no layer — inactive rule, click to edit'}
|
||||
onClick={onClick}
|
||||
>
|
||||
{all.map((l) => <span key={l} className={`rb-seg ${set.has(l) ? 'on' : ''}`} />)}
|
||||
</div>
|
||||
)
|
||||
}
|
||||
|
||||
/* Token field with live resolution (debounce) and clickable candidates. */
|
||||
function TokenField({ label, value, onChange, placeholder }) {
|
||||
const [cands, setCands] = useState([])
|
||||
const timerRef = useRef(null)
|
||||
|
||||
// clears the suggestion when the field is reset by the parent (after an add,
|
||||
// or an add from a view): otherwise the old candidate stays displayed
|
||||
useEffect(() => {
|
||||
if (!value.text.trim()) setCands([])
|
||||
}, [value.text])
|
||||
|
||||
function lookup(text) {
|
||||
clearTimeout(timerRef.current)
|
||||
if (!text.trim()) { setCands([]); return }
|
||||
timerRef.current = setTimeout(async () => {
|
||||
try {
|
||||
const body = await jsonFetch(`/api/token-lookup?q=${encodeURIComponent(text.trim())}`)
|
||||
setCands(body.candidates)
|
||||
const preferred = body.candidates.find((c) => c.str.startsWith(' ')) || body.candidates[0]
|
||||
if (preferred) onChange({ text, id: preferred.id, str: preferred.str })
|
||||
} catch { setCands([]) }
|
||||
}, 280)
|
||||
}
|
||||
|
||||
return (
|
||||
<>
|
||||
<div className="row"><label>{label}</label>
|
||||
<input type="text" value={value.text} placeholder={placeholder}
|
||||
onChange={(e) => { onChange({ text: e.target.value, id: null, str: '' }); lookup(e.target.value) }} />
|
||||
</div>
|
||||
{cands.length > 0 && (
|
||||
<div className="ed-cands">
|
||||
{cands.map((c) => (
|
||||
<button key={c.id} className={value.id === c.id ? 'ed-cand-on' : ''}
|
||||
onClick={() => { onChange({ text: value.text, id: c.id, str: c.str }); }}>
|
||||
{fmtTok(c.str)}
|
||||
</button>
|
||||
))}
|
||||
{value.id == null && <span className="src">no single token — multi-token word?</span>}
|
||||
</div>
|
||||
)}
|
||||
</>
|
||||
)
|
||||
}
|
||||
|
||||
// Layers as compact ranges: [20,21,22,24] → "20-22, 24".
|
||||
export function formatRanges(layers) {
|
||||
const s = [...layers].sort((a, b) => a - b)
|
||||
const out = []
|
||||
let start = null, prev = null
|
||||
for (const l of s) {
|
||||
if (start === null) { start = prev = l; continue }
|
||||
if (l === prev + 1) { prev = l; continue }
|
||||
out.push(start === prev ? `${start}` : `${start}-${prev}`)
|
||||
start = prev = l
|
||||
}
|
||||
if (start !== null) out.push(start === prev ? `${start}` : `${start}-${prev}`)
|
||||
return out.join(', ')
|
||||
}
|
||||
|
||||
// Multi-line tooltip for a rule: full (untruncated) tokens, ids, mode, factor,
|
||||
// layers — especially useful for replacement (words get cut off in the row).
|
||||
function ruleTitle(r) {
|
||||
const lines = [`token: "${fmtTok(r.token)}" (id ${r.token_id})`]
|
||||
if (r.mode === 'replace') {
|
||||
lines.push(`replacement: "${fmtTok(r.replacement)}" (id ${r.replacement_id})`)
|
||||
}
|
||||
lines.push(`mode: ${r.mode === 'replace' ? 'replace' : 'scale'} × ${r.factor}`)
|
||||
const ls = r.layers || []
|
||||
lines.push(`layers (${ls.length}): ${ls.length ? formatRanges(ls) : 'none → inactive'}`)
|
||||
if (r.enabled === false) lines.push('— rule disabled —')
|
||||
return lines.join('\n')
|
||||
}
|
||||
|
||||
// Two-position toggle: steering (exploration) ↔ the pure-weights mode the
|
||||
// architecture supports (read projection, or global projection on write-norm
|
||||
// models like Gemma). The "exact" mode stays reachable through the API; if it
|
||||
// is active, the thumb sits on the pure side and a click brings it back.
|
||||
function ModeToggle({ mode, onChange, pureMode = 'readthrough' }) {
|
||||
const isPure = mode !== 'standard'
|
||||
return (
|
||||
<div className={`mode-toggle ${isPure ? 'mt-read' : 'mt-steer'}`} role="radiogroup"
|
||||
aria-label="intervention mode">
|
||||
<div className="mt-thumb" />
|
||||
<button type="button" className={`mt-opt ${!isPure ? 'mt-on' : ''}`}
|
||||
aria-pressed={!isPure}
|
||||
onClick={() => mode !== 'standard' && onChange('standard')}>
|
||||
<svg viewBox="0 0 14 14" width="15" height="15" aria-hidden="true">
|
||||
<path d="M2.5 1v5M2.5 10.6V13M7 1v1.4M7 7V13M11.5 1v7.4M11.5 13v-2"
|
||||
stroke="currentColor" strokeWidth="1.4" strokeLinecap="round" fill="none" />
|
||||
<circle cx="2.5" cy="8.3" r="1.7" fill="currentColor" />
|
||||
<circle cx="7" cy="4.7" r="1.7" fill="currentColor" />
|
||||
<circle cx="11.5" cy="10.7" r="1.7" fill="currentColor" />
|
||||
</svg>
|
||||
<span className="mt-lab">Per-layer steering</span>
|
||||
<span className="mt-sub">preview only</span>
|
||||
</button>
|
||||
<button type="button" className={`mt-opt ${isPure ? 'mt-on' : ''}`}
|
||||
aria-pressed={isPure}
|
||||
onClick={() => mode !== pureMode && onChange(pureMode)}>
|
||||
<svg viewBox="0 0 14 14" width="15" height="15" aria-hidden="true">
|
||||
<path d="M1.2 7S3.4 3.2 7 3.2 12.8 7 12.8 7 10.6 10.8 7 10.8 1.2 7 1.2 7Z"
|
||||
fill="none" stroke="currentColor" strokeWidth="1.3" strokeLinejoin="round" />
|
||||
<circle cx="7" cy="7" r="1.9" fill="currentColor" />
|
||||
</svg>
|
||||
<span className="mt-lab">{pureMode === 'abliteration' ? 'Global projection' : 'Read projection'}</span>
|
||||
<span className="mt-sub">pure-weights · exportable</span>
|
||||
</button>
|
||||
</div>
|
||||
)
|
||||
}
|
||||
|
||||
const MODE_INFO = {
|
||||
standard: {
|
||||
label: 'Per-layer steering (preview only)',
|
||||
tag: 'not exportable',
|
||||
help: 'J-space hooks on the chosen layers: the most expressive way to explore, '
|
||||
+ 'but the export bake only captures ~1-2 % of the effect. Switch to '
|
||||
+ '"read projection" to export what you see.',
|
||||
exportHelp: 'no export in this mode — switch to "read projection" for a '
|
||||
+ 'faithful pure-weights bake.',
|
||||
},
|
||||
readthrough: {
|
||||
label: 'Read projection (faithful bake)',
|
||||
tag: 'pure-weights',
|
||||
help: 'every read of the residual downstream of the chosen layers (q/k/v, gate/up, lm_head) '
|
||||
+ 'sees the transformed residual: the preview = the exported checkpoint. Recommended for '
|
||||
+ 'removals and replacements. Regenerate after a change.',
|
||||
exportHelp: 'change of basis of the downstream reads + lm_head (untied if embeddings are tied). '
|
||||
+ 'Formats: full checkpoint (standard safetensors), modified layers, or LoRA '
|
||||
+ '(exact low-rank diff vs the original weights).',
|
||||
},
|
||||
exact: {
|
||||
label: 'Exact compensated (soft factors)',
|
||||
tag: 'pure-weights',
|
||||
help: 'read projection + counter-transform of the downstream writes: reproduces a '
|
||||
+ 'hook applied exactly once. ⚠ a full zap/replace makes the inverse singular '
|
||||
+ '(regularized ≈ read projection) — reserve this mode for partial factors.',
|
||||
exportHelp: 'downstream reads + writes transformed, lm_head untied if needed. '
|
||||
+ 'Formats: full checkpoint, modified layers, or LoRA (exact low-rank diff).',
|
||||
},
|
||||
abliteration: {
|
||||
label: 'Global projection (W_U abliteration)',
|
||||
tag: 'pure-weights',
|
||||
help: 'W_U projection on every residual write (embed + all layers): the pure-weights '
|
||||
+ 'mode for architectures where the read projection is unavailable (write norms, '
|
||||
+ 'Gemma style). Faithful for full removals/replacements; the rules\' layers are '
|
||||
+ 'ignored (global projection).',
|
||||
exportHelp: 'global abliteration × scale: removes/redirects the direction in embed + '
|
||||
+ 'every o_proj/down_proj. Formats: full checkpoint, modified layers, or LoRA '
|
||||
+ '(exact delta; embed omitted if embeddings are tied). ⚠ amplifying (factor > 1) '
|
||||
+ 'stays approximate.',
|
||||
},
|
||||
}
|
||||
|
||||
export default function Editor({
|
||||
open, onClose, rules, scale, mode, lensMeta, nLayers, genId, busy,
|
||||
prefill, onPrefillConsumed, onRules, onScale, onMode, onNotice,
|
||||
rebaseSupported = true, autoLayerRadius, llamaCppSet = false, ggufState,
|
||||
}) {
|
||||
// pure-weights mode this architecture can bake (cf. ModeToggle)
|
||||
const pureMode = rebaseSupported === false ? 'abliteration' : 'readthrough'
|
||||
const layerRadius = autoLayerRadius ?? AUTO_LAYER_RADIUS_DEFAULT
|
||||
// All the model's layers; those outside the lens use the direct logit lens.
|
||||
const allLayers = useMemo(() => {
|
||||
if (nLayers) return Array.from({ length: nLayers }, (_, i) => i)
|
||||
if (!lensMeta) return []
|
||||
if (lensMeta.fitted_layers_all?.length) return lensMeta.fitted_layers_all
|
||||
const [lo, hi] = lensMeta.fitted_layers || [0, -1]
|
||||
return Array.from({ length: hi - lo + 1 }, (_, i) => lo + i)
|
||||
}, [nLayers, lensMeta])
|
||||
|
||||
const fittedSet = useMemo(() => {
|
||||
if (!lensMeta?.fitted_layers_all?.length) return null
|
||||
return new Set(lensMeta.fitted_layers_all)
|
||||
}, [lensMeta])
|
||||
|
||||
const defaultLayers = useMemo(() => {
|
||||
const n = allLayers.length
|
||||
if (!n) return []
|
||||
const lo = Math.floor(n * DEFAULT_LAYERS_FRAC_LO)
|
||||
const hi = Math.min(Math.floor(n * DEFAULT_LAYERS_FRAC_HI), n - 1)
|
||||
return allLayers.filter((l) => l >= lo && l <= hi)
|
||||
}, [allLayers])
|
||||
|
||||
// --- optimistic factor editing + debounced PATCH with flush ---
|
||||
const [localFactors, setLocalFactors] = useState({})
|
||||
const pendingRef = useRef(new Map()) // ruleId -> {timer, body}
|
||||
|
||||
function firePatch(id) {
|
||||
const entry = pendingRef.current.get(id)
|
||||
if (!entry) return Promise.resolve()
|
||||
pendingRef.current.delete(id)
|
||||
clearTimeout(entry.timer)
|
||||
return patchJson(`/api/interventions/${id}`, entry.body)
|
||||
.then((r) => {
|
||||
onRules(r.rules)
|
||||
setLocalFactors((prev) => { const n = { ...prev }; delete n[id]; return n })
|
||||
})
|
||||
.catch((err) => {
|
||||
setLocalFactors((prev) => { const n = { ...prev }; delete n[id]; return n })
|
||||
onNotice(String(err.message || err))
|
||||
})
|
||||
}
|
||||
|
||||
function schedulePatch(id, body) {
|
||||
const prev = pendingRef.current.get(id)
|
||||
if (prev) { clearTimeout(prev.timer); body = { ...prev.body, ...body } }
|
||||
const timer = setTimeout(() => firePatch(id), 350)
|
||||
pendingRef.current.set(id, { timer, body })
|
||||
}
|
||||
|
||||
// --- global scale ---
|
||||
const [scaleEdit, setScaleEdit] = useState(null)
|
||||
const scaleTimer = useRef(null)
|
||||
const scaleShown = scaleEdit ?? scale ?? 1
|
||||
|
||||
function setGlobalScale(v) {
|
||||
setScaleEdit(v)
|
||||
clearTimeout(scaleTimer.current)
|
||||
scaleTimer.current = setTimeout(() => flushScale(v), 300)
|
||||
}
|
||||
|
||||
function flushScale(v) {
|
||||
clearTimeout(scaleTimer.current)
|
||||
scaleTimer.current = null
|
||||
return patchJson('/api/interventions', { scale: +v })
|
||||
.then((r) => { onScale(r.scale); setScaleEdit(null) })
|
||||
.catch((err) => { setScaleEdit(null); onNotice(String(err.message || err)) })
|
||||
}
|
||||
|
||||
function setMode(next) {
|
||||
patchJson('/api/interventions', { mode: next })
|
||||
.then((r) => {
|
||||
onMode(r.mode)
|
||||
onNotice(`${MODE_INFO[r.mode]?.label || r.mode} — regenerate to see the effect.`, 'ok')
|
||||
})
|
||||
.catch((err) => onNotice(String(err.message || err)))
|
||||
}
|
||||
|
||||
async function flushAll() {
|
||||
const jobs = [...pendingRef.current.keys()].map(firePatch)
|
||||
if (scaleTimer.current != null) jobs.push(flushScale(scaleEdit ?? scale ?? 1))
|
||||
await Promise.all(jobs)
|
||||
}
|
||||
|
||||
// --- multiple selection ---
|
||||
const [selected, setSelected] = useState(new Set())
|
||||
const [groupLayers, setGroupLayers] = useState([])
|
||||
const [groupFactor, setGroupFactor] = useState('')
|
||||
const selIds = [...selected].filter((id) => rules.some((r) => r.id === id))
|
||||
|
||||
async function applyGroup(body) {
|
||||
try {
|
||||
let last = null
|
||||
for (const id of selIds) last = await patchJson(`/api/interventions/${id}`, body)
|
||||
if (last) onRules(last.rules)
|
||||
onNotice(`${selIds.length} rule(s) updated`, 'ok')
|
||||
} catch (err) { onNotice(String(err.message || err)) }
|
||||
}
|
||||
|
||||
// --- per-rule layers (inline picker) ---
|
||||
const [expandedRule, setExpandedRule] = useState(null)
|
||||
|
||||
// --- add / edit form (editRuleId != null: the form UPDATES that rule) ---
|
||||
const [addToken, setAddToken] = useState({ text: '', id: null, str: '' })
|
||||
const [addRepl, setAddRepl] = useState({ text: '', id: null, str: '' })
|
||||
const [addMode, setAddMode] = useState('scale')
|
||||
const [addFactor, setAddFactor] = useState(0)
|
||||
const [addLayers, setAddLayers] = useState([])
|
||||
const [editRuleId, setEditRuleId] = useState(null)
|
||||
const [flash, setFlash] = useState(false)
|
||||
const addFormRef = useRef(null)
|
||||
|
||||
function startEditRule(r) {
|
||||
setEditRuleId(r.id)
|
||||
setAddToken({ text: (r.token || '').trim(), id: r.token_id, str: r.token })
|
||||
setAddRepl(r.replacement_id != null
|
||||
? { text: (r.replacement || '').trim(), id: r.replacement_id, str: r.replacement }
|
||||
: { text: '', id: null, str: '' })
|
||||
setAddMode(r.mode)
|
||||
setAddFactor(r.factor)
|
||||
setAddLayers(r.layers || [])
|
||||
setFlash(true)
|
||||
setTimeout(() => setFlash(false), 1600)
|
||||
setTimeout(() => addFormRef.current?.scrollIntoView({ behavior: 'smooth', block: 'center' }), 50)
|
||||
}
|
||||
|
||||
function resetAddForm() {
|
||||
setEditRuleId(null)
|
||||
setAddToken({ text: '', id: null, str: '' })
|
||||
setAddRepl({ text: '', id: null, str: '' })
|
||||
}
|
||||
|
||||
// Set the default slice ONCE per model (layer count). Definitely not on every
|
||||
// change of the defaultLayers reference: lensMeta is rebuilt on every
|
||||
// /api/status poll, and "no layer" would refill itself.
|
||||
const layersInitRef = useRef(0)
|
||||
useEffect(() => {
|
||||
if (!defaultLayers.length) return
|
||||
if (layersInitRef.current === allLayers.length) return
|
||||
layersInitRef.current = allLayers.length
|
||||
setAddLayers(defaultLayers)
|
||||
}, [defaultLayers])
|
||||
|
||||
// Auto-select the "peak" layer of the added token: as soon as a token is
|
||||
// resolved and a generation with frames exists, we ask for its per-layer ranks
|
||||
// (same data as the pins) and set the layers to peak ± 1. Silent if there is
|
||||
// no generation, the token was never seen, or the server is busy.
|
||||
useEffect(() => {
|
||||
if (addToken.id == null || genId == null) return
|
||||
let stale = false
|
||||
jsonFetch('/api/lens/pin', {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({ gen_id: genId, token_ids: [addToken.id] }),
|
||||
})
|
||||
.then((body) => {
|
||||
if (stale) return
|
||||
const d = body.pins?.[addToken.id]
|
||||
if (!d?.ranks?.length) return
|
||||
// most-relevant layer = the row that is lightest on average (the heatmap
|
||||
// metric: 1 - log10(rank+1)/5) — consistent with peakLayerOf in the J-lens
|
||||
let bestLi = 0, bestScore = -Infinity
|
||||
d.ranks.forEach((layerRanks, li) => {
|
||||
if (!layerRanks.length) return
|
||||
const score = layerRanks.reduce(
|
||||
(s, r) => s + (1 - Math.min(1, Math.log10(r + 1) / 5)), 0,
|
||||
) / layerRanks.length
|
||||
if (score > bestScore) { bestScore = score; bestLi = li }
|
||||
})
|
||||
const peak = body.layers[bestLi]
|
||||
const band = allLayers.filter((l) => Math.abs(l - peak) <= layerRadius)
|
||||
if (band.length) setAddLayers(band)
|
||||
})
|
||||
.catch(() => {})
|
||||
return () => { stale = true }
|
||||
}, [addToken.id, genId])
|
||||
|
||||
useEffect(() => {
|
||||
if (!prefill) return
|
||||
setAddToken({ text: (prefill.str || '').trim(), id: prefill.id, str: prefill.str })
|
||||
setAddMode('scale')
|
||||
setAddFactor(0)
|
||||
// Pre-select the most-relevant layer (± 1 for an effective edit), otherwise
|
||||
// keep the default slice already in place.
|
||||
if (prefill.layer != null && allLayers.length) {
|
||||
const band = allLayers.filter((l) => Math.abs(l - prefill.layer) <= layerRadius)
|
||||
if (band.length) setAddLayers(band)
|
||||
}
|
||||
setFlash(true)
|
||||
setTimeout(() => setFlash(false), 1600)
|
||||
setTimeout(() => addFormRef.current?.scrollIntoView({ behavior: 'smooth', block: 'center' }), 50)
|
||||
onPrefillConsumed()
|
||||
}, [prefill])
|
||||
|
||||
async function resolveId(field) {
|
||||
if (field.id != null) return field.id
|
||||
const body = await jsonFetch(`/api/token-lookup?q=${encodeURIComponent(field.text.trim())}`)
|
||||
if (!body.candidates.length) throw new Error(`no single token for "${field.text}"`)
|
||||
return (body.candidates.find((c) => c.str.startsWith(' ')) || body.candidates[0]).id
|
||||
}
|
||||
|
||||
async function addRule() {
|
||||
try {
|
||||
const tokenId = await resolveId(addToken)
|
||||
const replId = addMode === 'replace' ? await resolveId(addRepl) : null
|
||||
if (editRuleId != null) {
|
||||
// rewrite the existing rule in place (directions re-resolved server-side)
|
||||
const resp = await patchJson(`/api/interventions/${editRuleId}`, {
|
||||
token_id: tokenId, mode: addMode, factor: +addFactor,
|
||||
replacement_id: replId, layers: addLayers,
|
||||
})
|
||||
onRules(resp.rules)
|
||||
resetAddForm()
|
||||
onNotice('rule updated — regenerate to see the effect', 'ok')
|
||||
return
|
||||
}
|
||||
const r = await jsonFetch('/api/interventions', {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({
|
||||
token_id: tokenId, mode: addMode, factor: +addFactor,
|
||||
replacement_id: replId, layers: addLayers.length ? addLayers : null,
|
||||
}),
|
||||
})
|
||||
onRules(r.rules)
|
||||
resetAddForm()
|
||||
onNotice('rule added — regenerate to see the effect', 'ok')
|
||||
} catch (err) { onNotice(String(err.message || err)) }
|
||||
}
|
||||
|
||||
// --- presets ---
|
||||
const [presets, setPresets] = useState([])
|
||||
const [presetName, setPresetName] = useState('')
|
||||
const refreshPresets = () => jsonFetch('/api/presets').then((b) => setPresets(b.presets)).catch(() => {})
|
||||
useEffect(() => { if (open) refreshPresets() }, [open])
|
||||
|
||||
async function savePreset() {
|
||||
try {
|
||||
await flushAll() // ensures the in-flight edits are the ones being saved
|
||||
await jsonFetch(`/api/presets/${encodeURIComponent(presetName.trim())}`, { method: 'POST' })
|
||||
setPresetName('')
|
||||
refreshPresets()
|
||||
onNotice('preset saved', 'ok')
|
||||
} catch (err) { onNotice(String(err.message || err)) }
|
||||
}
|
||||
|
||||
async function applyPreset(name) {
|
||||
try {
|
||||
const r = await jsonFetch(`/api/presets/${encodeURIComponent(name)}/apply`, { method: 'POST' })
|
||||
onRules(r.rules)
|
||||
if (r.scale != null) onScale(r.scale)
|
||||
const warn = (r.warnings || []).join(' ; ')
|
||||
onNotice(warn || `preset "${name}" applied`, warn ? 'err' : 'ok')
|
||||
} catch (err) { onNotice(String(err.message || err)) }
|
||||
}
|
||||
|
||||
// --- export ---
|
||||
const [exportFmt, setExportFmt] = useState('full')
|
||||
const [exportName, setExportName] = useState('')
|
||||
const [ggufType, setGgufType] = useState('q4_k_m')
|
||||
useEffect(() => {
|
||||
// llama.cpp path removed from Options: don't leave an orphan gguf format
|
||||
if (!llamaCppSet && exportFmt === 'gguf') setExportFmt('full')
|
||||
}, [llamaCppSet, exportFmt])
|
||||
|
||||
async function doExport() {
|
||||
onNotice('exporting...', 'ok')
|
||||
try {
|
||||
await flushAll()
|
||||
if (exportFmt === 'gguf') {
|
||||
const r = await jsonFetch('/api/edit/export-gguf', {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({ name: exportName.trim(), gguf_type: ggufType }),
|
||||
})
|
||||
onNotice(
|
||||
`GGUF conversion started (checkpoint ${r.checkpoint}) — progress shown below`,
|
||||
'ok',
|
||||
)
|
||||
return
|
||||
}
|
||||
const r = await jsonFetch('/api/edit/export', {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({ format: exportFmt, name: exportName.trim() }),
|
||||
})
|
||||
const nParams = r.modified_params?.length ?? r.modified_params_count
|
||||
const warn = (r.warnings || []).join(' ; ')
|
||||
onNotice(
|
||||
`exported to ${r.out_dir} (${nParams} matrices${r.untied_lm_head ? ', lm_head untied' : ''})`
|
||||
+ (warn ? ` — ⚠ ${warn}` : ''),
|
||||
warn ? 'err' : 'ok',
|
||||
)
|
||||
if (!warn) setExportName('')
|
||||
} catch (err) { onNotice(String(err.message || err)) }
|
||||
}
|
||||
|
||||
async function cleanGgufCache() {
|
||||
try {
|
||||
const r = await jsonFetch('/api/edit/gguf-cache/delete', {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({ name: exportName.trim() }),
|
||||
})
|
||||
onNotice(`cache cleaned — ${(r.freed_bytes / 2 ** 30).toFixed(1)} GB freed`, 'ok')
|
||||
} catch (err) { onNotice(String(err.message || err)) }
|
||||
}
|
||||
|
||||
if (!open) return null
|
||||
|
||||
return (
|
||||
<div className="editor">
|
||||
<div className="ed-head">
|
||||
<span>☢ Token editor</span>
|
||||
<button className="ed-close" onClick={onClose}>✕</button>
|
||||
</div>
|
||||
|
||||
<div className="ed-body">
|
||||
<div className="ed-section">
|
||||
<h3>Global multiplier</h3>
|
||||
<div className="row">
|
||||
<input type="range" min="0" max="3" step="0.05" value={scaleShown}
|
||||
onChange={(e) => setGlobalScale(+e.target.value)} style={{ flex: 1 }} />
|
||||
<input type="number" step="0.05" value={scaleShown}
|
||||
onChange={(e) => setGlobalScale(e.target.value)}
|
||||
style={{ width: 64, flexShrink: 0 }} />
|
||||
</div>
|
||||
<div className="src">
|
||||
all alterations × {(+scaleShown).toFixed(2)}
|
||||
{+scaleShown === 1 ? ' (neutral)' : +scaleShown === 0 ? ' (all disabled)' : ''}
|
||||
</div>
|
||||
<div style={{ marginTop: 10 }}>
|
||||
<ModeToggle mode={mode || 'standard'} onChange={setMode} pureMode={pureMode} />
|
||||
</div>
|
||||
<div className="src" style={{ marginTop: 8 }}>{MODE_INFO[mode]?.help || ''}</div>
|
||||
</div>
|
||||
|
||||
<div className="ed-section">
|
||||
<h3>Active rules ({rules.length})</h3>
|
||||
{rules.length === 0 && <div className="src">none — add a token below or from the J-lens (☢)</div>}
|
||||
{rules.map((r) => (
|
||||
<div key={r.id} className={`ed-rule ${r.enabled === false ? 'ed-rule-off' : ''}`}>
|
||||
<div className="ed-rule-main">
|
||||
<input type="checkbox" checked={selected.has(r.id)}
|
||||
onChange={(e) => {
|
||||
const next = new Set(selected)
|
||||
e.target.checked ? next.add(r.id) : next.delete(r.id)
|
||||
setSelected(next)
|
||||
}} />
|
||||
<button className="ed-rule-toggle"
|
||||
title={r.enabled === false ? 'rule disabled — click to enable (layers kept)' : 'rule active — click to disable without losing the layers'}
|
||||
onClick={async () => {
|
||||
try {
|
||||
const resp = await patchJson(`/api/interventions/${r.id}`, { enabled: r.enabled === false })
|
||||
onRules(resp.rules)
|
||||
} catch (err) { onNotice(String(err.message || err)) }
|
||||
}}>{r.enabled === false ? '○' : '●'}</button>
|
||||
<span className="ed-rule-tok" title={ruleTitle(r)}>
|
||||
«{fmtTok(r.token)}»{r.mode === 'replace' ? ` → «${fmtTok(r.replacement)}»` : ''}
|
||||
</span>
|
||||
<span className="src">×</span>
|
||||
<input type="number" step="0.05" className="ed-rule-factor"
|
||||
value={localFactors[r.id] ?? r.factor}
|
||||
onChange={(e) => {
|
||||
setLocalFactors((prev) => ({ ...prev, [r.id]: e.target.value }))
|
||||
schedulePatch(r.id, { factor: +e.target.value })
|
||||
}} />
|
||||
<RuleLayerBar all={allLayers} layers={r.layers}
|
||||
onClick={() => setExpandedRule(expandedRule === r.id ? null : r.id)} />
|
||||
<button className="ed-rule-del" title="edit this rule (token, replacement, mode, factor, layers) in the form below"
|
||||
onClick={() => startEditRule(r)}>✎</button>
|
||||
<button className="ed-rule-del" title="delete this rule" onClick={async () => {
|
||||
try {
|
||||
const resp = await jsonFetch(`/api/interventions/${r.id}`, { method: 'DELETE' })
|
||||
onRules(resp.rules)
|
||||
if (editRuleId === r.id) resetAddForm()
|
||||
} catch (err) { onNotice(String(err.message || err)) }
|
||||
}}>✕</button>
|
||||
</div>
|
||||
{expandedRule === r.id && (
|
||||
<div className="ed-rule-layers">
|
||||
<LayerPicker all={allLayers} value={r.layers} defaults={defaultLayers} fitted={fittedSet}
|
||||
onChange={async (layers) => {
|
||||
try {
|
||||
const resp = await patchJson(`/api/interventions/${r.id}`, { layers })
|
||||
onRules(resp.rules)
|
||||
} catch (err) { onNotice(String(err.message || err)) }
|
||||
}} />
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
))}
|
||||
{rules.length > 1 && (
|
||||
<div className="row" style={{ marginTop: 4 }}>
|
||||
<button onClick={() => setSelected(new Set(rules.map((r) => r.id)))}>select all</button>
|
||||
{selIds.length > 0 && (
|
||||
<button onClick={() => setSelected(new Set())}>deselect all</button>
|
||||
)}
|
||||
<button onClick={async () => {
|
||||
try {
|
||||
const resp = await jsonFetch('/api/interventions', { method: 'DELETE' })
|
||||
onRules(resp.rules)
|
||||
} catch (err) { onNotice(String(err.message || err)) }
|
||||
}}>remove all</button>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
|
||||
{selIds.length > 0 && (
|
||||
<div className="ed-section ed-group">
|
||||
<h3>{selIds.length} rule(s) selected</h3>
|
||||
<div className="src">layers to apply (none = inactive rules):</div>
|
||||
<LayerPicker all={allLayers} value={groupLayers} defaults={defaultLayers} fitted={fittedSet}
|
||||
onChange={setGroupLayers} />
|
||||
<div className="row">
|
||||
<button className="primary"
|
||||
onClick={() => applyGroup({ layers: groupLayers })}>Apply layers</button>
|
||||
</div>
|
||||
<div className="row">
|
||||
<label>factor</label>
|
||||
<input type="number" step="0.05" value={groupFactor} placeholder="—"
|
||||
onChange={(e) => setGroupFactor(e.target.value)} />
|
||||
<button disabled={groupFactor === ''}
|
||||
onClick={() => applyGroup({ factor: +groupFactor })}>Apply</button>
|
||||
</div>
|
||||
<button onClick={() => setSelected(new Set())}>deselect</button>
|
||||
</div>
|
||||
)}
|
||||
|
||||
<div className={`ed-section ${flash ? 'ed-flash' : ''}`} ref={addFormRef}>
|
||||
<h3>{editRuleId != null
|
||||
? <>Edit the rule <button style={{ marginLeft: 8, fontSize: 11 }} onClick={resetAddForm}>cancel</button></>
|
||||
: 'Add a rule'}</h3>
|
||||
<TokenField label="token" value={addToken} onChange={setAddToken} placeholder="word (e.g. Euro)" />
|
||||
<div className="row"><label>mode</label>
|
||||
<select value={addMode} onChange={(e) => { setAddMode(e.target.value); setAddFactor(e.target.value === 'replace' ? 1 : 0) }}>
|
||||
<option value="scale">multiply (×0 = remove)</option>
|
||||
<option value="replace">replace with</option>
|
||||
</select>
|
||||
</div>
|
||||
{addMode === 'replace' && (
|
||||
<TokenField label="with" value={addRepl} onChange={setAddRepl} placeholder="replacement token" />
|
||||
)}
|
||||
<div className="row"><label>factor</label>
|
||||
<input type="number" step="0.05" value={addFactor} onChange={(e) => setAddFactor(e.target.value)} />
|
||||
</div>
|
||||
<div className="src">layers ({addLayers.length}):</div>
|
||||
<LayerPicker all={allLayers} value={addLayers} defaults={defaultLayers} fitted={fittedSet}
|
||||
onChange={setAddLayers} />
|
||||
<button className="primary"
|
||||
disabled={!addToken.text.trim() || (addMode === 'replace' && !addRepl.text.trim()) || !addLayers.length || !!busy}
|
||||
onClick={addRule}>{editRuleId != null ? 'Update' : 'Add'}</button>
|
||||
</div>
|
||||
|
||||
<div className="ed-section">
|
||||
<h3>Presets</h3>
|
||||
{presets.map((p) => (
|
||||
<div key={p.name} className="reg-item" title={p.model_id ? `saved for ${p.model_id}` : ''}>
|
||||
<span className="preset-info">
|
||||
<span className="preset-name">{p.name} · {p.n_rules} rule(s)</span>
|
||||
{p.model_id && <span className="src preset-model">{p.model_id.replace(/^local\//, '')}</span>}
|
||||
</span>
|
||||
<button disabled={!!busy} onClick={() => applyPreset(p.name)}>Apply</button>
|
||||
<button onClick={async () => {
|
||||
await jsonFetch(`/api/presets/${encodeURIComponent(p.name)}`, { method: 'DELETE' }).catch(() => {})
|
||||
refreshPresets()
|
||||
}}>✕</button>
|
||||
</div>
|
||||
))}
|
||||
<div className="row">
|
||||
<input type="text" placeholder="preset name" value={presetName} onChange={(e) => setPresetName(e.target.value)}
|
||||
onKeyDown={(e) => { if (e.key === 'Enter' && presetName.trim() && rules.length && !busy) savePreset() }} />
|
||||
<button disabled={!presetName.trim() || !rules.length} onClick={savePreset}>Save</button>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div className="ed-section">
|
||||
<h3>Export the edit
|
||||
<span className="exp-tag">{MODE_INFO[mode]?.tag || ''}</span>
|
||||
</h3>
|
||||
{mode === 'standard' ? (
|
||||
<div className="src exp-disabled-note" style={{ marginBottom: 6 }}>
|
||||
⚠ export disabled in "per-layer steering": no bake reproduces the
|
||||
per-layer hooks faithfully. Switch the mode to
|
||||
"{pureMode === 'abliteration' ? 'global projection' : 'read projection'}"
|
||||
above for a faithful checkpoint.
|
||||
</div>
|
||||
) : null}
|
||||
<div className={mode === 'standard' ? 'exp-grid exp-grid-off' : 'exp-grid'}>
|
||||
<div className="row"><label>format</label>
|
||||
<select value={exportFmt} onChange={(e) => setExportFmt(e.target.value)} disabled={mode === 'standard'}>
|
||||
<option value="full">full checkpoint</option>
|
||||
<option value="layers">modified layers (safetensors)</option>
|
||||
<option value="lora">LoRA (PEFT)</option>
|
||||
{llamaCppSet && <option value="gguf">GGUF (via llama.cpp)</option>}
|
||||
</select>
|
||||
</div>
|
||||
{exportFmt === 'gguf' && (
|
||||
<div className="row"><label title="bf16/f16 = plain conversion; q* = quantized with llama-quantize. The intermediate HF checkpoint is cached so other types don't re-bake.">type</label>
|
||||
<select value={ggufType} onChange={(e) => setGgufType(e.target.value)}>
|
||||
{['q4_k_m', 'q5_k_m', 'q6_k', 'q8_0', 'q3_k_m', 'bf16', 'f16'].map((t) => (
|
||||
<option key={t} value={t}>{t}</option>
|
||||
))}
|
||||
</select>
|
||||
</div>
|
||||
)}
|
||||
<div className="row"><label>name</label>
|
||||
<input type="text" placeholder="edit name" value={exportName}
|
||||
onChange={(e) => setExportName(e.target.value)} disabled={mode === 'standard'} />
|
||||
</div>
|
||||
<button className="primary"
|
||||
disabled={mode === 'standard' || !exportName.trim()
|
||||
|| (exportFmt !== 'gguf' && !rules.length) || !!busy
|
||||
|| ggufState?.state === 'running'}
|
||||
onClick={doExport}>Export</button>
|
||||
{!llamaCppSet && (
|
||||
<div className="src">tip: set the llama.cpp folder in the Options tab to
|
||||
unlock a direct GGUF export.</div>
|
||||
)}
|
||||
{exportFmt === 'gguf' && (
|
||||
<div className="row" style={{ marginTop: 2 }}>
|
||||
<button disabled={!exportName.trim() || ggufState?.state === 'running'}
|
||||
title="delete the cached intermediate HF checkpoint of this export (the .gguf files stay)"
|
||||
onClick={cleanGgufCache}>clean cache</button>
|
||||
</div>
|
||||
)}
|
||||
{ggufState?.state === 'running' && (
|
||||
<div className="src">⏳ GGUF “{ggufState.name}”: {ggufState.step}…</div>
|
||||
)}
|
||||
{ggufState?.state === 'done' && ggufState.result && (
|
||||
<div className="src ok">✔ GGUF ready: {ggufState.result.gguf}
|
||||
{' '}({(ggufState.result.size_bytes / 2 ** 30).toFixed(1)} GB) — the HF
|
||||
checkpoint stays cached for other types (clean cache to reclaim).</div>
|
||||
)}
|
||||
{ggufState?.state === 'error' && (
|
||||
<div className="src reg-reason">GGUF failed: {ggufState.error}</div>
|
||||
)}
|
||||
<div className="src">{MODE_INFO[mode]?.exportHelp || ''}</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
)
|
||||
}
|
||||
@@ -0,0 +1,900 @@
|
||||
import { useEffect, useLayoutEffect, useMemo, useRef, useState } from 'react'
|
||||
import * as d3 from 'd3'
|
||||
import { fmtTok } from './tok'
|
||||
import { LayerPicker, formatRanges } from './Editor.jsx'
|
||||
|
||||
// Mouse-following tooltip that never overflows the viewport: measured after
|
||||
// render (before paint), it flips to the left of / above the cursor if needed.
|
||||
function Tip({ x, y, children }) {
|
||||
const ref = useRef(null)
|
||||
useLayoutEffect(() => {
|
||||
const el = ref.current
|
||||
if (!el) return
|
||||
const { width, height } = el.getBoundingClientRect()
|
||||
let left = x + 12
|
||||
let top = y + 14
|
||||
if (left + width > window.innerWidth - 8) left = Math.max(8, x - width - 12)
|
||||
if (top + height > window.innerHeight - 8) top = Math.max(8, y - height - 14)
|
||||
el.style.left = `${left}px`
|
||||
el.style.top = `${top}px`
|
||||
})
|
||||
return (
|
||||
<div className="lv-tip" ref={ref} style={{ left: x + 12, top: y + 14 }}>
|
||||
{children}
|
||||
</div>
|
||||
)
|
||||
}
|
||||
|
||||
const trimTok = (s) => (s || '').replace(/^\s+|\s+$/g, '')
|
||||
|
||||
function copyText(s) {
|
||||
const clean = trimTok(s)
|
||||
if (!clean) return
|
||||
navigator.clipboard?.writeText(clean).catch(() => {})
|
||||
}
|
||||
|
||||
// Short name of the loaded lens, for the header badge ("which lens do I have?").
|
||||
function lensLabel(meta) {
|
||||
if (!meta) return ''
|
||||
if (meta.path) return meta.path.replace(/[/\\]+$/, '').split(/[/\\]/).pop()
|
||||
if (meta.filename) {
|
||||
const stem = meta.filename.split('/').pop().replace('_jacobian_lens', '').replace('.pt', '')
|
||||
return meta.revision ? `${stem} @${meta.revision}` : stem
|
||||
}
|
||||
return meta.repo_id || 'lens'
|
||||
}
|
||||
|
||||
const CELL_W = 58
|
||||
const CELL_H = 22
|
||||
const HEADER_H = 52
|
||||
const LABEL_W = 44
|
||||
const MAX_COLS = 400
|
||||
const PIN_COLORS = ['#e8a13c', '#6cb8e0', '#7ec97e', '#c98bd4', '#d0654f']
|
||||
|
||||
/* ============================================================================
|
||||
"PINS" VIEW SETTINGS (rank curves + heatmap of a pinned token)
|
||||
|
||||
⚠ After ANY change here you MUST rebuild the front-end and hard-reload the
|
||||
page with Ctrl+F5:
|
||||
cd ui && npm run build
|
||||
(without a rebuild nothing changes: the browser serves the compiled bundle,
|
||||
not this source file)
|
||||
============================================================================ */
|
||||
const PIN_FILL_WIDTH = true // true = the block stretches to the full available width
|
||||
const PIN_PX_PER_TOKEN = 6 // minimum width of a column (position), in px
|
||||
// (raise to 12-20 for wider cells: past the
|
||||
// available width, horizontal scroll kicks in)
|
||||
const PIN_MAX_W = 1600 // max graph width (px) — beyond it: horizontal scroll
|
||||
const PIN_MIN_W = 120 // minimum graph width (px)
|
||||
const PIN_CURVE_H = 90 // height of the curves graph (px)
|
||||
const PIN_ROW_H_MIN = 5 // minimum height of a row (layer) in the heatmap (px)
|
||||
const PIN_ROW_H_MAX = 12 // maximum height of a row (px)
|
||||
const PIN_MAP_H = 180 // target heatmap height: row height ≈ PIN_MAP_H / n_layers
|
||||
/* ========================================================================== */
|
||||
|
||||
function cellData(frame, layer, maskOn) {
|
||||
const d = frame.layers[layer]
|
||||
if (!d) return null
|
||||
const strs = maskOn ? d.m_strs : d.strs
|
||||
const ps = maskOn ? d.m_p : d.p
|
||||
const ids = maskOn ? d.m_ids : d.ids
|
||||
return {
|
||||
str: strs[0],
|
||||
rank: maskOn ? d.m_rank[0] : 0,
|
||||
p: ps[0],
|
||||
tid: ids[0],
|
||||
top: strs.map((s, i) => ({ s, p: ps[i], r: maskOn ? d.m_rank[i] : i, tid: ids[i] })),
|
||||
}
|
||||
}
|
||||
|
||||
/* First pinned token present in the cell's top-k (concept localization). */
|
||||
function pinHit(frame, layer, maskOn, pinnedIds) {
|
||||
const d = frame.layers[layer]
|
||||
if (!d || !pinnedIds.length) return null
|
||||
const ids = maskOn ? d.m_ids : d.ids
|
||||
for (let i = 0; i < ids.length; i++) {
|
||||
if (pinnedIds.includes(ids[i])) return { tid: ids[i], idx: i }
|
||||
}
|
||||
return null
|
||||
}
|
||||
|
||||
function rankColor(rank) {
|
||||
const t = 1 - Math.min(1, Math.log10(rank + 1) / 5)
|
||||
return d3.interpolateInferno(0.15 + 0.8 * t)
|
||||
}
|
||||
|
||||
// CJK, kana, hangul, cyrillic, arabic, hebrew...: candidates for "nearest tokens"
|
||||
const NONLATIN_RE = /[Ѐ-ӿ-ۿऀ--ᄀ-ᇿ⺀-鿿ꀀ-가-豈-︰-﹏]/
|
||||
|
||||
function parseLayerSpec(spec) {
|
||||
if (!spec.trim()) return null
|
||||
const set = new Set()
|
||||
for (const part of spec.split(',')) {
|
||||
const m = part.trim().match(/^(\d+)\s*-\s*(\d+)$/)
|
||||
if (m) for (let i = +m[1]; i <= +m[2]; i++) set.add(i)
|
||||
else if (part.trim() && !isNaN(+part.trim())) set.add(+part.trim())
|
||||
}
|
||||
return set.size ? set : null
|
||||
}
|
||||
|
||||
export default function LensView({ frames, tick, genId, lensMeta, onNotice, onEditToken, hidden, onHideToken, maxH, editorOpen, streaming }) {
|
||||
const [maskOn, setMaskOn] = useState(true)
|
||||
const [view, setView] = useState('agg')
|
||||
const [filterInput, setFilterInput] = useState('')
|
||||
const [dragLayers, setDragLayers] = useState(null)
|
||||
const [aggLimit, setAggLimit] = useState(() => {
|
||||
const v = +localStorage.getItem('jlens_agg_limit')
|
||||
return v > 0 ? v : 80
|
||||
})
|
||||
const [autoScroll, setAutoScroll] = useState(() => localStorage.getItem('jlens_autoscroll') !== '0')
|
||||
const gridScrollRef = useRef(null)
|
||||
const [pinInput, setPinInput] = useState('')
|
||||
const [pinCands, setPinCands] = useState([])
|
||||
const pinTimerRef = useRef(null)
|
||||
const [pinned, setPinned] = useState({})
|
||||
const [pinData, setPinData] = useState(null)
|
||||
const [tip, setTip] = useState(null)
|
||||
const [aggTip, setAggTip] = useState(null)
|
||||
const [pinHover, setPinHover] = useState(null) // { tid, li }: layer hovered in a pin block
|
||||
const transCache = useRef(new Map()) // tid -> [{id, str, sim}] | null (request in flight)
|
||||
const [, setTransTick] = useState(0)
|
||||
|
||||
// available width for the pin blocks (PIN_FILL_WIDTH): measured on the
|
||||
// .lv-pins container. editorOpen is an explicit dependency — the 400px
|
||||
// editor panel opening/closing is the main reason this width changes, and
|
||||
// ResizeObserver alone proved unreliable for it — plus window resizes.
|
||||
const pinsRef = useRef(null)
|
||||
const [pinsW, setPinsW] = useState(0)
|
||||
useEffect(() => {
|
||||
const el = pinsRef.current
|
||||
if (!el) return
|
||||
const measure = () => setPinsW(el.clientWidth)
|
||||
// synchronous first (reading clientWidth forces the reflow, and rAF /
|
||||
// ResizeObserver never fire in a hidden tab), then again next frame in
|
||||
// case the flex layout still settles
|
||||
measure()
|
||||
const raf = requestAnimationFrame(measure)
|
||||
const ro = new ResizeObserver(measure)
|
||||
ro.observe(el)
|
||||
window.addEventListener('resize', measure)
|
||||
return () => {
|
||||
cancelAnimationFrame(raf)
|
||||
ro.disconnect()
|
||||
window.removeEventListener('resize', measure)
|
||||
}
|
||||
}, [pinData, editorOpen])
|
||||
|
||||
useEffect(() => { localStorage.setItem('jlens_agg_limit', String(aggLimit)) }, [aggLimit])
|
||||
useEffect(() => { localStorage.setItem('jlens_autoscroll', autoScroll ? '1' : '0') }, [autoScroll])
|
||||
|
||||
// horizontal auto-scroll to the right during generation (new tokens appended
|
||||
// on the right) — Heatmap view only, toggleable. Deps without gridW (declared
|
||||
// below → TDZ): tick changes on every batch of frames and the effect reads
|
||||
// scrollWidth live at execution time.
|
||||
useEffect(() => {
|
||||
if (autoScroll && view === 'grid' && gridScrollRef.current) {
|
||||
const el = gridScrollRef.current
|
||||
el.scrollLeft = el.scrollWidth
|
||||
}
|
||||
}, [tick, autoScroll, view])
|
||||
|
||||
// manual pinning of a token by its text, with suggestions (variants with/without
|
||||
// a leading space): the pin colors the grid and draws the curves
|
||||
function lookupPin(text) {
|
||||
clearTimeout(pinTimerRef.current)
|
||||
if (!text.trim()) { setPinCands([]); return }
|
||||
pinTimerRef.current = setTimeout(async () => {
|
||||
try {
|
||||
const res = await fetch(`/api/token-lookup?q=${encodeURIComponent(text.trim())}`)
|
||||
const body = await res.json()
|
||||
setPinCands(body.candidates || [])
|
||||
} catch { setPinCands([]) }
|
||||
}, 250)
|
||||
}
|
||||
function addPin(c) {
|
||||
if (!c) return
|
||||
if (!pinned[c.id]) togglePin(c.id, c.str)
|
||||
setPinInput('')
|
||||
setPinCands([])
|
||||
}
|
||||
async function fetchTrans(entries) {
|
||||
const missing = [...new Set(
|
||||
entries.filter((e) => NONLATIN_RE.test(e.str || '')).map((e) => e.tid)
|
||||
)].filter((t) => !transCache.current.has(t))
|
||||
if (!missing.length) return
|
||||
missing.forEach((t) => transCache.current.set(t, null))
|
||||
try {
|
||||
const res = await fetch('/api/token-neighbors', {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({ token_ids: missing, k: 6 }),
|
||||
})
|
||||
if (!res.ok) throw new Error()
|
||||
const body = await res.json()
|
||||
Object.entries(body.neighbors).forEach(([tid, list]) => transCache.current.set(+tid, list))
|
||||
setTransTick((x) => x + 1)
|
||||
} catch {
|
||||
missing.forEach((t) => transCache.current.delete(t))
|
||||
}
|
||||
}
|
||||
|
||||
const transLabel = (tid) => {
|
||||
const list = transCache.current.get(tid)
|
||||
if (!list?.length) return null
|
||||
return list.slice(0, 2).map((n) => fmtTok(n.str)).join(', ')
|
||||
}
|
||||
|
||||
const shown = frames.length > MAX_COLS ? frames.slice(-MAX_COLS) : frames
|
||||
const layers = useMemo(() => {
|
||||
if (!shown.length) return []
|
||||
return Object.keys(shown[shown.length - 1].layers).map(Number).sort((a, b) => a - b)
|
||||
}, [tick, frames])
|
||||
|
||||
const filterSet = useMemo(() => parseLayerSpec(filterInput), [filterInput])
|
||||
const filtered = filterSet ? layers.filter((l) => filterSet.has(l)) : layers
|
||||
const filterEmpty = filterSet != null && filtered.length === 0
|
||||
const layersShown = filterEmpty ? layers : filtered
|
||||
|
||||
// Defer applying a layer drag-selection until the mouse is released, so the
|
||||
// Frequencies cloud reflows once (on mouseup) instead of on every dragged cell.
|
||||
const dragLayersRef = useRef(null)
|
||||
dragLayersRef.current = dragLayers
|
||||
const layersRef = useRef(layers)
|
||||
layersRef.current = layers
|
||||
useEffect(() => {
|
||||
const commit = () => {
|
||||
const sel = dragLayersRef.current
|
||||
if (sel == null) return
|
||||
dragLayersRef.current = null
|
||||
setFilterInput(sel.length >= layersRef.current.length ? '' : formatRanges(sel))
|
||||
setDragLayers(null)
|
||||
}
|
||||
window.addEventListener('mouseup', commit)
|
||||
return () => window.removeEventListener('mouseup', commit)
|
||||
}, [])
|
||||
|
||||
const pinnedIds = useMemo(() => Object.keys(pinned).map(Number).sort((a, b) => a - b), [pinned])
|
||||
const pinColor = (tid) => PIN_COLORS[pinnedIds.indexOf(tid) % PIN_COLORS.length]
|
||||
|
||||
// "Most-relevant" layer for a pinned token = where it reaches its best rank
|
||||
// (the lowest) across all positions. Used to prefill the editor on the right
|
||||
// slice rather than the default one.
|
||||
function peakLayerOf(tid) {
|
||||
const d = pinData?.pins?.[tid]
|
||||
if (!d?.ranks?.length) return null
|
||||
// "most-relevant" layer = the row that is LIGHTEST on average in the heatmap
|
||||
// (rankColor metric: 1 - log10(rank+1)/5), not the single best rank at one
|
||||
// position (a spike doesn't make the layer read as strong).
|
||||
let bestLi = 0, bestScore = -Infinity
|
||||
d.ranks.forEach((layerRanks, li) => {
|
||||
if (!layerRanks.length) return
|
||||
const score = layerRanks.reduce(
|
||||
(s, r) => s + (1 - Math.min(1, Math.log10(r + 1) / 5)), 0,
|
||||
) / layerRanks.length
|
||||
if (score > bestScore) { bestScore = score; bestLi = li }
|
||||
})
|
||||
return pinData.layers[bestLi]
|
||||
}
|
||||
|
||||
const agg = useMemo(() => {
|
||||
if (view !== 'agg') return []
|
||||
const visible = new Set(layersShown)
|
||||
const map = new Map()
|
||||
frames.forEach((f) => {
|
||||
const seenHere = new Set()
|
||||
Object.entries(f.layers).forEach(([layer, d]) => {
|
||||
if (!visible.has(+layer)) return
|
||||
const strs = maskOn ? d.m_strs : d.strs
|
||||
const ps = maskOn ? d.m_p : d.p
|
||||
const ids = maskOn ? d.m_ids : d.ids
|
||||
strs.forEach((s, i) => {
|
||||
const tid = ids[i]
|
||||
let entry = map.get(tid)
|
||||
if (!entry) {
|
||||
entry = { tid, str: s, appearances: 0, maxP: 0, peakLayer: null }
|
||||
map.set(tid, entry)
|
||||
}
|
||||
if (!seenHere.has(tid)) {
|
||||
seenHere.add(tid)
|
||||
entry.appearances++
|
||||
}
|
||||
if (ps[i] > entry.maxP) {
|
||||
entry.maxP = ps[i]
|
||||
entry.peakLayer = +layer
|
||||
}
|
||||
})
|
||||
})
|
||||
})
|
||||
return [...map.values()]
|
||||
.filter((e) => !hidden.has(trimTok(e.str)))
|
||||
.sort((a, b) => b.appearances - a.appearances || b.maxP - a.maxP)
|
||||
.slice(0, aggLimit)
|
||||
}, [tick, frames, maskOn, view, layersShown, hidden, aggLimit])
|
||||
|
||||
const gridW = shown.length * CELL_W
|
||||
const height = HEADER_H + layersShown.length * CELL_H
|
||||
|
||||
// Sequence token: two quick pin/unpin clicks race their POSTs — only the
|
||||
// LAST request may write pinData, or a stale response resurrects a token
|
||||
// that was just unpinned (ghost graph showing the raw token id).
|
||||
const pinReqRef = useRef(0)
|
||||
|
||||
async function refreshPins(nextPinned) {
|
||||
const ids = Object.keys(nextPinned).map(Number)
|
||||
setPinned(nextPinned)
|
||||
const reqId = ++pinReqRef.current
|
||||
if (!ids.length || genId == null) {
|
||||
setPinData(null)
|
||||
return
|
||||
}
|
||||
try {
|
||||
const res = await fetch('/api/lens/pin', {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({ gen_id: genId, token_ids: ids }),
|
||||
})
|
||||
const body = await res.json()
|
||||
if (!res.ok) throw new Error(body.detail || res.statusText)
|
||||
if (reqId === pinReqRef.current) setPinData(body)
|
||||
} catch (err) {
|
||||
if (reqId === pinReqRef.current) {
|
||||
onNotice?.(String(err.message || err))
|
||||
setPinData(null)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Keep the pin graphs in sync with the generation being VIEWED: refetch when
|
||||
// it changes (new message generated / another message selected) — debounced
|
||||
// on tick so a streaming generation triggers one fetch at the end, not one
|
||||
// per frame batch.
|
||||
const pinnedRef = useRef(pinned)
|
||||
pinnedRef.current = pinned
|
||||
useEffect(() => {
|
||||
// never fetch mid-stream: slow generations space their frame batches past
|
||||
// the debounce, and the server 409s /api/lens/pin while it generates —
|
||||
// when `streaming` flips back to false this effect refires and fetches once
|
||||
if (streaming) return
|
||||
if (genId == null || !Object.keys(pinnedRef.current).length) return
|
||||
const t = setTimeout(() => refreshPins(pinnedRef.current), 400)
|
||||
return () => clearTimeout(t)
|
||||
}, [genId, tick, streaming])
|
||||
|
||||
function togglePin(tid, str) {
|
||||
const next = { ...pinned }
|
||||
if (next[tid]) delete next[tid]
|
||||
else next[tid] = str
|
||||
if (genId == null) {
|
||||
// in replay: no server curves, but the grid coloring stays available
|
||||
setPinned(next)
|
||||
setPinData(null)
|
||||
return
|
||||
}
|
||||
refreshPins(next)
|
||||
}
|
||||
|
||||
// Memoized grid: doesn't depend on the tooltip (smooth hover even with 400 columns).
|
||||
const gridCells = useMemo(() => {
|
||||
if (view !== 'grid') return null
|
||||
return (
|
||||
<>
|
||||
{layersShown.map((layer, ri) => ri % 2 === 1 && (
|
||||
<rect key={`z${layer}`} x="0" y={HEADER_H + ri * CELL_H} width={gridW} height={CELL_H}
|
||||
fill="rgba(255,255,255,.025)" />
|
||||
))}
|
||||
{shown.map((f, ci) => (
|
||||
<g key={ci} transform={`translate(${ci * CELL_W},0)`}>
|
||||
<rect
|
||||
x="0" y="0" width={CELL_W - 1} height={HEADER_H - 6}
|
||||
fill={f.phase === 'reading' ? 'rgba(232,161,60,.10)' : 'rgba(108,184,224,.10)'}
|
||||
/>
|
||||
<text
|
||||
x={CELL_W / 2} y={HEADER_H - 12}
|
||||
className="lv-toklabel"
|
||||
transform={`rotate(-38 ${CELL_W / 2} ${HEADER_H - 12})`}
|
||||
>
|
||||
{fmtTok(f.tok).slice(0, 9) || '·'}
|
||||
</text>
|
||||
</g>
|
||||
))}
|
||||
{layersShown.map((layer, ri) => (
|
||||
<g key={layer} transform={`translate(0,${HEADER_H + ri * CELL_H})`}>
|
||||
{shown.map((f, ci) => {
|
||||
const c = cellData(f, String(layer), maskOn)
|
||||
if (!c) return null
|
||||
const opacity = Math.max(0.28, Math.min(1, Math.sqrt(c.p) * 2.4))
|
||||
const hit = pinHit(f, String(layer), maskOn, pinnedIds)
|
||||
return (
|
||||
<g
|
||||
key={ci}
|
||||
transform={`translate(${ci * CELL_W},0)`}
|
||||
className="lv-cell"
|
||||
onMouseMove={(e) => {
|
||||
setTip({ x: e.clientX, y: e.clientY, c, layer, ri, ci, f })
|
||||
fetchTrans(c.top.map((t) => ({ tid: t.tid, str: t.s })))
|
||||
}}
|
||||
onClick={() => togglePin(c.tid, c.str)}
|
||||
>
|
||||
{/* hover capture over the whole cell (without the 1px gap between
|
||||
visible rects). We do NOT clear the tooltip per cell: only a
|
||||
hover leaving the whole grid clears it (see lv-scroll). So
|
||||
moving from one cell to the next fires a single event (no
|
||||
intermediate null state = no flicker) and hovering a gap keeps
|
||||
the last info instead of jumping. */}
|
||||
<rect x="0" y="0" width={CELL_W} height={CELL_H} fill="transparent" />
|
||||
<rect
|
||||
x="0" y="0" width={CELL_W - 1} height={CELL_H - 1}
|
||||
fill={hit
|
||||
? d3.color(pinColor(hit.tid)).copy({ opacity: 0.14 + 0.42 * (1 - hit.idx / Math.max(1, (lensMeta?.k ?? 8) - 1)) }).formatRgb()
|
||||
: f.phase === 'reading' ? 'rgba(232,161,60,.06)' : 'rgba(108,184,224,.06)'}
|
||||
stroke={pinned[c.tid] ? pinColor(c.tid) : 'transparent'}
|
||||
/>
|
||||
<text x="3" y={CELL_H - 7} className="lv-word" opacity={opacity}>
|
||||
{fmtTok(c.str).slice(0, 7)}
|
||||
{maskOn && c.rank > 0 && <tspan className="lv-rank" dy="-4">{c.rank}</tspan>}
|
||||
</text>
|
||||
</g>
|
||||
)
|
||||
})}
|
||||
</g>
|
||||
))}
|
||||
</>
|
||||
)
|
||||
}, [tick, frames, maskOn, pinned, genId, view, layersShown, gridW])
|
||||
|
||||
// "Activations" view: L2 norm of the residual per layer/position, normalized
|
||||
// PER LAYER (the norm grows strongly with depth: without this, the last layers
|
||||
// would crush everything) between the min and the 95th PERCENTILE of the
|
||||
// layer: the first token is an "attention sink" with a norm ~100× larger than
|
||||
// the rest — without this clamp it would be the only thing visible and all the
|
||||
// rest uniformly dark. Light color = high norm for the layer; outliers saturate
|
||||
// to yellow (exact value in the tooltip).
|
||||
const actHasData = useMemo(
|
||||
() => view === 'act' && shown.some((f) => Object.values(f.layers).some((d) => d?.h_norm != null)),
|
||||
[view, tick, frames],
|
||||
)
|
||||
const actCells = useMemo(() => {
|
||||
if (view !== 'act') return null
|
||||
const extent = {}
|
||||
layersShown.forEach((layer) => {
|
||||
const vals = []
|
||||
shown.forEach((f) => {
|
||||
const v = f.layers[String(layer)]?.h_norm
|
||||
if (v != null) vals.push(v)
|
||||
})
|
||||
vals.sort((a, b) => a - b)
|
||||
const lo = vals[0] ?? Infinity
|
||||
const hi = vals.length ? vals[Math.floor(0.95 * (vals.length - 1))] : -Infinity
|
||||
extent[layer] = [lo, hi]
|
||||
})
|
||||
return (
|
||||
<>
|
||||
{shown.map((f, ci) => (
|
||||
<g key={ci} transform={`translate(${ci * CELL_W},0)`}>
|
||||
<rect
|
||||
x="0" y="0" width={CELL_W - 1} height={HEADER_H - 6}
|
||||
fill={f.phase === 'reading' ? 'rgba(232,161,60,.10)' : 'rgba(108,184,224,.10)'}
|
||||
/>
|
||||
<text
|
||||
x={CELL_W / 2} y={HEADER_H - 12}
|
||||
className="lv-toklabel"
|
||||
transform={`rotate(-38 ${CELL_W / 2} ${HEADER_H - 12})`}
|
||||
>
|
||||
{fmtTok(f.tok).slice(0, 9) || '·'}
|
||||
</text>
|
||||
</g>
|
||||
))}
|
||||
{layersShown.map((layer, ri) => (
|
||||
<g key={layer} transform={`translate(0,${HEADER_H + ri * CELL_H})`}>
|
||||
{shown.map((f, ci) => {
|
||||
const v = f.layers[String(layer)]?.h_norm
|
||||
const [lo, hi] = extent[layer]
|
||||
const t = v == null || !isFinite(lo) ? null
|
||||
: hi > lo ? Math.min(1, (v - lo) / (hi - lo)) : 0.5
|
||||
return (
|
||||
<g
|
||||
key={ci}
|
||||
transform={`translate(${ci * CELL_W},0)`}
|
||||
onMouseMove={(e) => setTip({ x: e.clientX, y: e.clientY, layer, ri, ci, f, act: { v, lo, hi } })}
|
||||
>
|
||||
<rect x="0" y="0" width={CELL_W} height={CELL_H} fill="transparent" />
|
||||
<rect
|
||||
x="0" y="0" width={CELL_W - 1} height={CELL_H - 1}
|
||||
fill={t == null ? 'rgba(255,255,255,.04)' : d3.interpolateInferno(0.08 + 0.84 * t)}
|
||||
/>
|
||||
</g>
|
||||
)
|
||||
})}
|
||||
</g>
|
||||
))}
|
||||
</>
|
||||
)
|
||||
}, [tick, frames, view, layersShown, gridW])
|
||||
|
||||
// No frames (e.g. while regenerating the very first reply): render an empty
|
||||
// shell instead of unmounting — unmounting would wipe the pinned tokens.
|
||||
if (!shown.length) return <div className="lensview" />
|
||||
|
||||
return (
|
||||
<div className="lensview" style={maxH ? { maxHeight: maxH, height: maxH } : undefined}>
|
||||
<div className="lv-controls">
|
||||
{/* ['act', 'Activations'] disabled (TODO) — the view==='act' render below
|
||||
stays in place, just re-add the entry here to re-enable it */}
|
||||
{[['agg', 'Frequencies'], ['grid', 'Heatmap']].map(([v, label]) => (
|
||||
<button key={v} className={view === v ? 'lv-view-on' : ''}
|
||||
onClick={() => { setView(v); setTip(null); setAggTip(null) }}>{label}</button>
|
||||
))}
|
||||
{lensMeta && (
|
||||
<span className="lv-lensbadge" title={lensMeta.path || `${lensMeta.repo_id || ''} ${lensMeta.filename || ''}`.trim()}>
|
||||
🔎 {lensLabel(lensMeta)}
|
||||
</span>
|
||||
)}
|
||||
<span className="lv-sep" />
|
||||
<label>
|
||||
<input type="checkbox" checked={maskOn} onChange={(e) => setMaskOn(e.target.checked)} />
|
||||
BPE/punctuation mask
|
||||
</label>
|
||||
{view === 'grid' && (
|
||||
<label title="follow the right edge (latest tokens) during generation">
|
||||
<input type="checkbox" checked={autoScroll} onChange={(e) => setAutoScroll(e.target.checked)} />
|
||||
auto-scroll →
|
||||
</label>
|
||||
)}
|
||||
<span className="lv-sep" />
|
||||
<label title="local display filter — doesn't affect the server capture">display</label>
|
||||
<input
|
||||
type="text"
|
||||
placeholder={`all (e.g. ${layers[0]}-${layers[layers.length - 1]})`}
|
||||
value={filterInput}
|
||||
onChange={(e) => setFilterInput(e.target.value)}
|
||||
style={{ width: 120 }}
|
||||
/>
|
||||
{filterSet != null && (
|
||||
<span className={`lv-filter-badge ${filterEmpty ? 'err' : ''}`}>
|
||||
{filterEmpty ? 'no layer matches' : `${layersShown.length}/${layers.length} layers`}
|
||||
<span className="lv-filter-clear" onClick={() => setFilterInput('')}> ✕</span>
|
||||
</span>
|
||||
)}
|
||||
<label title="pin a token by its text — pick the variant with (˽) or without a space">pin</label>
|
||||
<span className="lv-pinadd">
|
||||
<input type="text" placeholder="token…" value={pinInput}
|
||||
onChange={(e) => { setPinInput(e.target.value); lookupPin(e.target.value) }}
|
||||
onKeyDown={(e) => {
|
||||
if (e.key === 'Enter' && pinCands.length) {
|
||||
addPin(pinCands.find((c) => c.str.startsWith(' ')) || pinCands[0])
|
||||
}
|
||||
}}
|
||||
style={{ width: 90 }} />
|
||||
{pinCands.length > 0 && (
|
||||
<span className="lv-pinadd-cands">
|
||||
{pinCands.slice(0, 6).map((c) => (
|
||||
<button key={c.id} title={`id ${c.id}`} onClick={() => addPin(c)}>{fmtTok(c.str)}</button>
|
||||
))}
|
||||
</span>
|
||||
)}
|
||||
</span>
|
||||
<span className="lv-sep" />
|
||||
{pinnedIds.map((tid) => (
|
||||
<span key={tid} className="lv-pin" style={{ borderColor: pinColor(tid), color: pinColor(tid) }}
|
||||
title={transLabel(tid) ? `≈ ${transLabel(tid)}` : undefined}
|
||||
onMouseEnter={() => fetchTrans([{ tid, str: String(pinned[tid]) }])}>
|
||||
<span onClick={() => togglePin(tid, pinned[tid])}>{fmtTok(pinned[tid])} ✕</span>
|
||||
<span
|
||||
className="lv-ablate"
|
||||
title="open the editor with this token prefilled (most-relevant layer pre-selected)"
|
||||
onClick={() => onEditToken?.(+tid, String(pinned[tid]), peakLayerOf(+tid))}
|
||||
> ☢</span>
|
||||
</span>
|
||||
))}
|
||||
<span className="fcount" style={{ marginLeft: 'auto' }}>
|
||||
{frames.length} tokens · {layersShown.length}{layersShown.length !== layers.length ? `/${layers.length}` : ''} layers
|
||||
{frames.length > MAX_COLS ? ` · last ${MAX_COLS} shown` : ''}
|
||||
</span>
|
||||
</div>
|
||||
|
||||
{/* visual selection of the displayed layers (click/click-drag, number on hover) —
|
||||
synced with the "display" text field above via filterInput */}
|
||||
<div className="lv-layerrow">
|
||||
<span className="src">layers shown:</span>
|
||||
<LayerPicker all={layers} value={dragLayers ?? layersShown} compact
|
||||
onChange={setDragLayers} />
|
||||
</div>
|
||||
|
||||
{view === 'agg' && (
|
||||
<div className="lv-agg-wrap">
|
||||
<div className="lv-agg-bar">
|
||||
<label>show{' '}
|
||||
{/* min/default multiples of the step: an unaligned value makes the
|
||||
spinner's first click "snap" by 1 instead of stepping */}
|
||||
<input type="number" min="10" max="2000" step="10" value={aggLimit}
|
||||
onChange={(e) => setAggLimit(Math.max(10, Math.trunc(+e.target.value) || 10))}
|
||||
style={{ width: 62 }} /> tokens
|
||||
</label>
|
||||
<span className="lv-sep" />
|
||||
<span className="src">click = pin · right-click = hide</span>
|
||||
</div>
|
||||
<div className="lv-agg">
|
||||
{agg.map((e) => (
|
||||
<span
|
||||
key={e.tid}
|
||||
className={`lv-agg-word ${pinned[e.tid] ? 'lv-agg-pinned' : ''}`}
|
||||
style={{
|
||||
opacity: Math.max(0.35, Math.min(1, Math.sqrt(e.maxP) * 2.4)),
|
||||
fontSize: `${Math.min(19, 11 + Math.log2(e.appearances) * 1.6)}px`,
|
||||
...(pinned[e.tid] ? { color: pinColor(e.tid) } : {}),
|
||||
}}
|
||||
onMouseEnter={(ev) => {
|
||||
setAggTip({ x: ev.clientX, y: ev.clientY, e })
|
||||
fetchTrans([{ tid: e.tid, str: e.str }])
|
||||
}}
|
||||
onMouseLeave={() => setAggTip(null)}
|
||||
onClick={() => togglePin(e.tid, e.str)}
|
||||
onContextMenu={(ev) => { ev.preventDefault(); onHideToken(e.str); setAggTip(null) }}
|
||||
>
|
||||
{fmtTok(e.str)}<sup>{e.appearances}</sup>
|
||||
</span>
|
||||
))}
|
||||
{agg.length === 0 && <span className="src">no token to show (all hidden or no frames)</span>}
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{aggTip && (() => {
|
||||
const e = aggTip.e
|
||||
const trans = NONLATIN_RE.test(e.str || '') ? transCache.current.get(e.tid) : null
|
||||
return (
|
||||
<Tip x={aggTip.x} y={aggTip.y}>
|
||||
<div className="lv-tip-head">"{fmtTok(e.str)}" — click = pin</div>
|
||||
<div className="lv-tip-row"><span>appearances</span><span>{e.appearances}</span></div>
|
||||
<div className="lv-tip-row"><span>peak</span><span>L{e.peakLayer}</span></div>
|
||||
<div className="lv-tip-row"><span>max p</span><span>{(e.maxP * 100).toFixed(2)} %</span></div>
|
||||
{trans?.length > 0 && (
|
||||
<>
|
||||
<div className="lv-tip-head" style={{ marginTop: 6 }}>≈ nearest tokens (W_U cosine)</div>
|
||||
{trans.map((n) => (
|
||||
<div key={n.id} className="lv-tip-row">
|
||||
<span className="lv-trans">{fmtTok(n.str)}</span>
|
||||
<span>{n.sim.toFixed(2)}</span>
|
||||
</div>
|
||||
))}
|
||||
</>
|
||||
)}
|
||||
</Tip>
|
||||
)
|
||||
})()}
|
||||
|
||||
{view === 'grid' && (
|
||||
<div className="lv-gridwrap" ref={gridScrollRef}>
|
||||
<svg width={LABEL_W} height={height} className="lv-svg lv-labels">
|
||||
{layersShown.map((layer, ri) => (
|
||||
<g key={layer} transform={`translate(0,${HEADER_H + ri * CELL_H})`}>
|
||||
{ri % 2 === 1 && <rect x="0" y="0" width={LABEL_W} height={CELL_H} fill="rgba(255,255,255,.025)" />}
|
||||
<text
|
||||
x={LABEL_W - 6} y={CELL_H / 2 + 4}
|
||||
className={`lv-laylabel ${tip?.layer === layer ? 'lv-laylabel-hot' : ''}`}
|
||||
>L{layer}</text>
|
||||
</g>
|
||||
))}
|
||||
</svg>
|
||||
<div className="lv-scroll" onMouseLeave={() => setTip(null)}>
|
||||
<svg width={gridW} height={height} className="lv-svg">
|
||||
{gridCells}
|
||||
{tip && tip.ri != null && (
|
||||
<g pointerEvents="none">
|
||||
<rect x="0" y={HEADER_H + tip.ri * CELL_H} width={gridW} height={CELL_H - 1}
|
||||
fill="rgba(108,184,224,.08)" stroke="rgba(108,184,224,.35)" />
|
||||
<rect x={tip.ci * CELL_W} y="0" width={CELL_W - 1} height={height}
|
||||
fill="rgba(108,184,224,.06)" />
|
||||
</g>
|
||||
)}
|
||||
</svg>
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{view === 'act' && (
|
||||
<div className="lv-gridwrap">
|
||||
<svg width={LABEL_W} height={height} className="lv-svg lv-labels">
|
||||
{layersShown.map((layer, ri) => (
|
||||
<g key={layer} transform={`translate(0,${HEADER_H + ri * CELL_H})`}>
|
||||
{ri % 2 === 1 && <rect x="0" y="0" width={LABEL_W} height={CELL_H} fill="rgba(255,255,255,.025)" />}
|
||||
<text
|
||||
x={LABEL_W - 6} y={CELL_H / 2 + 4}
|
||||
className={`lv-laylabel ${tip?.layer === layer ? 'lv-laylabel-hot' : ''}`}
|
||||
>L{layer}</text>
|
||||
</g>
|
||||
))}
|
||||
</svg>
|
||||
<div className="lv-scroll" onMouseLeave={() => setTip(null)}>
|
||||
{!actHasData && (
|
||||
<div className="status-line" style={{ padding: '6px 8px' }}>
|
||||
activation norms are absent from these frames — regenerate (new
|
||||
generations capture them)
|
||||
</div>
|
||||
)}
|
||||
<svg width={gridW} height={height} className="lv-svg">
|
||||
{actCells}
|
||||
{tip && tip.ri != null && (
|
||||
<g pointerEvents="none">
|
||||
<rect x="0" y={HEADER_H + tip.ri * CELL_H} width={gridW} height={CELL_H - 1}
|
||||
fill="none" stroke="rgba(108,184,224,.45)" />
|
||||
<rect x={tip.ci * CELL_W} y="0" width={CELL_W - 1} height={height}
|
||||
fill="rgba(108,184,224,.06)" />
|
||||
</g>
|
||||
)}
|
||||
</svg>
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{tip && tip.act && (
|
||||
<Tip x={tip.x} y={tip.y}>
|
||||
<div className="lv-tip-head">
|
||||
pos {tip.f.pos} · L{tip.layer} · {tip.f.phase === 'reading' ? 'reading' : 'thinking'} · tok «{fmtTok(tip.f.tok)}»
|
||||
</div>
|
||||
{tip.act.v == null ? (
|
||||
<div className="lv-tip-row"><span>‖h‖ unavailable (older generation)</span></div>
|
||||
) : (
|
||||
<>
|
||||
<div className="lv-tip-row"><span>‖h‖ (residual norm)</span><span>{tip.act.v}</span></div>
|
||||
<div className="lv-tip-row"><span>layer range</span><span>{tip.act.lo} – {tip.act.hi}</span></div>
|
||||
</>
|
||||
)}
|
||||
</Tip>
|
||||
)}
|
||||
|
||||
{tip && !tip.act && (() => {
|
||||
// non-latin tokens of the cell whose "nearest tokens" we show at the tooltip's end
|
||||
const nonLatin = tip.c.top.filter((t) => NONLATIN_RE.test(t.s || ''))
|
||||
return (
|
||||
<Tip x={tip.x} y={tip.y}>
|
||||
<div className="lv-tip-head">
|
||||
pos {tip.f.pos} · L{tip.layer} · {tip.f.phase === 'reading' ? 'reading' : 'thinking'} · tok «{fmtTok(tip.f.tok)}»
|
||||
</div>
|
||||
{tip.c.top.map((t, i) => (
|
||||
<div key={i} className="lv-tip-row" style={pinned[t.tid] ? { color: pinColor(t.tid) } : undefined}>
|
||||
<span>{fmtTok(t.s)}</span>
|
||||
<span>r{t.r} · {(t.p * 100).toFixed(2)}%</span>
|
||||
</div>
|
||||
))}
|
||||
{nonLatin.map((t) => {
|
||||
const list = transCache.current.get(t.tid)
|
||||
if (!list?.length) return null
|
||||
return (
|
||||
<div key={`tr${t.tid}`} className="lv-tip-trans">
|
||||
<div className="lv-tip-head">"{fmtTok(t.s)}" ≈ nearest tokens (W_U cosine)</div>
|
||||
{list.slice(0, 6).map((n) => (
|
||||
<div key={n.id} className="lv-tip-row">
|
||||
<span className="lv-trans">{fmtTok(n.str)}</span>
|
||||
<span>{n.sim.toFixed(2)}</span>
|
||||
</div>
|
||||
))}
|
||||
</div>
|
||||
)
|
||||
})}
|
||||
</Tip>
|
||||
)
|
||||
})()}
|
||||
|
||||
{pinData && (
|
||||
<div className="lv-pins" ref={pinsRef}>
|
||||
{/* only tokens still pinned: a stale pinData must never resurrect an
|
||||
unpinned token's graph */}
|
||||
{Object.entries(pinData.pins).filter(([tid]) => pinned[tid]).map(([tid, data]) => {
|
||||
const N = pinData.positions.length
|
||||
const L = pinData.layers.length
|
||||
// width: fills the container (PIN_FILL_WIDTH), at least
|
||||
// PIN_PX_PER_TOKEN per position, bounded by PIN_MIN_W / PIN_MAX_W
|
||||
const fillW = PIN_FILL_WIDTH && pinsW ? pinsW - 46 - 10 : 0
|
||||
const w = Math.min(PIN_MAX_W, Math.max(PIN_MIN_W, N * PIN_PX_PER_TOKEN, fillW))
|
||||
const curveH = PIN_CURVE_H
|
||||
const rowH = Math.max(PIN_ROW_H_MIN, Math.min(PIN_ROW_H_MAX, Math.floor(PIN_MAP_H / L)))
|
||||
const mapH = L * rowH
|
||||
// column bands: position pi occupies [xL(pi), xL(pi)+cw] — the
|
||||
// curve passes through the centers. A [0, N-1] → [0, w] point
|
||||
// scale would push the LAST heatmap column past the svg edge
|
||||
// (clipped) and misalign the hover column math.
|
||||
const cw = w / N
|
||||
const xL = (pi) => 40 + pi * cw
|
||||
const xC = (pi) => 40 + (pi + 0.5) * cw
|
||||
const y = d3.scaleLinear([0, 5.2], [2, curveH - 2])
|
||||
const layerColor = (li) => d3.interpolateCool(0.15 + 0.7 * (li / Math.max(1, L - 1)))
|
||||
const hov = pinHover?.tid === tid ? pinHover.li : null
|
||||
const hovPi = pinHover?.tid === tid ? pinHover.pi : null
|
||||
// layer labels: all if few, otherwise sampled
|
||||
const labelEvery = L <= 16 ? 1 : Math.ceil(L / 12)
|
||||
return (
|
||||
<div key={tid} className="lv-pinblock">
|
||||
<div className="lv-pinname" style={{ color: pinColor(+tid) }}>
|
||||
"{fmtTok(pinned[tid] ?? tid)}"
|
||||
<button className="lv-copybtn" title="copy the token to the clipboard (without spaces)"
|
||||
onClick={() => copyText(String(pinned[tid] ?? ''))}>⧉</button>
|
||||
<span className="lv-pinsub"> — token rank per layer and position</span>
|
||||
{hov != null && (
|
||||
<span className="lv-pinhov">
|
||||
{' '}· hover: L{pinData.layers[hov]}
|
||||
{hovPi != null && pinData.tokens?.[hovPi] != null
|
||||
? ` — ${fmtTok(pinData.tokens[hovPi])}` : ''}
|
||||
</span>
|
||||
)}
|
||||
</div>
|
||||
{/* legend: one chip per layer, hover = highlights the curve and the row */}
|
||||
<div className="lv-laylegend">
|
||||
{pinData.layers.map((layer, li) => (
|
||||
<span
|
||||
key={layer}
|
||||
className={`lv-laychip ${hov != null && hov !== li ? 'dim' : ''}`}
|
||||
style={{ borderColor: layerColor(li), color: layerColor(li) }}
|
||||
onMouseEnter={() => setPinHover({ tid, li })}
|
||||
onMouseLeave={() => setPinHover(null)}
|
||||
>L{layer}</span>
|
||||
))}
|
||||
</div>
|
||||
<svg width={w + 46} height={curveH + 10}>
|
||||
{[0, 1, 2, 3, 4, 5].map((d) => (
|
||||
<g key={d}>
|
||||
<line x1="40" x2={w + 40} y1={y(d)} y2={y(d)} className="lv-grid" />
|
||||
<text x="36" y={y(d) + 3} className="lv-axis">{d === 0 ? '1' : `1e${d}`}</text>
|
||||
</g>
|
||||
))}
|
||||
{data.ranks.map((layerRanks, li) => (
|
||||
<polyline
|
||||
key={li}
|
||||
fill="none"
|
||||
stroke={layerColor(li)}
|
||||
strokeWidth={hov === li ? 2.4 : 1.2}
|
||||
opacity={hov == null ? 0.75 : hov === li ? 1 : 0.12}
|
||||
points={layerRanks.map((r, pi) => `${xC(pi)},${y(Math.log10(r + 1))}`).join(' ')}
|
||||
onMouseEnter={() => setPinHover({ tid, li })}
|
||||
onMouseLeave={() => setPinHover(null)}
|
||||
style={{ cursor: 'pointer' }}
|
||||
/>
|
||||
))}
|
||||
</svg>
|
||||
<svg width={w + 46} height={mapH + 16}
|
||||
onMouseMove={(e) => {
|
||||
// continuous hover: each Y pixel of the heatmap falls on a row
|
||||
// (layer), with no gap between rows → no more jumping
|
||||
const rect = e.currentTarget.getBoundingClientRect()
|
||||
const ry = e.clientY - rect.top
|
||||
const li = Math.floor(ry / rowH)
|
||||
// column = token position (X axis), to show the hovered token
|
||||
const pi = Math.max(0, Math.min(N - 1,
|
||||
Math.floor((e.clientX - rect.left - 40) / cw)))
|
||||
if (li >= 0 && li < L) setPinHover({ tid, li, pi })
|
||||
}}
|
||||
onMouseLeave={() => setPinHover(null)}
|
||||
>
|
||||
{data.ranks.map((layerRanks, li) => (
|
||||
<g key={li}>
|
||||
{(li % labelEvery === 0 || hov === li) && (
|
||||
<text x="36" y={li * rowH + rowH - 1}
|
||||
className={`lv-axis ${hov === li ? 'lv-axis-hot' : ''}`}>L{pinData.layers[li]}</text>
|
||||
)}
|
||||
{layerRanks.map((r, pi) => (
|
||||
<rect
|
||||
key={pi}
|
||||
x={xL(pi)} y={li * rowH}
|
||||
width={Math.max(1.5, cw)} height={rowH - 1}
|
||||
fill={rankColor(r)}
|
||||
opacity={hov == null || hov === li ? 1 : 0.25}
|
||||
/>
|
||||
))}
|
||||
{hov === li && (
|
||||
<rect x="40" y={li * rowH} width={w} height={rowH - 1}
|
||||
fill="none" stroke="var(--accent2)" strokeWidth="1" pointerEvents="none" />
|
||||
)}
|
||||
</g>
|
||||
))}
|
||||
<text x="40" y={mapH + 12} className="lv-axis" style={{ textAnchor: 'start' }}>
|
||||
color = rank (light = rank 1) · X axis = token position · hover a layer above
|
||||
</text>
|
||||
</svg>
|
||||
</div>
|
||||
)
|
||||
})}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
)
|
||||
}
|
||||
@@ -0,0 +1,6 @@
|
||||
import React from 'react'
|
||||
import { createRoot } from 'react-dom/client'
|
||||
import App from './App.jsx'
|
||||
import './styles.css'
|
||||
|
||||
createRoot(document.getElementById('root')).render(<App />)
|
||||
@@ -0,0 +1,455 @@
|
||||
* { box-sizing: border-box; margin: 0; padding: 0; }
|
||||
|
||||
:root {
|
||||
--bg: #14161a;
|
||||
--panel: #1c1f26;
|
||||
--panel2: #23272f;
|
||||
--border: #333842;
|
||||
--text: #d8dce3;
|
||||
--muted: #8b93a1;
|
||||
--accent: #e8a13c;
|
||||
--accent2: #6cb8e0;
|
||||
--danger: #d0654f;
|
||||
}
|
||||
|
||||
body { background: var(--bg); color: var(--text); font: 14px/1.45 'Segoe UI', system-ui, sans-serif; }
|
||||
|
||||
#root { display: flex; height: 100vh; }
|
||||
|
||||
.sidebar {
|
||||
width: 300px; flex-shrink: 0; background: var(--panel);
|
||||
border-right: 1px solid var(--border); padding: 14px;
|
||||
display: flex; flex-direction: column; gap: 12px; overflow-y: auto;
|
||||
}
|
||||
.sidebar h1 { font-size: 17px; color: var(--accent); }
|
||||
.sidebar h2 { font-size: 12px; text-transform: uppercase; letter-spacing: .08em; color: var(--muted); margin-top: 4px; }
|
||||
|
||||
.model-list { display: flex; flex-direction: column; gap: 4px; }
|
||||
.model-item {
|
||||
padding: 7px 9px; border: 1px solid var(--border); border-radius: 6px;
|
||||
cursor: pointer; font-size: 13px; word-break: break-all;
|
||||
}
|
||||
.model-item:hover { background: var(--panel2); }
|
||||
.model-item.selected { border-color: var(--accent); background: var(--panel2); }
|
||||
.model-item .src { color: var(--muted); font-size: 11px; }
|
||||
.model-item.loaded { border-color: var(--accent2); }
|
||||
.model-item-head { display: flex; justify-content: space-between; align-items: flex-start; gap: 6px; }
|
||||
.model-del {
|
||||
padding: 0 5px; font-size: 12px; border: none; background: transparent;
|
||||
flex-shrink: 0; opacity: .55; display: inline-flex; align-items: center;
|
||||
}
|
||||
.model-item:hover .model-del { opacity: 1; }
|
||||
.model-del:disabled { opacity: .2; }
|
||||
/* red = deletes the FILES from disk; blue = only forgets a registered entry */
|
||||
.model-trash { color: var(--danger); }
|
||||
.model-trash:hover:not(:disabled) { color: #ff7a6b; border-color: transparent; }
|
||||
.model-unreg { color: var(--accent2); }
|
||||
.model-unreg:hover:not(:disabled) { color: #9fd4f2; border-color: transparent; }
|
||||
|
||||
.row { display: flex; gap: 6px; align-items: center; }
|
||||
.row label { color: var(--muted); font-size: 12px; width: 58px; flex-shrink: 0; }
|
||||
select, input[type="text"], input[type="number"], textarea {
|
||||
background: var(--panel2); color: var(--text); border: 1px solid var(--border);
|
||||
border-radius: 5px; padding: 5px 7px; font: inherit; width: 100%;
|
||||
}
|
||||
textarea { resize: vertical; }
|
||||
|
||||
button {
|
||||
background: var(--panel2); color: var(--text); border: 1px solid var(--border);
|
||||
border-radius: 6px; padding: 6px 12px; cursor: pointer; font: inherit;
|
||||
}
|
||||
button:hover:not(:disabled) { border-color: var(--accent); }
|
||||
button:disabled { opacity: .45; cursor: default; }
|
||||
button.primary { background: var(--accent); color: #14161a; border-color: var(--accent); font-weight: 600; }
|
||||
button.danger { border-color: var(--danger); color: var(--danger); }
|
||||
|
||||
.gpu { margin-bottom: 8px; }
|
||||
.gpu .name { font-size: 12px; color: var(--muted); display: flex; justify-content: space-between; }
|
||||
.gpu .bar { height: 8px; background: var(--panel2); border-radius: 4px; overflow: hidden; margin-top: 3px; }
|
||||
.gpu .fill { height: 100%; background: var(--accent2); transition: width .5s; }
|
||||
.dl-bar { height: 5px; background: var(--panel2); border-radius: 3px; overflow: hidden; margin-top: 3px; }
|
||||
.dl-fill { height: 100%; background: var(--accent); transition: width .4s; }
|
||||
|
||||
/* flex-shrink 0: as a flex item of the (full) sidebar it used to get squeezed
|
||||
to min-height, and the download progress bar overflowed onto "Browse" */
|
||||
.status-line { font-size: 12px; color: var(--muted); min-height: 16px; word-break: break-word; flex-shrink: 0; }
|
||||
.status-line .ok { color: var(--accent2); }
|
||||
.status-line .err { color: var(--danger); }
|
||||
|
||||
.main { flex: 1; display: flex; flex-direction: column; min-width: 0; }
|
||||
|
||||
.messages { flex: 1; overflow-y: auto; padding: 18px 22px; display: flex; flex-direction: column; gap: 12px; }
|
||||
.msg { max-width: 780px; padding: 10px 14px; border-radius: 10px; white-space: pre-wrap; word-break: break-word; }
|
||||
.msg.user { background: var(--panel2); align-self: flex-end; }
|
||||
.msg.assistant { background: var(--panel); border: 1px solid var(--border); align-self: flex-start; }
|
||||
.msg .msgmeta { font-size: 11px; color: var(--muted); margin-top: 6px; }
|
||||
|
||||
/* markdown rendering inside assistant bubbles (the bubble keeps pre-wrap for
|
||||
plain text; .md content manages its own whitespace) */
|
||||
.msg .md { white-space: normal; }
|
||||
.msg .md > :first-child { margin-top: 0; }
|
||||
.msg .md > :last-child { margin-bottom: 0; }
|
||||
.msg .md p, .msg .md ul, .msg .md ol, .msg .md pre, .msg .md blockquote,
|
||||
.msg .md h1, .msg .md h2, .msg .md h3, .msg .md h4, .msg .md table { margin: 6px 0; }
|
||||
.msg .md h1, .msg .md h2 { font-size: 15px; }
|
||||
.msg .md h3, .msg .md h4 { font-size: 13.5px; }
|
||||
.msg .md ul, .msg .md ol { padding-left: 20px; }
|
||||
.msg .md code {
|
||||
background: var(--panel2); padding: 1px 4px; border-radius: 4px;
|
||||
font-size: 12px; font-family: Consolas, 'Fira Code', monospace;
|
||||
}
|
||||
.msg .md pre {
|
||||
background: var(--panel2); border: 1px solid var(--border); border-radius: 6px;
|
||||
padding: 8px 10px; overflow-x: auto;
|
||||
}
|
||||
.msg .md pre code { background: none; padding: 0; }
|
||||
.msg .md table { border-collapse: collapse; }
|
||||
.msg .md th, .msg .md td { border: 1px solid var(--border); padding: 3px 8px; font-size: 12.5px; }
|
||||
.msg .md blockquote { border-left: 3px solid var(--border); padding-left: 10px; color: var(--muted); }
|
||||
.msg .md a { color: var(--accent2); }
|
||||
|
||||
/* inline edit of an assistant reply */
|
||||
.msg-edit textarea { width: 100%; min-width: 420px; resize: vertical; }
|
||||
.msg-edit-actions { display: flex; gap: 8px; margin-top: 6px; }
|
||||
|
||||
.composer { border-top: 1px solid var(--border); padding: 12px 22px; display: flex; flex-direction: column; gap: 8px; }
|
||||
.composer .controls { display: flex; gap: 10px; flex-wrap: wrap; align-items: center; }
|
||||
.composer .controls label { font-size: 12px; color: var(--muted); }
|
||||
.composer .controls input { width: 64px; }
|
||||
.composer .inputrow { display: flex; gap: 8px; }
|
||||
.composer .inputrow textarea { flex: 1; min-height: 44px; max-height: 160px; }
|
||||
|
||||
.tabs { display: flex; gap: 4px; }
|
||||
.tabs button { flex: 1; padding: 5px 0; font-size: 12px; border-radius: 6px 6px 0 0; border-bottom: 2px solid transparent; }
|
||||
.tabs button.tab-on { border-bottom-color: var(--accent); color: var(--accent); font-weight: 600; }
|
||||
|
||||
.exp-tag {
|
||||
font-size: 9px; color: var(--danger); border: 1px solid var(--danger);
|
||||
border-radius: 6px; padding: 0 5px; vertical-align: middle; letter-spacing: .05em;
|
||||
}
|
||||
/* export greyed out when the "per-layer steering" mode doesn't produce a faithful checkpoint */
|
||||
.exp-grid { display: flex; flex-direction: column; gap: 7px; transition: opacity .18s; }
|
||||
.exp-grid-off { opacity: .4; pointer-events: none; }
|
||||
.exp-disabled-note { color: var(--danger); }
|
||||
|
||||
.diff-btn { padding: 0 6px; font-size: 11px; margin-left: 6px; }
|
||||
.diff-table { border-collapse: collapse; font-size: 11px; }
|
||||
.diff-table th { color: var(--muted); font-weight: 400; padding: 2px 4px; max-width: 60px; overflow: hidden; }
|
||||
.diff-table td { padding: 1px 4px; border: 1px solid rgba(255,255,255,.04); white-space: nowrap; }
|
||||
.diff-lay { color: var(--muted); }
|
||||
.diff-same { opacity: .45; }
|
||||
.diff-diff { outline: 1px solid var(--accent); }
|
||||
.diff-a { color: var(--accent2); }
|
||||
.diff-b { color: var(--accent); }
|
||||
.lv-ablate { cursor: pointer; }
|
||||
.lv-ablate:hover { color: var(--danger); }
|
||||
|
||||
.reg-list { display: flex; flex-direction: column; gap: 3px; }
|
||||
.reg-item {
|
||||
display: flex; justify-content: space-between; align-items: center; gap: 6px;
|
||||
font-size: 12px; padding: 4px 7px; border: 1px solid var(--border); border-radius: 5px;
|
||||
}
|
||||
.reg-item span { overflow: hidden; text-overflow: ellipsis; white-space: nowrap; }
|
||||
.reg-item button { padding: 2px 8px; font-size: 12px; flex-shrink: 0; }
|
||||
.preset-info { display: flex; flex-direction: column; gap: 1px; min-width: 0; flex: 1; }
|
||||
.preset-name { overflow: hidden; text-overflow: ellipsis; white-space: nowrap; }
|
||||
.preset-model { font-size: 10.5px; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; }
|
||||
|
||||
.reg-item2 {
|
||||
display: flex; justify-content: space-between; align-items: center; gap: 8px;
|
||||
padding: 6px 8px; border: 1px solid var(--border); border-radius: 6px;
|
||||
}
|
||||
.reg-item2.reg-warn { border-color: var(--danger); }
|
||||
.reg-item2 .reg-main { min-width: 0; flex: 1; display: flex; flex-direction: column; gap: 1px; }
|
||||
.reg-item2 button { padding: 3px 9px; font-size: 12px; flex-shrink: 0; align-self: center; }
|
||||
.reg-name { font-size: 12.5px; font-weight: 600; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; }
|
||||
.reg-tag {
|
||||
font-size: 9px; text-transform: uppercase; letter-spacing: .05em; padding: 0 5px;
|
||||
border-radius: 7px; background: var(--panel2); color: var(--muted); font-weight: 600;
|
||||
}
|
||||
.reg-tag.hub { background: rgba(108,184,224,.18); color: var(--accent2); }
|
||||
.reg-tag.base { background: rgba(140,220,160,.16); color: #7fce8f; }
|
||||
.reg-reason { color: var(--danger) !important; white-space: normal; }
|
||||
|
||||
.linkbtn {
|
||||
background: none; border: none; padding: 0; color: var(--accent2);
|
||||
font-size: 11px; cursor: pointer; text-decoration: underline;
|
||||
}
|
||||
.linkbtn:hover { color: var(--accent); }
|
||||
.linkbtn:disabled { color: var(--muted); cursor: default; text-decoration: none; }
|
||||
|
||||
.browser { border: 1px solid var(--border); border-radius: 6px; padding: 6px; }
|
||||
.browser-path { display: flex; align-items: center; gap: 6px; flex-wrap: wrap; margin-bottom: 5px; }
|
||||
.browser-path .src { overflow: hidden; text-overflow: ellipsis; white-space: nowrap; font-size: 11px; }
|
||||
.browser-list { max-height: 260px; overflow-y: auto; display: flex; flex-direction: column; gap: 1px; }
|
||||
.browser-row {
|
||||
display: flex; justify-content: space-between; align-items: center; gap: 6px;
|
||||
padding: 2px 4px; border-radius: 4px; font-size: 12px;
|
||||
}
|
||||
.browser-row:hover { background: var(--panel2); }
|
||||
.browser-row.is-model { color: var(--accent); }
|
||||
.browser-name { cursor: pointer; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; flex: 1; }
|
||||
.browser-row button { padding: 1px 8px; font-size: 11px; flex-shrink: 0; }
|
||||
|
||||
.fit-adv { border-left: 2px solid var(--border); padding-left: 8px; display: flex; flex-direction: column; gap: 6px; }
|
||||
|
||||
.conv-list { display: flex; flex-direction: column; gap: 4px; max-height: 220px; overflow-y: auto; }
|
||||
.conv-item { padding: 6px 8px; border: 1px solid var(--border); border-radius: 6px; cursor: pointer; font-size: 12px; }
|
||||
.conv-item:hover { background: var(--panel2); }
|
||||
.conv-item.selected { border-color: var(--accent2); }
|
||||
.conv-title { display: flex; justify-content: space-between; align-items: center; gap: 6px; }
|
||||
.conv-del { padding: 0 5px; font-size: 10px; border: none; color: var(--muted); }
|
||||
.conv-del:hover { color: var(--danger); }
|
||||
.conv-snippet { color: var(--muted); font-size: 11px; margin-top: 2px; }
|
||||
|
||||
.conv-header {
|
||||
display: flex; gap: 8px; align-items: center; padding: 8px 22px;
|
||||
border-bottom: 1px solid var(--border); background: var(--panel);
|
||||
}
|
||||
.conv-title-input { flex: 1; font-weight: 600; }
|
||||
.conv-tags-input { width: 220px; }
|
||||
.conv-header a { color: var(--accent2); font-size: 12px; text-decoration: none; border: 1px solid var(--border); border-radius: 5px; padding: 4px 9px; }
|
||||
.conv-header a:hover { border-color: var(--accent2); }
|
||||
|
||||
.tree-panel {
|
||||
border-bottom: 1px solid var(--border); background: var(--panel);
|
||||
padding: 8px 14px; max-height: 220px; overflow-y: auto; font-size: 12px;
|
||||
}
|
||||
.tree-node { display: flex; gap: 7px; align-items: baseline; padding: 2px 4px; cursor: pointer; border-radius: 4px; white-space: nowrap; overflow: hidden; }
|
||||
.tree-node:hover { background: var(--panel2); }
|
||||
.tree-node.on-path { background: rgba(108,184,224,.10); }
|
||||
.tn-role { width: 14px; text-align: center; border-radius: 3px; font-size: 10px; font-weight: 700; flex-shrink: 0; }
|
||||
.tn-role.user { color: var(--accent); }
|
||||
.tn-role.assistant { color: var(--accent2); }
|
||||
.tn-role.system { color: var(--muted); }
|
||||
.tn-text { overflow: hidden; text-overflow: ellipsis; color: var(--text); }
|
||||
.tn-frames { color: var(--accent); }
|
||||
.tree-hint { color: var(--muted); font-size: 11px; margin-top: 5px; }
|
||||
.fork-note { color: var(--accent); font-size: 12px; }
|
||||
|
||||
.msg.has-frames { cursor: pointer; }
|
||||
.msg.msg-selected { outline: 1px solid var(--accent2); }
|
||||
|
||||
.lensview {
|
||||
border-top: 1px solid var(--border); background: var(--panel);
|
||||
max-height: 46vh; overflow-y: auto; padding: 6px 12px; flex-shrink: 0;
|
||||
}
|
||||
/* drag handle above the lens view: sets its height manually */
|
||||
.v-resizer {
|
||||
height: 6px; cursor: ns-resize; flex-shrink: 0;
|
||||
background: linear-gradient(var(--border), transparent);
|
||||
}
|
||||
.v-resizer:hover { background: var(--accent2); opacity: .5; }
|
||||
.lv-controls { display: flex; align-items: center; gap: 10px; font-size: 12px; color: var(--muted); padding: 4px 0 8px; flex-wrap: wrap; }
|
||||
.lv-controls input[type="text"] { padding: 3px 6px; font-size: 12px; }
|
||||
.lv-controls button { padding: 3px 10px; font-size: 12px; }
|
||||
.lv-sep { width: 1px; height: 16px; background: var(--border); }
|
||||
.lv-pin { background: rgba(232,161,60,.15); color: var(--accent); border: 1px solid var(--accent); border-radius: 9px; padding: 1px 8px; cursor: pointer; font-size: 12px; }
|
||||
/* single scroll container (H + V): the bars stay at the edges so they're always
|
||||
visible, even when the grid overflows in height. The label column is pinned to
|
||||
the left (sticky) and scrolls vertically with the grid. */
|
||||
.lv-gridwrap { display: flex; align-items: flex-start; overflow: auto; max-height: 34vh; }
|
||||
.lv-labels { position: sticky; left: 0; z-index: 2; background: var(--panel); flex-shrink: 0; }
|
||||
.lv-scroll { flex-shrink: 0; }
|
||||
.lv-laylabel-hot { fill: var(--accent2); font-weight: 700; }
|
||||
.lv-filter-badge { color: var(--accent2); font-size: 11px; white-space: nowrap; }
|
||||
.lv-filter-badge.err { color: var(--danger); }
|
||||
.lv-filter-clear { cursor: pointer; }
|
||||
.lv-filter-clear:hover { color: var(--danger); }
|
||||
.lv-agg { display: flex; flex-wrap: wrap; gap: 8px 12px; align-items: baseline; padding: 6px 2px; }
|
||||
.lv-agg-word { cursor: pointer; }
|
||||
.lv-agg-word:hover { color: var(--accent2); }
|
||||
.lv-agg-word.lv-agg-pinned { color: var(--accent); text-decoration: underline; }
|
||||
.lv-agg-word sup { font-size: 9px; color: var(--muted); }
|
||||
.lv-agg-bar { display: flex; align-items: center; gap: 10px; flex-wrap: wrap; font-size: 12px; color: var(--muted); padding: 2px 2px 6px; }
|
||||
.lv-agg-bar input[type="number"], .lv-agg-bar input[type="text"] { padding: 3px 6px; font-size: 12px; }
|
||||
.lv-hidden-row { display: flex; align-items: center; gap: 6px; flex-wrap: wrap; padding: 0 2px 8px; }
|
||||
.lv-hidden-chip {
|
||||
font-size: 11px; color: var(--muted); border: 1px solid var(--border);
|
||||
border-radius: 5px; padding: 1px 6px; cursor: pointer; max-width: 160px;
|
||||
overflow: hidden; text-overflow: ellipsis; white-space: nowrap;
|
||||
}
|
||||
.lv-hidden-chip:hover { border-color: var(--accent2); color: var(--accent2); }
|
||||
.lv-view-on { border-color: var(--accent) !important; color: var(--accent); }
|
||||
.lv-pinadd { position: relative; display: inline-flex; }
|
||||
.lv-pinadd input { padding: 3px 6px; font-size: 12px; }
|
||||
.lv-pinadd-cands {
|
||||
position: absolute; top: 100%; left: 0; margin-top: 3px; z-index: 12;
|
||||
display: flex; gap: 3px; flex-wrap: wrap; max-width: 240px;
|
||||
background: #0f1114; border: 1px solid var(--border); border-radius: 6px;
|
||||
padding: 4px; box-shadow: 0 4px 14px rgba(0,0,0,.5);
|
||||
}
|
||||
.lv-pinadd-cands button { padding: 2px 7px; font-size: 12px; }
|
||||
.lv-layerrow { display: flex; align-items: center; gap: 8px; flex-wrap: wrap; padding: 0 2px 6px; }
|
||||
.lv-layerrow .layerpicker { flex: 1; min-width: 0; }
|
||||
.lv-lensbadge {
|
||||
font-size: 11px; color: var(--accent2); border: 1px solid var(--border);
|
||||
border-radius: 6px; padding: 1px 8px; max-width: 240px; overflow: hidden;
|
||||
text-overflow: ellipsis; white-space: nowrap;
|
||||
}
|
||||
.lv-svg { display: block; }
|
||||
.lv-toklabel { fill: var(--muted); font-size: 10px; text-anchor: start; }
|
||||
.lv-laylabel { fill: var(--muted); font-size: 10px; text-anchor: end; }
|
||||
.lv-word { fill: var(--text); font-size: 11px; }
|
||||
.lv-rank { fill: var(--muted); font-size: 7.5px; }
|
||||
.lv-cell { cursor: pointer; }
|
||||
.lv-cell:hover rect { stroke: var(--accent2); }
|
||||
.lv-tip {
|
||||
position: fixed; z-index: 10; background: #0f1114; border: 1px solid var(--border);
|
||||
border-radius: 6px; padding: 8px 10px; font-size: 12px; pointer-events: none;
|
||||
max-width: 260px; box-shadow: 0 4px 18px rgba(0,0,0,.5);
|
||||
}
|
||||
.lv-tip-head { color: var(--muted); font-size: 11px; margin-bottom: 5px; }
|
||||
.lv-trans { color: var(--accent2); font-size: 11px; font-style: italic; }
|
||||
.lv-tip-row { display: flex; justify-content: space-between; gap: 14px; }
|
||||
.lv-tip-row span:last-child { color: var(--muted); }
|
||||
.lv-tip-trans { border-top: 1px solid var(--border); margin-top: 6px; padding-top: 5px; }
|
||||
.lv-tip-trans .lv-tip-head { margin-bottom: 3px; }
|
||||
.lv-pins { border-top: 1px solid var(--border); margin-top: 8px; padding-top: 8px; }
|
||||
.lv-pinblock { margin-bottom: 12px; overflow-x: auto; }
|
||||
/* graph on top, heatmap below — never side by side (svgs are inline by default) */
|
||||
.lv-pinblock svg { display: block; }
|
||||
.lv-pinname { font-size: 12px; color: var(--accent); margin-bottom: 4px; }
|
||||
.lv-pinsub { color: inherit; }
|
||||
.lv-copybtn {
|
||||
padding: 0 5px; font-size: 12px; margin: 0 2px; line-height: 1.4;
|
||||
border: 1px solid var(--border); border-radius: 4px; color: var(--muted);
|
||||
background: var(--panel2); cursor: pointer; vertical-align: baseline;
|
||||
}
|
||||
.lv-copybtn:hover { color: var(--accent2); border-color: var(--accent2); }
|
||||
.lv-pinhov { color: var(--accent2); font-weight: 600; }
|
||||
.lv-grid { stroke: rgba(255,255,255,.06); }
|
||||
.lv-axis { fill: var(--muted); font-size: 9px; text-anchor: end; }
|
||||
.lv-axis-hot { fill: var(--accent2); font-weight: 700; }
|
||||
.lv-laylegend { display: flex; flex-wrap: wrap; gap: 3px; margin-bottom: 5px; }
|
||||
.lv-laychip {
|
||||
font-size: 10px; padding: 0 5px; border: 1px solid; border-radius: 8px;
|
||||
cursor: pointer; user-select: none; line-height: 1.5;
|
||||
}
|
||||
.lv-laychip.dim { opacity: .3; }
|
||||
|
||||
.lens-live {
|
||||
border-top: 1px solid var(--border); background: var(--panel);
|
||||
padding: 8px 22px; max-height: 240px; overflow-y: auto;
|
||||
}
|
||||
.lens-head { display: flex; gap: 12px; align-items: center; font-size: 12px; color: var(--muted); margin-bottom: 6px; }
|
||||
.phase { padding: 1px 8px; border-radius: 8px; font-weight: 600; }
|
||||
.phase.reading { background: rgba(232, 161, 60, .18); color: var(--accent); }
|
||||
.phase.thinking { background: rgba(108, 184, 224, .18); color: var(--accent2); }
|
||||
.lens-rows { display: flex; flex-direction: column; gap: 2px; }
|
||||
.lens-row { display: flex; gap: 10px; align-items: baseline; white-space: nowrap; overflow: hidden; }
|
||||
.lens-row .lnum { color: var(--muted); font-size: 11px; width: 34px; flex-shrink: 0; }
|
||||
.lword { font-size: 13px; }
|
||||
.lword sup { font-size: 9px; color: var(--muted); }
|
||||
|
||||
.sys { padding: 8px 22px; border-bottom: 1px solid var(--border); background: var(--panel); }
|
||||
.sys summary { cursor: pointer; font-size: 12px; color: var(--muted); }
|
||||
.sys textarea { margin-top: 6px; min-height: 40px; }
|
||||
|
||||
/* ---- Token editor ---- */
|
||||
.editor {
|
||||
width: 400px; flex-shrink: 0; background: var(--panel);
|
||||
border-left: 1px solid var(--border);
|
||||
display: flex; flex-direction: column; min-height: 0;
|
||||
}
|
||||
.ed-head {
|
||||
display: flex; justify-content: space-between; align-items: center;
|
||||
padding: 10px 14px; border-bottom: 1px solid var(--border);
|
||||
color: var(--accent); font-weight: 600;
|
||||
}
|
||||
.ed-close { padding: 2px 9px; }
|
||||
.ed-body { flex: 1; overflow-y: auto; padding: 12px 14px; display: flex; flex-direction: column; gap: 16px; }
|
||||
.ed-section { display: flex; flex-direction: column; gap: 7px; }
|
||||
.ed-section h3 {
|
||||
font-size: 12px; text-transform: uppercase; letter-spacing: .08em;
|
||||
color: var(--muted); font-weight: 600;
|
||||
}
|
||||
.editor .src { color: var(--muted); font-size: 11px; }
|
||||
.ed-flash { animation: ed-flash 1.5s ease-out; border-radius: 8px; }
|
||||
@keyframes ed-flash {
|
||||
0% { background: rgba(232,161,60,.25); box-shadow: 0 0 0 2px var(--accent); }
|
||||
100% { background: transparent; box-shadow: none; }
|
||||
}
|
||||
.ed-toggle-on { border-color: var(--accent) !important; color: var(--accent); }
|
||||
|
||||
/* Mode toggle: steering (exploration) ↔ read projection (exportable).
|
||||
Sliding thumb, two halves tiling the track (4px insets, seam at the center). */
|
||||
.mode-toggle {
|
||||
position: relative; display: grid; grid-template-columns: 1fr 1fr;
|
||||
background: var(--panel2); border: 1px solid var(--border); border-radius: 10px;
|
||||
padding: 4px; isolation: isolate; user-select: none;
|
||||
}
|
||||
.mt-thumb {
|
||||
position: absolute; z-index: 0; top: 4px; bottom: 4px; left: 4px;
|
||||
width: calc(50% - 4px); border-radius: 7px;
|
||||
transition: transform .24s cubic-bezier(.34,1.2,.44,1), background .2s, box-shadow .2s;
|
||||
}
|
||||
.mt-steer .mt-thumb {
|
||||
transform: translateX(0);
|
||||
background: linear-gradient(135deg, rgba(108,184,224,.92), rgba(108,184,224,.6));
|
||||
box-shadow: 0 2px 12px rgba(108,184,224,.4);
|
||||
}
|
||||
.mt-read .mt-thumb {
|
||||
transform: translateX(100%);
|
||||
background: linear-gradient(135deg, rgba(126,201,126,.94), rgba(126,201,126,.62));
|
||||
box-shadow: 0 2px 12px rgba(126,201,126,.4);
|
||||
}
|
||||
.mt-opt {
|
||||
position: relative; z-index: 1; background: none; border: none; border-radius: 7px;
|
||||
display: flex; flex-direction: column; align-items: center; gap: 3px;
|
||||
padding: 9px 6px 8px; cursor: pointer; color: var(--muted);
|
||||
transition: color .18s; text-align: center; line-height: 1.15;
|
||||
}
|
||||
.mt-opt:hover:not(.mt-on) { color: var(--text); }
|
||||
.mt-opt svg { opacity: .92; margin-bottom: 1px; }
|
||||
.mt-lab { font-size: 12px; font-weight: 600; }
|
||||
.mt-sub { font-size: 9.5px; letter-spacing: .02em; opacity: .82; }
|
||||
.mt-on { color: #14161a; }
|
||||
.mt-on .mt-sub { opacity: .74; }
|
||||
|
||||
.ed-rule { border: 1px solid var(--border); border-radius: 6px; padding: 5px 7px; }
|
||||
.ed-rule-off { opacity: .5; }
|
||||
.ed-rule-toggle {
|
||||
padding: 0 5px; font-size: 12px; border: none; background: transparent;
|
||||
color: var(--accent2); flex-shrink: 0; line-height: 1;
|
||||
}
|
||||
.ed-rule-off .ed-rule-toggle { color: var(--muted); }
|
||||
.ed-rule-toggle:hover { color: var(--accent); border-color: transparent; }
|
||||
.ed-rule-main { display: flex; align-items: center; gap: 6px; }
|
||||
.ed-rule-tok {
|
||||
font-size: 13px; overflow: hidden; text-overflow: ellipsis; white-space: nowrap;
|
||||
flex: 1; min-width: 96px;
|
||||
}
|
||||
.ed-rule-factor { width: 62px !important; flex-shrink: 0; padding: 3px 5px !important; font-size: 12px !important; }
|
||||
.ed-rule-del { padding: 1px 7px; font-size: 11px; border: none; color: var(--muted); flex-shrink: 0; }
|
||||
.ed-rule-del:hover { color: var(--danger); }
|
||||
.ed-rule-layers { margin-top: 6px; border-top: 1px solid var(--border); padding-top: 6px; }
|
||||
|
||||
.rulebar {
|
||||
display: flex; gap: 1px; cursor: pointer; flex-shrink: 1;
|
||||
min-width: 42px; max-width: 118px; overflow: hidden;
|
||||
padding: 3px 2px; border-radius: 4px; align-items: center;
|
||||
}
|
||||
.rulebar:hover { background: var(--panel2); }
|
||||
.rb-seg { width: 2px; min-width: 1px; height: 12px; background: var(--panel2); border-radius: 1px; }
|
||||
.rb-seg.on { background: var(--accent); }
|
||||
|
||||
.layerpicker { display: flex; flex-direction: column; gap: 5px; }
|
||||
.lp-cells { display: flex; flex-wrap: wrap; gap: 3px; }
|
||||
.lp-cell {
|
||||
min-width: 22px; text-align: center; font-size: 11px; padding: 2px 3px;
|
||||
border: 1px solid var(--border); border-radius: 4px; cursor: pointer;
|
||||
color: var(--muted); user-select: none;
|
||||
}
|
||||
.lp-cell:hover { border-color: var(--accent2); }
|
||||
.lp-cell.on { background: rgba(232,161,60,.2); border-color: var(--accent); color: var(--accent); }
|
||||
.lp-cell.lp-approx { border-style: dashed; opacity: .75; }
|
||||
.lp-quick { display: flex; gap: 5px; }
|
||||
.lp-quick button { padding: 2px 9px; font-size: 11px; }
|
||||
|
||||
.ed-cands { display: flex; flex-wrap: wrap; gap: 4px; align-items: center; font-size: 11px; }
|
||||
.ed-cands button { padding: 2px 8px; font-size: 12px; }
|
||||
.ed-cands .ed-cand-on { border-color: var(--accent); color: var(--accent); background: rgba(232,161,60,.12); }
|
||||
|
||||
.ed-group { border: 1px solid var(--accent2); border-radius: 8px; padding: 9px; }
|
||||
.ed-group h3 { color: var(--accent2); }
|
||||
@@ -0,0 +1,7 @@
|
||||
// Token display: leading/trailing spaces made visible with ˽ (like the original
|
||||
// JLens) — ' Paris' → '˽Paris', 'foo ' → 'foo˽'. Without this it's impossible to
|
||||
// tell ' Euro' from 'Euro' in the UI.
|
||||
export function fmtTok(s) {
|
||||
if (s == null) return ''
|
||||
return String(s).replace(/^ +| +$/g, (m) => '˽'.repeat(m.length))
|
||||
}
|
||||
@@ -0,0 +1,16 @@
|
||||
import { defineConfig } from 'vite'
|
||||
import react from '@vitejs/plugin-react'
|
||||
|
||||
// Dev proxy target: the backend port (run.py --port), overridable so a
|
||||
// non-default instance can be hot-reload developed too.
|
||||
const backend = `127.0.0.1:${process.env.JWASH_PORT || 8381}`
|
||||
|
||||
export default defineConfig({
|
||||
plugins: [react()],
|
||||
server: {
|
||||
proxy: {
|
||||
'/api': `http://${backend}`,
|
||||
'/ws': { target: `ws://${backend}`, ws: true },
|
||||
},
|
||||
},
|
||||
})
|
||||
Reference in New Issue
Block a user