gpt2 SAE recipe: interactive feature-explorer dashboard

interpretability/sparse_autoencoders/pyproject.toml

@@ -32,4 +32,4 @@ where = ["."]
 include = []
 
 [tool.uv.workspace]
-members = ["sae", "recipes/esm2"]
+members = ["sae", "recipes/esm2", "recipes/gpt2"]

interpretability/sparse_autoencoders/recipes/evo2/feature_explorer/src/Dashboard.jsx

@@ -4,7 +4,7 @@ import GenerativeSteering from './GenerativeSteering'
 import SequenceInspector from './SequenceInspector'
 import SequenceUMAPView, { EMBEDDINGS_URL } from './SequenceUMAPView'
 import { Sun, Moon } from 'lucide-react'
-import { useHealth } from './backend'
+import { useHealth, UI } from './backend'
 
 // Four-tab shell with graceful degradation when there's no live backend:
 //   offline:true     -> always available (reads static files only)
@@ -64,6 +64,12 @@ export default function Dashboard() {
     document.documentElement.classList.toggle('dark', dark)
   }, [dark])
 
+  // Keep the tab title in sync with the (possibly server-overridden) brand. `health` re-renders
+  // after useHealth merges the /health `ui` block into UI, so the title picks up UI.brand then.
+  useEffect(() => {
+    document.title = UI.brand
+  }, [health])
+
   // Probe for a precomputed Sequence-UMAP bundle so that tab survives without a backend.
   useEffect(() => {
     fetch(EMBEDDINGS_URL, { method: 'HEAD' })
@@ -84,7 +90,7 @@ export default function Dashboard() {
   return (
     <div style={S.shell}>
       <div style={S.tabBar}>
-        <span style={S.brand}>Evo 2 SAE Feature Explorer</span>
+        <span style={S.brand}>{UI.brand}</span>
         {TABS.map((t) => (
           <button key={t.id} onClick={() => setTab(t.id)} style={tab === t.id ? S.tabOn : S.tabOff}>
             {t.label}

interpretability/sparse_autoencoders/recipes/evo2/feature_explorer/src/backend.js

@@ -6,6 +6,12 @@ import { useEffect, useRef, useState } from 'react'
 
 export const BACKEND = (import.meta.env && import.meta.env.VITE_BACKEND) || '/api'
 
+// UI config, defaulting to the DNA (Evo 2) build. The server can override this at runtime by
+// returning a `ui` block from /health (see useHealth) — e.g. the GPT-2 build flips textMode on so
+// pasted natural-language text is not stripped to A/C/G/T/N, and sets its own brand. With no `ui`
+// block in health this stays DNA, so the Evo 2 build behaves exactly as before.
+export const UI = { textMode: false, brand: 'Evo 2 SAE Feature Explorer' }
+
 // Per-nucleotide letter colors (shared with the steering strips).
 export const BASE_COLORS = { A: '#59A14F', C: '#4E79A7', G: '#F28E2B', T: '#E15759', N: '#888', U: '#E15759' }
 
@@ -21,6 +27,9 @@ export function useHealth(pollMs = 4000) {
         const r = await fetchWithTimeout(`${BACKEND}/health`, { cache: 'no-store' }, 8000)
         if (!r.ok) throw new Error(`HTTP ${r.status}`)
         const info = await r.json()
+        // Merge any server-provided UI config (textMode/brand) so one built dashboard can serve both
+        // the DNA (Evo 2) and the text (GPT-2) backends. Absent `ui` block -> defaults unchanged.
+        if (info && info.ui) Object.assign(UI, info.ui)
         if (alive) setHealth({ status: info.ready ? 'ready' : 'loading', info })
       } catch (e) {
         if (alive) setHealth({ status: 'offline', error: String(e) })
@@ -70,7 +79,11 @@ export function activationColor(value, max) {
   return `rgba(${r}, ${g}, ${b}, ${(0.22 + 0.78 * t).toFixed(3)})`
 }
 
+// Sanitize pasted input before sending. DNA build (default): strip to A/C/G/T/N. Text build
+// (UI.textMode, set from the /health `ui` block): pass the raw text through untouched, since the
+// A/C/G/T/N filter would gut natural language ("LOOK AT THIS TEXT" -> "ATTTT").
 export function cleanDNA(raw) {
+  if (UI.textMode) return raw || ''
   return (raw || '').toUpperCase().replace(/[^ACGTN]/g, '')
 }
 

interpretability/sparse_autoencoders/recipes/gpt2/README.md

@@ -0,0 +1,86 @@
+# GPT-2 SAE Recipe
+
+An interactive **sparse-autoencoder feature explorer for GPT-2 small** — the natural-language sibling
+of the `evo2` / `esm2` / `codonfm` recipes, built on the same model-agnostic [`sae`](../../sae) library
+and the same feature-explorer dashboard.
+
+Unlike the other recipes, this one needs **no training and no gated model**: it loads OpenAI GPT-2 small
+(124M) through [TransformerLens](https://github.com/TransformerLensOrg/TransformerLens) and a
+**pretrained residual-stream SAE** (Joseph Bloom's `blocks.7.hook_resid_pre`, 24,576 features). It runs
+on a single GPU or on **CPU** (GPT-2 small is small), so it doubles as a zero-setup demo of the toolchain.
+
+## Quick start
+
+```bash
+# from the repo root (uv workspace)
+uv sync
+
+# serve the dashboard + API (loads model + SAE once; GPU if available, else CPU)
+gpt2-sae serve            # -> http://localhost:8749
+```
+
+The dashboard is **reused from the evo2 recipe** (`../evo2/feature_explorer/dist`) — the same built
+explorer, flipped into text mode at runtime via the `ui` block in `/api/health`, so we don't ship a
+near-duplicate copy. GPT-2 ships only its backend (`src/gpt2_sae/`) and its own `dashboard_data/` (the
+atlas parquets + 24,570 semantic feature labels); the server serves evo2's dist and shadows its data
+files with GPT-2's. Point `GPT2_SAE_DASHBOARD` at a dist built elsewhere to override.
+
+## Dashboard
+
+Four tabs (the same explorer used across recipes, re-skinned for text via the health `ui` block):
+
+- **Feature atlas** — browse every SAE feature: firing rate, decoder-space UMAP, top-activating text
+  snippets, and its label. Static — no backend needed.
+- **Text inspector** — paste text, see per-token feature activations (top-k or features you pick).
+- **Generative steering** — generate from a prompt while clamping chosen features on the continuation,
+  vs. an unsteered baseline.
+- **Text UMAP** — embed a set of texts into per-feature SAE vectors and lay them out in 2-D.
+
+Feature labels come from [Neuronpedia](https://www.neuronpedia.org/)'s auto-interp explanations for this
+exact SAE (index-aligned 1:1), so any concept is searchable in the pickers. A curated set of ⭐
+verified-steerable features (France/Paris, music, ocean, football, …) leads the catalog.
+
+## CLI
+
+```bash
+gpt2-sae annotate "the cat sat on the mat" -k 8            # top-k features per token
+gpt2-sae generate "I think the best thing to do is" \
+    -f 20174:45 -f 21634:40 -n 24 --baseline               # steer ocean + Paris features
+gpt2-sae serve 8749                                        # dashboard + API
+```
+
+## Rebuilding the bundled data (optional)
+
+The parquets + labels are committed under `dashboard_data/`. To regenerate them from scratch:
+
+```bash
+uv sync --extra build                       # scikit-learn + umap-learn
+python scripts/build_atlas.py               # features_atlas.parquet (UMAP + stats)
+python scripts/build_meta.py                # feature_metadata.parquet + feature_examples.parquet
+python scripts/fetch_neuronpedia.py         # neuronpedia_labels.json + bakes labels into the parquets
+python scripts/curate_steerable.py          # (server running) probe -> steer-test -> user_labels.json
+```
+
+These scripts write to `dashboard_data/` by default (override with `GPT2_SAE_DATA`); the server picks
+them up automatically. The dashboard front-end itself is not built here — it's reused from
+`../evo2/feature_explorer` (see `recipes/evo2/scripts/launch_dashboard.py` to build that dist).
+
+## Why TransformerLens
+
+The Bloom SAE was trained on **TransformerLens** activations (with `fold_ln` / centering), not raw
+HuggingFace hidden states. Encoding with raw HF activations gives garbage reconstructions (FVU ≫ 1);
+TransformerLens gives FVU ≈ 0.001. The `encode_fn` runtime must match the SAE's training runtime — hence
+`HookedTransformer` here rather than a bare `transformers` forward pass.
+
+## Layout
+
+```
+gpt2/
+├── src/gpt2_sae/        # server.py (Engine + FastAPI), cli.py
+├── dashboard_data/      # bundled data served to the dashboard: atlas parquets + labels
+├── scripts/             # build_atlas, build_meta, fetch_neuronpedia, curate_steerable
+└── tests/
+```
+
+The dashboard front-end is not vendored here — the server serves the shared, config-driven build from
+`../evo2/feature_explorer/dist` and shadows its data files with `dashboard_data/`.

interpretability/sparse_autoencoders/recipes/gpt2/dashboard_data/sequence_library.json

@@ -0,0 +1,26 @@
+[
+  {"symbol": "cat", "label": "animals", "species": "text", "sequence": "The cat stretched lazily in the warm afternoon sun and began to purr."},
+  {"symbol": "dog", "label": "animals", "species": "text", "sequence": "The loyal dog waited by the front door, tail wagging, for its owner to return."},
+  {"symbol": "bird", "label": "animals", "species": "text", "sequence": "A small sparrow hopped along the fence, chirping brightly at the morning light."},
+  {"symbol": "whale", "label": "animals", "species": "text", "sequence": "The enormous blue whale surfaced slowly, spraying a great plume of mist into the air."},
+  {"symbol": "python", "label": "code", "species": "text", "sequence": "def train(model, data): return model.fit(data, epochs=10, batch_size=32)"},
+  {"symbol": "javascript", "label": "code", "species": "text", "sequence": "const total = items.reduce((acc, x) => acc + x.price, 0);"},
+  {"symbol": "sql", "label": "code", "species": "text", "sequence": "SELECT name, COUNT(*) FROM orders GROUP BY name HAVING COUNT(*) > 5;"},
+  {"symbol": "bash", "label": "code", "species": "text", "sequence": "for f in *.txt; do grep -c error \"$f\" >> counts.log; done"},
+  {"symbol": "physics", "label": "science", "species": "text", "sequence": "Energy is conserved: the total energy of an isolated system remains constant over time."},
+  {"symbol": "biology", "label": "science", "species": "text", "sequence": "Ribosomes read messenger RNA and assemble amino acids into proteins during translation."},
+  {"symbol": "chemistry", "label": "science", "species": "text", "sequence": "When an acid reacts with a base, they neutralize each other to form water and a salt."},
+  {"symbol": "astronomy", "label": "science", "species": "text", "sequence": "Light from distant galaxies is redshifted, evidence that the universe is expanding."},
+  {"symbol": "soccer", "label": "sports", "species": "text", "sequence": "The striker curled a stunning free kick into the top corner in the final minute."},
+  {"symbol": "basketball", "label": "sports", "species": "text", "sequence": "She drained a three-pointer at the buzzer to win the championship game."},
+  {"symbol": "tennis", "label": "sports", "species": "text", "sequence": "His serve topped 130 miles per hour, acing the opponent to hold the set."},
+  {"symbol": "running", "label": "sports", "species": "text", "sequence": "The marathon runners paced themselves carefully over the first twenty miles."},
+  {"symbol": "stocks", "label": "finance", "species": "text", "sequence": "The stock market rallied after the central bank signaled it would cut interest rates."},
+  {"symbol": "banking", "label": "finance", "species": "text", "sequence": "She deposited her paycheck and transferred savings into a higher-yield account."},
+  {"symbol": "crypto", "label": "finance", "species": "text", "sequence": "Bitcoin's price swung wildly as traders reacted to the new regulatory announcement."},
+  {"symbol": "taxes", "label": "finance", "species": "text", "sequence": "Remember to file your quarterly estimated taxes before the April deadline."},
+  {"symbol": "baking", "label": "cooking", "species": "text", "sequence": "Cream the butter and sugar, then fold in the flour until the dough just comes together."},
+  {"symbol": "grilling", "label": "cooking", "species": "text", "sequence": "Sear the steak over high heat for two minutes per side, then let it rest before slicing."},
+  {"symbol": "soup", "label": "cooking", "species": "text", "sequence": "Simmer the onions, carrots, and celery in broth until tender and fragrant."},
+  {"symbol": "salad", "label": "cooking", "species": "text", "sequence": "Toss the crisp greens with olive oil, lemon juice, and a pinch of sea salt."}
+]

interpretability/sparse_autoencoders/recipes/gpt2/dashboard_data/user_labels.json

@@ -0,0 +1 @@
+{"21634": "\u2b50 France / Paris", "21275": "\u2b50 music / guitar", "23383": "\u2b50 food / chicken", "10193": "\u2b50 space / rockets", "15993": "\u2b50 ocean / waves", "24310": "cat", "479": "code / return", "20174": "\u2b50 ocean / the sea", "13744": "\u2b50 London / the city", "919": "\u2b50 football / touchdowns", "21815": "\u2b50 religion / priests", "9821": "\u2b50 art / painting", "15194": "\u2b50 debugging / code", "21505": "\u2b50 cooking / sauces", "3683": "\u2b50 darkness / the dark", "16995": "\u2b50 heaviness / weight", "24507": "\u2b50 legal argument", "17876": "\u2b50 politics / senators"}

interpretability/sparse_autoencoders/recipes/gpt2/pyproject.toml

@@ -0,0 +1,43 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "gpt2-sae"
+version = "0.1.0"
+description = "Sparse-autoencoder feature explorer for the GPT-2 small language model"
+requires-python = ">=3.10"
+
+dependencies = [
+    "sae",
+    "torch>=2.0",
+    "transformer-lens>=2.0",  # HookedTransformer — SAEs were trained in this runtime (fold_ln/center)
+    "huggingface-hub>=0.20",  # pulls the Bloom SAE weights
+    "safetensors>=0.4",
+    "numpy>=1.20",
+    "pyarrow>=23.0.0",  # atlas / metadata parquets
+    "fastapi>=0.110",  # server.py
+    "uvicorn>=0.29",  # serve entry point
+    "anyio>=4.0",  # server.py: threadpool concurrency cap (ships with fastapi, declared explicitly)
+]
+
+[project.optional-dependencies]
+# Only needed to (re)build the static atlas from scratch (scripts/); the recipe ships prebuilt parquets.
+build = [
+    "scikit-learn>=1.3",
+    "umap-learn>=0.5",
+]
+
+[project.scripts]
+gpt2-sae = "gpt2_sae.cli:main"
+gpt2-sae-serve = "gpt2_sae.server:main"
+
+# The `gpt2_sae` package (src/) holds the live inference engine, the feature-clamp steering hook, and
+# the FastAPI server + CLI. scripts/ (build_atlas, build_meta, fetch_neuronpedia, curate_steerable) are
+# standalone entry points that regenerate the bundled data in dashboard_data/. The dashboard itself is
+# reused from ../evo2/feature_explorer/dist (config-driven into text mode), not shipped here.
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.uv.sources]
+sae = { workspace = true }

interpretability/sparse_autoencoders/recipes/gpt2/scripts/build_atlas.py

@@ -0,0 +1,80 @@
+import os
+from pathlib import Path
+
+import pyarrow.parquet as pq
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from huggingface_hub import hf_hub_download
+from sae.analysis import compute_feature_stats, compute_feature_umap, save_feature_atlas
+from safetensors.torch import load_file
+from transformer_lens import HookedTransformer
+
+
+# Where to write the parquets — defaults to the dashboard's bundled data dir.
+OUT = Path(os.environ.get("GPT2_SAE_DATA", Path(__file__).resolve().parents[1] / "dashboard_data"))
+OUT.mkdir(parents=True, exist_ok=True)
+
+REPO = "jbloom/GPT2-Small-SAEs-Reformatted"
+HP = "blocks.7.hook_resid_pre"
+model = HookedTransformer.from_pretrained("gpt2")
+model.eval()
+w = load_file(hf_hub_download(REPO, f"{HP}/sae_weights.safetensors"))
+W_enc, W_dec, b_enc, b_dec = [w[k].float() for k in ("W_enc", "W_dec", "b_enc", "b_dec")]
+
+
+class SAEWrap(nn.Module):
+    """Minimal nn.Module wrapper over the Bloom SAE weights for the atlas builder."""
+
+    def __init__(s):
+        """Register the SAE weight/bias buffers and a decoder Linear."""
+        super().__init__()
+        s.register_buffer("W_enc", W_enc)
+        s.register_buffer("W_dec", W_dec)
+        s.register_buffer("b_enc", b_enc)
+        s.register_buffer("b_dec", b_dec)
+        s.hidden_dim = W_enc.shape[1]
+        s.input_dim = W_enc.shape[0]
+        s.decoder = nn.Linear(s.hidden_dim, s.input_dim, bias=False)
+        s.decoder.weight.data = W_dec.t().contiguous()  # (input_dim, hidden_dim)
+
+    def encode(s, x):
+        """Residual activations -> ReLU SAE codes."""
+        return F.relu((x - s.b_dec) @ s.W_enc + s.b_enc)
+
+    def decode(s, c):
+        """SAE codes -> reconstructed residual activations."""
+        return c @ s.W_dec + s.b_dec
+
+
+sae = SAEWrap()
+
+# activations: cache to disk so re-runs are fast
+if os.path.exists(str(OUT / "acts.pt")):
+    X = torch.load(str(OUT / "acts.pt"))
+else:
+    from datasets import load_dataset
+
+    ds = load_dataset("wikitext", "wikitext-2-raw-v1", split="test")
+    lines = [t for t in ds["text"] if len(t.strip()) > 40][:120]
+    acts = []
+    with torch.no_grad():
+        for t in lines:
+            _, cache = model.run_with_cache(t[:600], names_filter=HP)
+            acts.append(cache[HP][0].float().cpu())
+    X = torch.cat(acts, 0)[:6000]
+    torch.save(X, str(OUT / "acts.pt"))
+print(f"  activations {tuple(X.shape)}")
+
+stats, top_examples = compute_feature_stats(sae, X, device="cpu", top_k=10, batch_size=4096, show_progress=False)
+print(f"  stats: {len(stats)} features")
+geom = compute_feature_umap(sae, compute_clusters=False)
+print(f"  umap: {geom.umap_x.shape}")
+save_feature_atlas(stats, geom, str(OUT / "features_atlas.parquet"), top_examples=top_examples)
+
+t = pq.read_table(str(OUT / "features_atlas.parquet"))
+print(f"  ✅ features_atlas.parquet: {t.num_rows} rows, cols={t.column_names}")
+df = t.to_pandas()
+alive = (df["activation_freq"] > 0).sum()
+print(f"  alive features: {alive}/{len(df)}   x/y range: x[{df.x.min():.1f},{df.x.max():.1f}]")
+print(df[["feature_id", "activation_freq", "max_activation", "x", "y"]].head(3).to_string(index=False))

interpretability/sparse_autoencoders/recipes/gpt2/scripts/build_meta.py

@@ -0,0 +1,65 @@
+import os
+from pathlib import Path
+
+import pyarrow.parquet as pq
+import torch.nn as nn
+import torch.nn.functional as F
+from datasets import load_dataset
+from huggingface_hub import hf_hub_download
+from sae.analysis import export_text_features_parquet
+from sae.collector import TokenActivationCollector
+from safetensors.torch import load_file
+from transformer_lens import HookedTransformer
+
+
+# Where to write the parquets — defaults to the dashboard's bundled data dir.
+OUT = Path(os.environ.get("GPT2_SAE_DATA", Path(__file__).resolve().parents[1] / "dashboard_data"))
+OUT.mkdir(parents=True, exist_ok=True)
+
+REPO = "jbloom/GPT2-Small-SAEs-Reformatted"
+HP = "blocks.7.hook_resid_pre"
+model = HookedTransformer.from_pretrained("gpt2")
+model.eval()
+w = load_file(hf_hub_download(REPO, f"{HP}/sae_weights.safetensors"))
+W_enc, W_dec, b_enc, b_dec = [w[k].float() for k in ("W_enc", "W_dec", "b_enc", "b_dec")]
+
+
+class SAEWrap(nn.Module):
+    """Minimal nn.Module wrapper over the Bloom SAE weights for the metadata builder."""
+
+    def __init__(s):
+        """Register the SAE weight/bias buffers."""
+        super().__init__()
+        for n, t in [("W_enc", W_enc), ("W_dec", W_dec), ("b_enc", b_enc), ("b_dec", b_dec)]:
+            s.register_buffer(n, t)
+        s.hidden_dim = W_enc.shape[1]
+
+    def encode(s, x):
+        """Residual activations -> ReLU SAE codes."""
+        return F.relu((x - s.b_dec) @ s.W_enc + s.b_enc)
+
+
+sae = SAEWrap()
+
+
+def encode_fn(text):
+    """Encode one text -> (str tokens, SAE codes) for the collector."""
+    _, cache = model.run_with_cache(text, names_filter=HP)
+    x = cache[HP][0].float().cpu()
+    codes = sae.encode(x)
+    labels = model.to_str_tokens(text)
+    return labels, codes
+
+
+ds = load_dataset("wikitext", "wikitext-2-raw-v1", split="test")
+lines = [t[:600] for t in ds["text"] if len(t.strip()) > 40][:120]
+
+col = TokenActivationCollector(encode_fn, n_features=sae.hidden_dim)
+print("  collecting over corpus ...")
+res = col.collect(lines)
+print(f"  collected: {len(res.feature_stats)} features")
+export_text_features_parquet(res, output_dir=str(OUT), n_examples=5)
+
+for f in (str(OUT / "feature_metadata.parquet"), str(OUT / "feature_examples.parquet")):
+    t = pq.read_table(f)
+    print(f"  ✅ {f}: {t.num_rows} rows, cols={t.column_names}")