evo2 SAE recipe: feature-explorer dashboard (viz)

interpretability/sparse_autoencoders/recipes/evo2/Dockerfile

@@ -5,12 +5,40 @@
 # (inheriting its pinned megatron-bridge / causal-conv1d / TransformerEngine versions), then
 # installs the sae library + this recipe on top. We don't reimplement the megatron build here.
 #
+# By DEFAULT this builds the engine + server only (no Node, no front-end) — the SAE-serving image.
+# Pass --build-arg WITH_DASHBOARD=1 to ALSO build the dashboard front-end (a Node build stage) and
+# bake it in, so one container serves the UI (at /) and the API (at /api). The Node toolchain is
+# only pulled when WITH_DASHBOARD=1, and never ends up in the runtime image either way.
+#
 # Build from the REPO ROOT (the context must include the recipes/evo2_megatron sibling):
+#   # engine + server only (default):
 #   docker build -f interpretability/sparse_autoencoders/recipes/evo2/Dockerfile -t evo2-sae .
-# Run (needs a GPU + checkpoints; see the recipe README):
+#   # engine + server + dashboard (single-container UI + API):
+#   docker build --build-arg WITH_DASHBOARD=1 -f interpretability/sparse_autoencoders/recipes/evo2/Dockerfile -t evo2-sae .
+#
+# Serve (needs a GPU + checkpoints; see the recipe README). With WITH_DASHBOARD=1 the UI is at / too:
+#   docker run --gpus all -p 8001:8001 \
+#     -e EVO2_CKPT_DIR=/ckpt/evo2 -e SAE_CKPT_PATH=/ckpt/sae.pt -e EMBEDDING_LAYER=26 \
+#     -v /my/checkpoints:/ckpt evo2-sae scripts/launch_inference.sh serve   # -> http://localhost:8001
+# Or run the tests:
 #   docker run --gpus all -it evo2-sae bash -lc "source .ci_test_env.sh && pytest tests/"
+
+ARG WITH_DASHBOARD=0
 ARG BASE_IMAGE=nvcr.io/nvidia/pytorch:26.04-py3
-FROM ${BASE_IMAGE}
+
+# Front-end build stage — compiles the dashboard to static files. Built ONLY when WITH_DASHBOARD=1
+# (nothing references it otherwise, so BuildKit skips it). Node lives only here, never at runtime.
+FROM node:20-slim AS frontend
+WORKDIR /frontend
+# Copy manifests first so `npm ci` is cached unless dependencies change.
+COPY interpretability/sparse_autoencoders/recipes/evo2/feature_explorer/package.json \
+     interpretability/sparse_autoencoders/recipes/evo2/feature_explorer/package-lock.json ./
+RUN npm ci
+COPY interpretability/sparse_autoencoders/recipes/evo2/feature_explorer/ ./
+RUN npm run build  # -> /frontend/dist
+
+# Engine + server image (no front-end). This is the default build.
+FROM ${BASE_IMAGE} AS engine
 
 WORKDIR /workspace
 # Two layers so editing SAE code doesn't rebuild the ~30-min megatron stack:
@@ -34,3 +62,15 @@ RUN bash .ci_build.sh install
 # Default to the built venv so `docker run … pytest tests/` (or importing Evo2SAE) just works.
 ENV VIRTUAL_ENV=/workspace/recipes/evo2_megatron/.venv
 ENV PATH="$VIRTUAL_ENV/bin:$PATH"
+
+# Engine + the baked dashboard: serve the UI at / alongside the API at /api from one container.
+FROM engine AS engine-with-dashboard
+COPY --from=frontend /frontend/dist \
+     /workspace/interpretability/sparse_autoencoders/recipes/evo2/feature_explorer/dist
+ENV DASHBOARD_DIST=/workspace/interpretability/sparse_autoencoders/recipes/evo2/feature_explorer/dist
+
+# Select the final image by WITH_DASHBOARD (0 = engine only [default], 1 = + dashboard). Only the
+# selected stage's dependencies are built, so the default never pulls Node / builds the front-end.
+FROM engine-with-dashboard AS final-1
+FROM engine AS final-0
+FROM final-${WITH_DASHBOARD} AS final

interpretability/sparse_autoencoders/recipes/evo2/feature_explorer/.gitignore

@@ -0,0 +1,7 @@
+node_modules/
+package-lock.json
+dist/
+.vite/
+
+# generated dashboard data (build it from your SAE/annotations, do not commit)
+public/*.parquet

interpretability/sparse_autoencoders/recipes/evo2/feature_explorer/README.md

@@ -0,0 +1,115 @@
+# Evo2 SAE Feature Explorer (front-end)
+
+Interactive dashboard for Evo2 SAE features, with four tabs:
+**Feature atlas**, **Sequence inspector**, **Generative steering**, and **Sequence UMAP**.
+
+This directory is the **front-end only** (React + Vite). Its backend is the standalone
+[`evo2_sae`](../src/evo2_sae) engine — the viz is just a UI over its HTTP API, so there is no
+model code here. The front-end always calls the API under **`/api`** (same path in dev and
+production), so the only thing that ever changes is *where* `/api` is served from.
+
+## Three ways to run
+
+### 1. One container — recommended for sharing / deploy
+
+Pass `--build-arg WITH_DASHBOARD=1` to the recipe [`Dockerfile`](../Dockerfile) and it builds this
+front-end to static files and bakes them in, so a **single container serves the dashboard and the
+API on one port** — no Node and no second process at runtime. This is what to hand a coworker or
+put behind an SSO proxy. (Without the flag, the image is engine + server only — see the note below.)
+
+```bash
+# build from the REPO ROOT (the build context needs the recipes/evo2_megatron sibling):
+docker build --build-arg WITH_DASHBOARD=1 \
+  -f interpretability/sparse_autoencoders/recipes/evo2/Dockerfile -t evo2-sae .
+
+# run with a GPU + your checkpoints, then open http://localhost:8001
+docker run --gpus all -p 8001:8001 \
+  -e EVO2_CKPT_DIR=/ckpt/evo2 -e SAE_CKPT_PATH=/ckpt/sae.pt -e EMBEDDING_LAYER=26 \
+  -v /path/to/checkpoints:/ckpt evo2-sae scripts/launch_inference.sh serve
+# -> dashboard + API both on http://localhost:8001  (/ = UI, /api = backend)
+```
+
+The first build compiles the megatron stack (~30 min) and is layer-cached afterward; with the flag,
+a Node build stage produces the static bundle and the server mounts it at `/` via `DASHBOARD_DIST`.
+No Node ends up in the runtime image. **Default build (no flag) is engine + server only** — it never
+pulls Node or builds the front-end, so an SAE-only deployment isn't coupled to the dashboard
+toolchain. See the [recipe Dockerfile](../Dockerfile) for the layer layout.
+
+> **Feature-atlas data:** the atlas tab loads three parquets (`features_atlas`, `feature_metadata`,
+> `feature_examples`) from the served directory. They aren't baked in (per-SAE generated data), so put
+> them in `$DASHBOARD_DIST` — either generate them in the same container with
+> `scripts/dashboard.py atlas|examples --output-dir "$DASHBOARD_DIST" …` (see mode 3), or `cp` a
+> pre-made set in before `serve`. Without them, the other three tabs still work via `/api`.
+
+### 2. Local dev — UI iteration with hot reload
+
+Needs **Node ≥ 18** (for Vite), plus a GPU + checkpoints for the live tabs. Two processes:
+
+```bash
+# backend: loads Evo2 + the SAE, serves the API under /api on :8001
+../scripts/launch_inference.sh serve            # or: python -m evo2_sae.cli serve
+
+# front-end: Vite dev server on :5176 (hot reload)
+../scripts/launch_dashboard.py                  # stages atlas data if --data-dir is given, then runs Vite
+#   or, for raw front-end dev:  npm install && npm run dev
+```
+
+Vite proxies `/api` **straight through** to `http://localhost:8001` (no path rewrite — see
+`vite.config.js`), so dev hits the same `/api/*` paths as the single-container build. Point it at a
+different backend with `VITE_BACKEND`. Configure the backend via the env vars in `launch_inference.sh`.
+
+To reach a remote box, tunnel the Vite port only (Vite proxies `/api` on the box):
+
+```bash
+tsh ssh -L 5176:localhost:5176 <gpu-box>        # then open http://localhost:5176
+```
+
+If instead you open the dev server through a **public hostname** (brev `*.brevlab.com`, ngrok,
+Codespaces `*.github.dev`) rather than a `localhost` tunnel, Vite's Host check applies. Those
+domains are pre-allowed in `vite.config.js` (`server.allowedHosts`); for a different tunnel, add
+its host there (or set `allowedHosts: true` to allow any).
+
+### 3. Offline / static — no backend
+
+The dashboard degrades gracefully: it probes `/api/health`, and when there's **no live backend** it
+hides the tabs that need the model and keeps the ones that read static files.
+
+| Tab                     | Needs backend?           | Offline source                                                |
+| ----------------------- | ------------------------ | ------------------------------------------------------------- |
+| **Feature atlas**       | no                       | the atlas parquets (`--data-dir`)                             |
+| **Sequence UMAP**       | no *iff* a bundle exists | `sequmap_embeddings.json` (precomputed; UMAP runs in-browser) |
+| **Generative steering** | yes                      | hidden offline                                                |
+| **Sequence inspector**  | yes                      | hidden offline                                                |
+
+So with no backend you always get the **Feature atlas**, plus **Sequence UMAP** if you precompute its
+bundle. Steering and the live inspector require `serve`.
+
+```bash
+# (one-time, needs the model) precompute the static artifacts into one dir:
+python ../scripts/dashboard.py atlas      --activations-dir $STORE --output-dir dashboard_data   # atlas tab
+python ../scripts/dashboard.py examples   --examples-fasta lib.fa  --output-dir dashboard_data   # example cards
+python ../scripts/dashboard.py embeddings --examples-fasta lib.fa  --output-dir dashboard_data   # Sequence-UMAP bundle
+#   (env: SAE_CKPT_PATH, EVO2_CKPT_DIR, FEATURE_ANNOTATIONS — same as launch_inference.sh)
+
+# serve the static dashboard — NO backend, NO GPU (needs Node for the dev server, or `npm run build`):
+python ../scripts/launch_dashboard.py --data-dir dashboard_data
+```
+
+`dashboard.py embeddings` writes `sequmap_embeddings.json` (the same shape `/api/gene_embed` returns).
+The viz auto-detects it (override the path with `?embeddings=<url>`). With all artifacts staged,
+`npm run build` produces a fully static site you can host anywhere (HF Spaces static / S3 / Pages) —
+the interactive steering/inspector tabs light up automatically if a backend later becomes reachable.
+
+## Tabs
+
+- **Feature atlas** — browse every SAE feature: firing rate, decoder-space UMAP, top-activating
+  examples, labels. Reads precomputed static files.
+- **Sequence inspector** — paste a sequence, see per-base SAE activations (top-k or picked features).
+  Live encode.
+- **Generative steering** — generate DNA while clamping chosen features on the continuation, vs. the
+  unsteered baseline. Live generation.
+- **Sequence UMAP** — embed a set of sequences (pooled per-feature vectors), UMAP them, color/re-project
+  by a feature. Live backend or a precomputed bundle.
+
+Each tab shows its own **Limitations** note in-app (context-length caps, the ±300 steering clamp,
+stochastic 2-D UMAP, etc.).

interpretability/sparse_autoencoders/recipes/evo2/feature_explorer/index.html

@@ -0,0 +1,107 @@
+<!DOCTYPE html>
+<html lang="en">
+  <head>
+    <meta charset="UTF-8" />
+    <meta name="viewport" content="width=device-width, initial-scale=1.0" />
+    <title>Evo 2 SAE Feature Explorer</title>
+    <style>
+      @font-face {
+        font-family: 'NVIDIA Sans';
+        font-style: normal;
+        src: url(https://brand-assets.cne.ngc.nvidia.com/assets/fonts/nvidia-sans/1.0.0/NVIDIASans_Lt.woff2);
+        font-weight: 300;
+      }
+      @font-face {
+        font-family: 'NVIDIA Sans';
+        font-style: italic;
+        src: url(https://brand-assets.cne.ngc.nvidia.com/assets/fonts/nvidia-sans/1.0.0/NVIDIASans_LtIt.woff2);
+        font-weight: 300;
+      }
+      @font-face {
+        font-family: 'NVIDIA Sans';
+        font-style: normal;
+        src: url(https://brand-assets.cne.ngc.nvidia.com/assets/fonts/nvidia-sans/1.0.0/NVIDIASans_Rg.woff2);
+        font-weight: normal;
+      }
+      @font-face {
+        font-family: 'NVIDIA Sans';
+        font-style: italic;
+        src: url(https://brand-assets.cne.ngc.nvidia.com/assets/fonts/nvidia-sans/1.0.0/NVIDIASans_It.woff2);
+        font-weight: normal;
+      }
+      @font-face {
+        font-family: 'NVIDIA Sans';
+        font-style: normal;
+        src: url(https://brand-assets.cne.ngc.nvidia.com/assets/fonts/nvidia-sans/1.0.0/NVIDIASans_Bd.woff2);
+        font-weight: bold;
+      }
+      @font-face {
+        font-family: 'NVIDIA Sans';
+        font-style: italic;
+        src: url(https://brand-assets.cne.ngc.nvidia.com/assets/fonts/nvidia-sans/1.0.0/NVIDIASans_BdIt.woff2);
+        font-weight: bold;
+      }
+      :root {
+        --bg: #f5f5f5;
+        --bg-card: #fff;
+        --bg-card-expanded: #fafafa;
+        --bg-example: #fff;
+        --bg-input: #fff;
+        --border: #e0e0e0;
+        --border-light: #eee;
+        --border-input: #ddd;
+        --text: #333;
+        --text-secondary: #666;
+        --text-tertiary: #888;
+        --text-muted: #999;
+        --text-heading: #000;
+        --accent: #76b900;
+        --highlight-border: #222;
+        --highlight-shadow: rgba(0,0,0,0.15);
+        --link: #2563eb;
+        --loading-bar-bg: #e0e0e0;
+        --density-bar-bg: #e0e0e0;
+        --scrollbar-thumb: #ccc;
+      }
+      :root.dark {
+        --bg: #000;
+        --bg-card: #000;
+        --bg-card-expanded: #000;
+        --bg-example: #0a0a0a;
+        --bg-input: #0a0a0a;
+        --border: #444;
+        --border-light: #3a3a3a;
+        --border-input: #4a4a4a;
+        --text: #E0E0E0;
+        --text-secondary: #bbb;
+        --text-tertiary: #999;
+        --text-muted: #777;
+        --text-heading: #fff;
+        --accent: #76b900;
+        --highlight-border: #76b900;
+        --highlight-shadow: rgba(118,185,0,0.3);
+        --link: #76b900;
+        --loading-bar-bg: #444;
+        --density-bar-bg: #444;
+        --scrollbar-thumb: #555;
+      }
+      * {
+        margin: 0;
+        padding: 0;
+        box-sizing: border-box;
+      }
+      body {
+        font-family: 'NVIDIA Sans', -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
+        background: var(--bg);
+        color: var(--text);
+      }
+    </style>
+    <!-- KaTeX (math rendering for the steering equation) — CDN, mirrors the font loading above -->
+    <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.11/dist/katex.min.css" />
+    <script defer src="https://cdn.jsdelivr.net/npm/katex@0.16.11/dist/katex.min.js"></script>
+  </head>
+  <body>
+    <div id="root"></div>
+    <script type="module" src="/src/index.jsx"></script>
+  </body>
+</html>

interpretability/sparse_autoencoders/recipes/evo2/feature_explorer/package.json

@@ -0,0 +1,25 @@
+{
+  "name": "evo2-sae-dashboard",
+  "version": "0.1.0",
+  "private": true,
+  "type": "module",
+  "scripts": {
+    "dev": "vite",
+    "build": "vite build",
+    "preview": "vite preview"
+  },
+  "dependencies": {
+    "@uwdata/mosaic-core": "^0.21.1",
+    "@uwdata/mosaic-sql": "^0.21.1",
+    "@uwdata/vgplot": "^0.21.1",
+    "embedding-atlas": "^0.16.1",
+    "lucide-react": "^0.577.0",
+    "react": "^18.2.0",
+    "react-dom": "^18.2.0",
+    "umap-js": "^1.4.0"
+  },
+  "devDependencies": {
+    "@vitejs/plugin-react": "^4.2.0",
+    "vite": "^5.0.0"
+  }
+}