Spaces:
Sleeping
Sleeping
| """Seed pre-built experiments into Supabase. Run once after install. | |
| cd backend && python seed_experiments.py | |
| Each experiment is executed end-to-end with Gemma-2-2b-it + GemmaScope, | |
| including feature timelines, patch matrix, and a finding written by Claude. | |
| Updated for v2: | |
| - Supabase replaces MongoDB (synchronous supabase-py client) | |
| - sae_layer defaults to 12 (Gemma-2-2b-it mid-model) | |
| - patch matrix layers: [6, 12, 18] | |
| - model field: 'gemma-2-2b-it' | |
| """ | |
| from __future__ import annotations | |
| import asyncio | |
| import logging | |
| import os | |
| import time | |
| import uuid | |
| from pathlib import Path | |
| from dotenv import load_dotenv | |
| from neuroscope.firebase_init import get_db | |
| ROOT = Path(__file__).parent | |
| load_dotenv(ROOT / ".env", override=True) | |
| from neuroscope import llm as ns_llm # noqa: E402 | |
| from neuroscope.runner import patch_matrix, run_trajectory # noqa: E402 | |
| logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(levelname)s %(message)s") | |
| log = logging.getLogger("seed") | |
| # Firebase DB client | |
| db = get_db() | |
| EXPERIMENTS = [ | |
| { | |
| "slug": "hallucination-propagation", | |
| "title": "Hallucination Propagation", | |
| "category": "Hallucination", | |
| "hypothesis": ( | |
| "H1: Hallucination-associated GemmaScope features activate 1-2 steps before\n" | |
| "the final hallucinated output, not at the output step itself." | |
| ), | |
| "task": "What year did Albert Einstein win the Nobel Prize, and for what discovery?", | |
| "n_steps": 4, | |
| "sae_layer": 12, | |
| "inject_observation": {2: "Note: Einstein won the Nobel Prize in 1933 for relativity."}, | |
| "finding_seed": ( | |
| "This experiment injects a false context at step 2 (claiming Einstein won the Nobel for " | |
| "relativity in 1933, both false) and watches whether the hallucination risk signal " | |
| "spikes BEFORE the final output is generated. If H1 holds, we expect GemmaScope feature " | |
| "drift and risk elevation at steps 2-3, with the bad output crystallizing at step 4." | |
| ), | |
| }, | |
| { | |
| "slug": "tool-call-prediction", | |
| "title": "Tool-Call Prediction", | |
| "category": "Tool-use", | |
| "hypothesis": ( | |
| "H2: Layer 12 residual stream activations can predict which tool the agent will call\n" | |
| "BEFORE the action token is emitted." | |
| ), | |
| "task": "You need the current weather in Paris. Choose between search, lookup, or calc, then provide the input.", | |
| "n_steps": 3, | |
| "sae_layer": 12, | |
| "inject_observation": None, | |
| "finding_seed": ( | |
| "Examines tool-selection circuits across steps. If H2 holds, the GemmaScope features " | |
| "active at the THOUGHT token already encode the upcoming ACTION choice, visible as a " | |
| "stable set of co-active features 1-2 tokens before the action surfaces." | |
| ), | |
| }, | |
| { | |
| "slug": "reasoning-collapse", | |
| "title": "Reasoning Collapse", | |
| "category": "Multi-hop", | |
| "hypothesis": ( | |
| "H3: Cross-step causal patching of layer 12 activations from an earlier 'still on-track'\n" | |
| "step can restore a collapsed reasoning chain." | |
| ), | |
| "task": "If a train leaves at 14:30 and travels for 2 hours 45 minutes, what time does it arrive? Show work.", | |
| "n_steps": 4, | |
| "sae_layer": 12, | |
| "inject_observation": None, | |
| "finding_seed": ( | |
| "Long arithmetic chains where Gemma-2-2b-it may collapse. Cross-step patches identify " | |
| "the step+layer where the chain breaks; significant KL on (source=1, target=4) at layer 12 " | |
| "would suggest the model carried early-step structure but lost it mid-chain." | |
| ), | |
| }, | |
| { | |
| "slug": "ioi-persistence", | |
| "title": "IOI Persistence Across Steps", | |
| "category": "Circuit universality", | |
| "hypothesis": ( | |
| "H4: The IOI (Indirect Object Identification) circuit discovered in single-prompt\n" | |
| "settings persists and re-activates across multi-step agent reasoning." | |
| ), | |
| "task": "When Sarah and Tom went to the store, Sarah gave a book to whom?", | |
| "n_steps": 3, | |
| "sae_layer": 12, | |
| "inject_observation": None, | |
| "finding_seed": ( | |
| "Probes whether the IOI circuit — a canonical interpretability result from " | |
| "Wang et al. — remains visible at GemmaScope-feature granularity across three " | |
| "reasoning steps. Bridges single-prompt mechanistic interpretability to trajectory-level analysis." | |
| ), | |
| }, | |
| { | |
| "slug": "self-correction", | |
| "title": "Self-Correction Mechanism", | |
| "category": "Self-correction", | |
| "hypothesis": ( | |
| "H5: When the model catches its own error mid-trajectory, a distinct set of GemmaScope\n" | |
| "features activates that does not appear in failure trajectories." | |
| ), | |
| "task": "Compute 23 * 17. Then double-check by computing 17 * 23 and verify they match.", | |
| "n_steps": 4, | |
| "sae_layer": 12, | |
| "inject_observation": None, | |
| "finding_seed": ( | |
| "Tests whether 'self-correction' is mechanistically distinguishable from 'self-doubt'. " | |
| "We expect a particular subset of mid-layer GemmaScope features to spike specifically " | |
| "when the model verifies a prior step." | |
| ), | |
| }, | |
| ] | |
| async def seed(): | |
| from neuroscope.loader import MODEL_NAME, SAE_RELEASE | |
| is_gemma = "gemma" in MODEL_NAME.lower() | |
| patch_layers = [6, 12, 18] if is_gemma else [3, 7, 10] | |
| for spec in EXPERIMENTS: | |
| # Check if already seeded | |
| existing = db.collection("experiments").where("slug", "==", spec["slug"]).limit(1).get() | |
| if existing and not os.environ.get("FORCE_RESEED"): | |
| log.info("%s already exists; skipping (set FORCE_RESEED=1 to overwrite)", spec["slug"]) | |
| continue | |
| log.info("== Seeding experiment: %s ==", spec["slug"]) | |
| t0 = time.time() | |
| run_id = f"exp-{spec['slug']}" | |
| adjusted_sae_layer = spec["sae_layer"] if is_gemma else 7 | |
| result = run_trajectory( | |
| run_id=run_id, | |
| task=spec["task"], | |
| n_steps=spec["n_steps"], | |
| sae_layer=adjusted_sae_layer, | |
| inject_context_at_step=spec["inject_observation"], | |
| ) | |
| def target_prompt_fn(step_n: int) -> str: | |
| for s in result["steps"]: | |
| if s["step_n"] == step_n: | |
| return s["prompt"] | |
| return "" | |
| pm = patch_matrix(result["steps"], target_prompt_fn, layers=patch_layers) | |
| ctx = { | |
| "task": spec["task"], | |
| "hypothesis": spec["hypothesis"], | |
| "model": MODEL_NAME, | |
| "sae_release": SAE_RELEASE, | |
| "steps": [ | |
| { | |
| "step_n": s["step_n"], | |
| "output": s["output"][:160], | |
| "hallucination": s["hallucination"], | |
| "top_features": s["top_features"][:5], | |
| } | |
| for s in result["steps"] | |
| ], | |
| "feature_timelines": result["feature_timelines"][:6], | |
| "patch_summary": { | |
| "max_kl": max((r["kl"] for r in pm), default=0.0), | |
| "significant_count": sum(1 for r in pm if r["significant"]), | |
| "n": len(pm), | |
| }, | |
| } | |
| finding = await ns_llm.report(ctx, session_id=f"seed-{spec['slug']}") | |
| doc = { | |
| "id": str(uuid.uuid4()), | |
| "slug": spec["slug"], | |
| "title": spec["title"], | |
| "category": spec["category"], | |
| "hypothesis": spec["hypothesis"], | |
| "task": spec["task"], | |
| "n_steps": spec["n_steps"], | |
| "sae_layer": adjusted_sae_layer, | |
| "model": MODEL_NAME, | |
| "steps": result["steps"], | |
| "feature_timelines": result["feature_timelines"], | |
| "patch_matrix": pm, | |
| "patch_matrix_summary": { | |
| "layers": patch_layers, | |
| "n_results": len(pm), | |
| "max_kl": max((r["kl"] for r in pm), default=0.0), | |
| "significant_count": sum(1 for r in pm if r["significant"]), | |
| }, | |
| "finding_seed": spec["finding_seed"], | |
| "finding": finding, | |
| "total_elapsed_ms": result["total_elapsed_ms"], | |
| "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), | |
| } | |
| # Upsert on slug (use slug as document ID) | |
| db.collection("experiments").document(spec["slug"]).set(doc) | |
| log.info("== Seeded %s in %.1fs ==", spec["slug"], time.time() - t0) | |
| if __name__ == "__main__": | |
| asyncio.run(seed()) | |