""" Data loader for swap experiment results (domain-agnostic). Reads from output/_batch/ directory structure: - _swaps/runs/{run_id}/by_source/*/to_*.json - Swap details (run-specific) - _swaps/runs/{run_id}/config_resolved.json - Domain & model metadata - _swaps/_analysis_v3/*.json - Analysis data - */manifest.json - Neuronpedia URLs Supports multiple experiment runs and auto-detects domain fields (e.g. state/capital, character/author) from config_resolved.json. Multi-dataset discovery: DemoRegistry scans output/* for run_manifest.json files that carry "display_demo": true and groups them by dataset (graphs_root). The demo can then switch both dataset and run without restarting the server. """ from pathlib import Path from typing import Any, Dict, List, Optional, Set import os import json import csv import re from functools import lru_cache from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit # --------------------------------------------------------------------------- # Control variant helpers # --------------------------------------------------------------------------- _VARIANT_SUFFIX_RE = re.compile(r"__(?:r\d+|add_.+|ctrl_.+|m_tuned)$") # --------------------------------------------------------------------------- # Logit-shift regime classification (Section 4.4 of METHODOLOGY_REPORT.md) # --------------------------------------------------------------------------- _REGIME_FLAT_THRESHOLD = 1.0 # |delta| below this is considered FLAT _REGIME_INFO = { "A": {"label": "Clean Redirection", "short": "A"}, "B": {"label": "Both Boosted", "short": "B"}, "C": {"label": "Differential Disruption", "short": "C"}, "D": {"label": "Generic Disruption", "short": "D"}, "E": {"label": "Pure Suppression", "short": "E"}, } def _classify_regime(pos0: Dict[str, Any]) -> Optional[str]: """Classify a swap into a logit-shift regime from position-0 comparison. Returns one of A/B/C/D/E or None if classification is impossible. """ if not pos0: return None tgt = pos0.get("target_logit_delta") src = pos0.get("source_logit_delta") flip = pos0.get("flip_at_0", False) if tgt is None or src is None: return None tgt_up = tgt > _REGIME_FLAT_THRESHOLD tgt_down = tgt < -_REGIME_FLAT_THRESHOLD src_down = src < -_REGIME_FLAT_THRESHOLD src_up = src > _REGIME_FLAT_THRESHOLD if tgt_up and src_down and flip: return "A" if tgt_up and src_up: return "B" if tgt_down and src_down and flip: return "C" if tgt_down and src_down and not flip: return "D" if not tgt_up and not tgt_down and src_down and flip: return "E" return None def _parse_to_slug(filename_stem: str) -> tuple: """ Extract ``(entity_slug, variant_suffix)`` from a swap result filename stem. Normal files: ``to_georgia_savannah`` -> ``("georgia_savannah", "")`` Variant files: ``to_georgia_savannah__r0`` -> ``("georgia_savannah", "r0")`` ``to_georgia_savannah__add_source`` -> ``("georgia_savannah", "add_source")`` Double suffix: ``to_buzz__add_sound__m_tuned`` -> ``("buzz", "add_sound__m_tuned")`` """ raw = filename_stem.replace("to_", "", 1) # Handle double suffix: strip __m_tuned, parse the inner variant, re-append m_tuned_tail = "" if raw.endswith("__m_tuned") and "__add_" in raw: m_tuned_tail = "__m_tuned" raw = raw[: -len("__m_tuned")] m = _VARIANT_SUFFIX_RE.search(raw) if m: return raw[: m.start()], raw[m.start() + 2 :] + m_tuned_tail return raw, "" def _is_control_variant(filename_stem: str) -> bool: """Return True if the filename contains a control variant suffix.""" return bool(_VARIANT_SUFFIX_RE.search(filename_stem.replace("to_", "", 1))) # --------------------------------------------------------------------------- # Run metadata helpers # --------------------------------------------------------------------------- _CONTROL_MODE_LABELS = { "labeled": "Labeled", "random_feature_matched": "Random x3", "additivity": "Field Additivity", "m_tuned": "Adaptive M", } # Run-id / experiment_name substrings that hide a run from the demo # selector. These are calibration sweeps or "labeled" control runs that # would otherwise dominate the dropdown without adding signal for the # end user. The cross-run-best aggregator applies the exact same filter # so "Best across runs" is restricted to runs the user can inspect. _RUN_BLACKLIST_SUBSTRINGS = ( "sweep_", # m-sweep calibration runs "entropy_study", # entropy calibration runs "highm_", # high-M ablations "_labeled", # *_labeled control replicas (kept on disk for analysis) ) def _is_dropdown_run(run_id: str, manifest: Dict[str, Any]) -> bool: """Return True if a run should appear in the demo run-selector. The aim is to keep only the three "headline" run families: - Random x3 (control_mode == random_feature_matched) - Field Additivity (control_mode == additivity) - Classic / canonical swap runs (everything else not blacklisted) """ rid = run_id.lower() exp_name = ( manifest.get("config", {}).get("experiment_name", "") or "" ).lower() for blocked in _RUN_BLACKLIST_SUBSTRINGS: if blocked in rid or blocked in exp_name: return False return True def _extract_control_mode(manifest: Dict[str, Any], run_dir: Optional[Path] = None) -> str: """Return the control mode from run metadata. Checks the manifest config first, then falls back to ``config_resolved.json`` on disk (which contains the full YAML config). """ mode = manifest.get("config", {}).get("control", {}).get("mode", "") if mode: return mode if run_dir is not None: resolved_path = run_dir / "config_resolved.json" if resolved_path.exists(): try: with open(resolved_path, "r", encoding="utf-8") as fh: resolved = json.load(fh) mode = resolved.get("control", {}).get("mode", "") if mode: return mode except (json.JSONDecodeError, IOError): pass return "labeled" def _semantic_run_label( run_id: str, manifest: Dict[str, Any], run_dir: Optional[Path] = None, ) -> str: """Build a concise, human-readable label that distinguishes control modes.""" mode = _extract_control_mode(manifest, run_dir=run_dir) mode_label = _CONTROL_MODE_LABELS.get(mode, mode.replace("_", " ").title()) exp_name = manifest.get("config", {}).get("experiment_name", "") if exp_name.startswith("fullscale_"): return mode_label if exp_name: pretty = exp_name.replace("_", " ").title() pretty = pretty.replace("50states", "50 States").replace("6states", "6 States") return pretty pretty = run_id.replace("_", " ").title() pretty = pretty.replace("50states", "50 States").replace("6states", "6 States") pretty = pretty.replace("Gpu", "GPU") return pretty def _is_subpath(child: Path, parent: Path) -> bool: """Return True if *child* is equal to or nested under *parent*.""" try: child.relative_to(parent) return True except ValueError: return False # --------------------------------------------------------------------------- # DemoRegistry -- multi-dataset discovery # --------------------------------------------------------------------------- class DemoRegistry: """Scan an output root for demo-enabled runs and manage the active dataset. Discovery rules: - Walk output_root/*/_swaps/runs/*/run_manifest.json - Keep only manifests where display_demo == True - Derive dataset_dir from manifest["config"]["inputs"]["graphs_root"] falling back to the grandparent of the run dir (_swaps/runs/ -> dataset) - Group runs by resolved dataset_dir (normalised to an absolute Path) The active dataset+run is exposed through ``active_loader``. """ DEFAULT_DATASET_ID = "usa_states_batch" # Default run for the homepage: the field-additivity sweep on USA # states (paper's main result, drives the 73% T5 / 83% T3+ headline # numbers). Falls through to ``_pick_default_run``'s fallback chain # for other datasets that don't ship this exact run id. DEFAULT_RUN_ID = "fullscale_usa_field_add" # Datasets present in the output root but hidden from the public demo. # The data still exists on disk (kept for research / local use); it is # just not surfaced in the dataset dropdown. Overrides: # DEMO_SHOW_HIDDEN=1 -> show everything (local research) # DEMO_HIDDEN_DATASETS=a,b,c -> replace this default list HIDDEN_DATASET_IDS = frozenset({"sounds_colors_batch"}) @classmethod def _resolve_hidden_ids(cls) -> Set[str]: if (os.environ.get("DEMO_SHOW_HIDDEN") or "").strip().lower() in ( "1", "true", "yes", ): return set() env = (os.environ.get("DEMO_HIDDEN_DATASETS") or "").strip() if env: return {s.strip() for s in env.split(",") if s.strip()} return set(cls.HIDDEN_DATASET_IDS) def __init__(self, output_root: Path, initial_data_dir: Optional[Path] = None): self.output_root = Path(output_root) # {dataset_id: {"dir": Path, "label": str, "runs": [run_dict, ...]}} self._datasets: Dict[str, Dict[str, Any]] = {} self._active_dataset_id: Optional[str] = None self.active_loader: Optional["DataLoader"] = None self._best_aggregator = None # lazy init after scan # Cached entity-slug intersection per dataset. A slug only stays # if every dropdown-visible run has a ``by_source//`` dir # with at least one ``to_*.json`` swap. Used to propagate manual # entity deletions (e.g. flawed targets) across all runs. self._allowed_slugs_cache: Dict[str, Set[str]] = {} self._scan() if self._datasets: # If a specific data_dir was given, prefer the dataset that matches it preferred_id = self._resolve_preferred_dataset_id(initial_data_dir) self._activate_dataset(preferred_id) def _resolve_preferred_dataset_id(self, initial_data_dir: Optional[Path]) -> str: """Pick a deterministic default dataset for the demo.""" if initial_data_dir is not None: initial_data_dir = Path(initial_data_dir).resolve() for ds_id, ds in self._datasets.items(): if ds["dir"].resolve() == initial_data_dir: return ds_id env_dataset_id = (os.environ.get("DEMO_DEFAULT_DATASET") or "").strip() if env_dataset_id and env_dataset_id in self._datasets: return env_dataset_id if self.DEFAULT_DATASET_ID in self._datasets: return self.DEFAULT_DATASET_ID # Fall back to the first dataset after label sort. return next(iter(self._datasets)) # ------------------------------------------------------------------ # Discovery # ------------------------------------------------------------------ def _scan(self) -> None: """Walk output_root looking for demo-enabled run_manifest.json files.""" if not self.output_root.exists(): return datasets: Dict[Path, Dict[str, Any]] = {} for manifest_path in sorted( self.output_root.glob("*/_swaps/runs/*/run_manifest.json"), reverse=True, ): try: with open(manifest_path, "r", encoding="utf-8") as fh: manifest = json.load(fh) except (json.JSONDecodeError, IOError): continue if not manifest.get("display_demo"): continue # Resolve dataset dir: prefer the explicit graphs_root in config graphs_root_raw = ( manifest.get("config", {}) .get("inputs", {}) .get("graphs_root", "") ) dataset_dir = None if graphs_root_raw: candidate = Path(graphs_root_raw) if not candidate.is_absolute(): candidate = self.output_root.parent / candidate resolved_root = self.output_root.resolve() if candidate.exists() and _is_subpath(candidate.resolve(), resolved_root): dataset_dir = candidate if dataset_dir is None: # Physical fallback: output//_swaps/runs//run_manifest.json dataset_dir = manifest_path.parent.parent.parent.parent dataset_dir = dataset_dir.resolve() if dataset_dir not in datasets: label = self._derive_dataset_label(dataset_dir.name) datasets[dataset_dir] = { "dir": dataset_dir, "label": label, "runs": [], } run_id = manifest_path.parent.name swap_count = sum( 1 for _ in (manifest_path.parent / "by_source").glob("*/to_*.json") ) if (manifest_path.parent / "by_source").exists() else 0 run_dir = manifest_path.parent datasets[dataset_dir]["runs"].append({ "id": run_id, "manifest": manifest, "swap_count": swap_count, "control_mode": _extract_control_mode(manifest, run_dir=run_dir), "semantic_label": _semantic_run_label(run_id, manifest, run_dir=run_dir), }) # Build ordered dict: datasets sorted by label, runs already reverse-sorted hidden = self._resolve_hidden_ids() for dataset_dir, ds in sorted(datasets.items(), key=lambda x: x[1]["label"]): ds_id = dataset_dir.name if ds_id in hidden: continue self._datasets[ds_id] = ds @staticmethod def _derive_dataset_label(folder_name: str) -> str: """Turn 'book_characters_authors_batch' -> 'Book Characters Authors'.""" name = folder_name.replace("_batch", "").replace("_", " ").title() return name # ------------------------------------------------------------------ # Public API # ------------------------------------------------------------------ def list_datasets(self) -> List[Dict[str, Any]]: """Return summary list of all demo datasets.""" result = [] for ds_id, ds in self._datasets.items(): result.append({ "id": ds_id, "label": ds["label"], "run_count": len(ds["runs"]), "is_active": ds_id == self._active_dataset_id, }) return result def list_runs_for_dataset(self, dataset_id: str) -> List[Dict[str, Any]]: """Return run summaries for a specific dataset (filtered for the UI).""" ds = self._datasets.get(dataset_id) if not ds: return [] result = [] for r in ds["runs"]: manifest = r["manifest"] if not _is_dropdown_run(r["id"], manifest): continue result.append({ "id": r["id"], "dataset_id": dataset_id, "dataset_label": ds["label"], "name": self._format_run_name(r["id"]), "semantic_label": r.get("semantic_label", ""), "control_mode": r.get("control_mode", "labeled"), "swap_count": r["swap_count"], "status": manifest.get("status", "unknown"), "timestamp": manifest.get("timestamp_started", manifest.get("timestamp", "")), "has_trajectory": ( (ds["dir"] / "_swaps" / "runs" / r["id"] / "_trajectory_analysis").exists() ), }) return result def list_all_runs(self) -> List[Dict[str, Any]]: """Return all demo-enabled runs across all datasets, grouped by dataset.""" all_runs = [] for ds_id in self._datasets: all_runs.extend(self.list_runs_for_dataset(ds_id)) return all_runs def set_dataset(self, dataset_id: str) -> bool: """Switch the active dataset (picks its best run automatically).""" if dataset_id not in self._datasets: return False self._activate_dataset(dataset_id) return True def set_run(self, run_id: str) -> bool: """Switch run within the current dataset.""" if self.active_loader is None: return False return self.active_loader.set_run(run_id) def set_run_global(self, run_id: str) -> bool: """Switch to any run across all datasets. Finds which dataset owns the run, activates that dataset, then selects the specific run. """ for ds_id, ds in self._datasets.items(): for r in ds["runs"]: if r["id"] == run_id: if ds_id != self._active_dataset_id: self._activate_dataset(ds_id) self.active_loader.set_run(run_id) return True return False @property def active_dataset_id(self) -> Optional[str]: return self._active_dataset_id @property def best_aggregator(self): """Lazy-init cross-run best aggregator.""" if self._best_aggregator is None: from .cross_run_best import CrossRunBestAggregator self._best_aggregator = CrossRunBestAggregator(self) return self._best_aggregator def get_best_cross_run_matrix(self) -> Dict[str, Any]: """Return best-per-cell matrix + winners for the active dataset.""" if not self._active_dataset_id: return {"matrix": {}, "winners": {}, "considered_runs": []} return self.best_aggregator.get_best_matrix(self._active_dataset_id) def get_allowed_slugs(self, dataset_id: str) -> Optional[Set[str]]: """Return the entity-slug intersection across visible runs. A slug is "allowed" only if every dropdown-visible run for the dataset has a non-empty ``by_source//`` directory. This propagates manual entity deletions in any one run to all the others (the user removes flawed targets in Field Additivity and the Classic Swap matrix automatically hides the same rows / columns). Returns ``None`` when the dataset is unknown or has no visible runs (so the caller should fall back to no filtering). """ if dataset_id in self._allowed_slugs_cache: return self._allowed_slugs_cache[dataset_id] ds = self._datasets.get(dataset_id) if not ds: return None visible_runs = [ r for r in ds["runs"] if _is_dropdown_run(r["id"], r.get("manifest", {})) ] if not visible_runs: return None per_run_slugs: List[Set[str]] = [] for r in visible_runs: run_dir = ds["dir"] / "_swaps" / "runs" / r["id"] slugs = self._collect_run_slugs(run_dir) if slugs: per_run_slugs.append(slugs) if not per_run_slugs: return None allowed = set.intersection(*per_run_slugs) self._allowed_slugs_cache[dataset_id] = allowed return allowed @staticmethod def _collect_run_slugs(run_dir: Path) -> Set[str]: """Return the set of entity slugs that have at least one swap file in ``run_dir/by_source//to_*.json`` (canonical or variant). An empty source folder (because every ``to_*.json`` was deleted) is treated as "entity removed" and excluded. """ by_source = run_dir / "by_source" if not by_source.exists(): return set() slugs: Set[str] = set() for source_dir in by_source.iterdir(): if not source_dir.is_dir(): continue if any(source_dir.glob("to_*.json")): slugs.add(source_dir.name) return slugs def get_swap_detail_for_run( self, run_id: str, from_slug: str, to_slug: str, variant: Optional[str] = None, ) -> Optional[Dict]: """Load a swap detail from a specific run (for cross-run best mode).""" for ds in self._datasets.values(): for r in ds["runs"]: if r["id"] == run_id: loader = DataLoader(ds["dir"], run_id=run_id) return loader.get_swap_detail( from_slug, to_slug, variant=variant ) return None # ------------------------------------------------------------------ # Internal helpers # ------------------------------------------------------------------ def _activate_dataset(self, dataset_id: str) -> None: ds = self._datasets[dataset_id] best_run_id = self._pick_default_run(ds["runs"]) allowed = self.get_allowed_slugs(dataset_id) self.active_loader = DataLoader( ds["dir"], run_id=best_run_id, allowed_slugs=allowed, ) self._active_dataset_id = dataset_id @classmethod def _pick_default_run(cls, runs: List[Dict[str, Any]]) -> Optional[str]: """Choose the default run for a dataset. Priority: 1. ``DEMO_DEFAULT_RUN`` env var (exact run id match) 2. Class constant ``DEFAULT_RUN_ID`` 3. Largest "classic" swap run that's visible in the dropdown 4. Largest visible run regardless of mode 5. First run in list (newest) """ if not runs: return None env_run = (os.environ.get("DEMO_DEFAULT_RUN") or "").strip() if env_run: for r in runs: if r["id"] == env_run: return r["id"] if cls.DEFAULT_RUN_ID: for r in runs: if r["id"] == cls.DEFAULT_RUN_ID: return r["id"] visible = [ r for r in runs if _is_dropdown_run(r["id"], r.get("manifest", {})) ] if visible: classic = [ r for r in visible if r.get("control_mode") not in ("random_feature_matched", "additivity") ] pool = classic or visible return max(pool, key=lambda r: r.get("swap_count", 0))["id"] return runs[0]["id"] @staticmethod def _format_run_name(run_id: str) -> str: name = run_id.replace("_", " ").title() name = name.replace("50states", "50 States").replace("6states", "6 States") return name class DataLoader: """Load and cache swap experiment data with multi-run support.""" def __init__( self, data_dir: Path, run_id: Optional[str] = None, allowed_slugs: Optional[Set[str]] = None, ): """Initialise a loader for one dataset directory. ``allowed_slugs`` is an optional whitelist of entity slugs. When provided, every public matrix / state getter prunes its output to only those slugs. This is how cross-run entity deletions propagate: ``DemoRegistry`` computes the intersection of slugs across all visible runs and forwards it here. """ self.data_dir = Path(data_dir) self.base_swaps_dir = self.data_dir / "_swaps" self.run_id = run_id self.allowed_slugs: Optional[Set[str]] = ( set(allowed_slugs) if allowed_slugs is not None else None ) self._set_swaps_dir() self._matrix_cache: Dict[Optional[str], Dict] = {} self._flip_matrix_cache: Dict[Optional[str], Dict] = {} self._regime_matrix_cache: Dict[Optional[str], Dict] = {} self._vsmax_matrix_cache: Dict[Optional[str], Dict] = {} self._variant_index_cache: Optional[Dict] = None self._states_cache: Optional[List[Dict]] = None self._analysis_cache: Optional[Dict] = None self._stats_cache: Optional[Dict] = None self._domain_config_cache: Optional[Dict] = None self._gpu_batch_features_index: Optional[Dict[str, List]] = None def _set_swaps_dir(self): """Set swaps_dir based on current run_id.""" if self.run_id: run_dir = self.base_swaps_dir / "runs" / self.run_id if run_dir.exists(): self.swaps_dir = run_dir else: # Fallback to base swaps dir self.swaps_dir = self.base_swaps_dir else: # Check if runs directory exists and has runs runs_dir = self.base_swaps_dir / "runs" if runs_dir.exists(): # Prefer the full 50-state run on a fresh load. default_run_id = self._get_default_run_id() if default_run_id: self.run_id = default_run_id self.swaps_dir = self.base_swaps_dir / "runs" / default_run_id else: self.swaps_dir = self.base_swaps_dir else: self.swaps_dir = self.base_swaps_dir def _get_default_run_id(self) -> Optional[str]: """Choose the default run for a fresh app load (most recent non-legacy). ``list_runs`` already filters out sweeps and ``*_labeled`` controls, so we only need to drop legacy runs here. """ runs = self.list_runs() if not runs: return None for run in runs: if "legacy" not in run["id"].lower(): return run["id"] return runs[0]["id"] def list_runs(self) -> List[Dict[str, Any]]: """List available experiment runs.""" runs = [] runs_dir = self.base_swaps_dir / "runs" if not runs_dir.exists(): return runs for run_dir in sorted(runs_dir.iterdir(), reverse=True): if not run_dir.is_dir(): continue by_source = run_dir / "by_source" if not by_source.exists(): continue # Load run manifest for metadata manifest_path = run_dir / "run_manifest.json" manifest = {} if manifest_path.exists(): try: with open(manifest_path, 'r', encoding='utf-8') as f: manifest = json.load(f) except (json.JSONDecodeError, IOError): pass if not _is_dropdown_run(run_dir.name, manifest): continue # Count swaps swap_count = sum(1 for _ in by_source.glob("*/to_*.json")) # Check for trajectory data has_trajectory = (run_dir / "_trajectory_analysis").exists() # Get config info config = manifest.get('config', {}) ct_config = config.get('ct_steering', {}) runs.append({ 'id': run_dir.name, 'name': self._format_run_name(run_dir.name), 'semantic_label': _semantic_run_label(run_dir.name, manifest, run_dir=run_dir), 'control_mode': _extract_control_mode(manifest, run_dir=run_dir), 'swap_count': swap_count, 'has_trajectory': has_trajectory, 'timestamp': manifest.get('timestamp_started', ''), 'status': manifest.get('status', 'unknown'), 'M_ablate': ct_config.get('M_ablate', -2), 'M_amplify': ct_config.get('M_amplify', 20), 'is_current': run_dir.name == self.run_id, }) return runs def _format_run_name(self, run_id: str) -> str: """Format run ID into human-readable name.""" # full_50states_v1 -> "Full 50 States v1" # pilot_6states_v1 -> "Pilot 6 States v1" # gpu_test_8 -> "GPU Test 8" name = run_id.replace('_', ' ').title() name = name.replace('50states', '50 States') name = name.replace('6states', '6 States') name = name.replace('Gpu', 'GPU') return name def set_run(self, run_id: str) -> bool: """Switch to a different run.""" run_dir = self.base_swaps_dir / "runs" / run_id if not run_dir.exists(): return False self.run_id = run_id self._set_swaps_dir() self._clear_caches() return True def get_current_run(self) -> Optional[str]: """Get the current run ID.""" return self.run_id def _clear_caches(self): """Clear all cached data.""" self._matrix_cache = {} self._flip_matrix_cache = {} self._regime_matrix_cache = {} self._vsmax_matrix_cache = {} self._variant_index_cache = None self._states_cache = None self._analysis_cache = None self._stats_cache = None self._domain_config_cache = None self._gpu_batch_features_index = None self._methodology_stats = None # ------------------------------------------------------------------ # Cross-run entity filtering # ------------------------------------------------------------------ def _apply_slug_filter(self, matrix: Dict[str, Dict]) -> Dict[str, Dict]: """Return a copy of *matrix* containing only allowed_slugs. Filters both the row dimension (sources) and column dimension (targets). When ``allowed_slugs`` is None this is a no-op (the original dict is returned untouched, since callers don't mutate the result). """ if self.allowed_slugs is None: return matrix allowed = self.allowed_slugs filtered: Dict[str, Dict] = {} for from_slug, row in matrix.items(): if from_slug not in allowed: continue filtered[from_slug] = { to_slug: val for to_slug, val in row.items() if to_slug in allowed } return filtered # Per-domain blacklists of common words that should be ignored by the # word-boundary tier heuristic (e.g. "city" in "Salt Lake City" would # otherwise match generic geography text like "is a city"). TIER_WORD_BLACKLIST: Dict[str, List[str]] = { 'usa_states': ['city'], } # ------------------------------------------------------------------ # Domain & model configuration (loaded from config_resolved.json) # ------------------------------------------------------------------ def get_domain_config(self) -> Dict[str, Any]: """Return domain and model metadata for the active run. Reads config_resolved.json once and caches the result. Falls back to USA-states defaults when the file is missing (legacy runs). """ if self._domain_config_cache is not None: return self._domain_config_cache config = self._load_config_resolved() if not config: self._domain_config_cache = self._build_fallback_domain_config() return self._domain_config_cache entities = config.get('_entities', []) entity_index = {e['slug']: e for e in entities} concept_fields = config.get('swap', {}).get('concept_fields', []) entity_fields: List[str] = [] if entities: entity_fields = [k for k in entities[0] if k != 'slug'] primary_field = entity_fields[0] if entity_fields else 'slug' for f in entity_fields: if f not in concept_fields: primary_field = f break answer_field = concept_fields[-1] if concept_fields else '' ct = config.get('ct_steering', {}) model_id = ct.get('model_id', '') transcoder_set = ct.get('transcoder_set', '') np_model = model_id.split('/')[-1] if model_id else 'gemma-2-2b' experiment_name = config.get('experiment_name', '') is_usa_states = 'usa_states' in experiment_name or 'state' in concept_fields display_name = experiment_name.replace('_swap', '').replace('_', ' ').title() blacklist: List[str] = [] for key, words in self.TIER_WORD_BLACKLIST.items(): if key in experiment_name: blacklist = [w.lower() for w in words] break self._domain_config_cache = { 'experiment_name': experiment_name, 'display_name': display_name, 'concept_fields': concept_fields, 'answer_field': answer_field, 'primary_field': primary_field, 'entity_fields': entity_fields, 'entity_index': entity_index, 'entity_count': len(entities), 'model_id': model_id, 'transcoder_set': transcoder_set, 'np_model': np_model, 'is_usa_states': is_usa_states, 'tier_word_blacklist': blacklist, 'm_ablate': ct.get('M_ablate', -2), 'm_amplify': ct.get('M_amplify', 20), 'temperature': ct.get('temperature', 0.3), 'n_tokens': ct.get('n_tokens', 10), } return self._domain_config_cache def _load_config_resolved(self) -> Optional[Dict]: """Load config_resolved.json from the active run directory.""" path = self.swaps_dir / "config_resolved.json" if path.exists(): try: with open(path, 'r', encoding='utf-8') as f: return json.load(f) except (json.JSONDecodeError, IOError): pass return None def _build_fallback_domain_config(self) -> Dict[str, Any]: """Fallback domain config for legacy runs without config_resolved.json.""" return { 'experiment_name': 'usa_states_swap', 'display_name': 'USA States', 'concept_fields': ['state', 'capital'], 'answer_field': 'capital', 'primary_field': 'city', 'entity_fields': ['city', 'state', 'capital'], 'entity_index': {}, 'entity_count': 0, 'model_id': 'google/gemma-2-2b', 'transcoder_set': 'mntss/clt-gemma-2-2b-2.5M', 'np_model': 'gemma-2-2b', 'is_usa_states': True, 'tier_word_blacklist': ['city'], 'm_ablate': -2, 'm_amplify': 20, 'temperature': 0.3, 'n_tokens': 10, } @staticmethod def _slug_to_label(slug: str) -> str: """Convert a slug to a human-readable label.""" return slug.replace('_', ' ').title() def _make_abbr(self, slug: str, label: str) -> str: """Generate a short abbreviation for matrix headers.""" dc = self.get_domain_config() if dc['is_usa_states']: return self._state_to_abbr(label) parts = slug.split('_') if len(parts) >= 2: return ''.join(p[0].upper() for p in parts[:3]) return slug[:3].upper() # ------------------------------------------------------------------ # Variant index # ------------------------------------------------------------------ def _build_variant_index(self) -> Dict: """Group swap files by canonical (from_slug, to_slug) pair. Returns ``{from_slug: {to_slug: [record, ...]}}``. Each record is:: { "path": Path, "variant_suffix": str, # "" for canonical, "r0" / "add_book" / etc. "is_canonical": bool, } """ if self._variant_index_cache is not None: return self._variant_index_cache index: Dict[str, Dict[str, list]] = {} by_source_dir = self.swaps_dir / "by_source" if not by_source_dir.exists(): self._variant_index_cache = index return index for source_dir in by_source_dir.iterdir(): if not source_dir.is_dir(): continue from_slug = source_dir.name if from_slug not in index: index[from_slug] = {} for swap_file in source_dir.glob("to_*.json"): to_slug, variant_suffix = _parse_to_slug(swap_file.stem) if not to_slug: continue index[from_slug].setdefault(to_slug, []).append({ "path": swap_file, "variant_suffix": variant_suffix, "is_canonical": variant_suffix == "", }) self._variant_index_cache = index return index def get_variants_for_pair( self, from_slug: str, to_slug: str ) -> List[Dict[str, Any]]: """Return lightweight variant descriptors for a single pair. Each item includes ``variant_suffix``, ``tier``, ``flip_position``, and control metadata extracted from the swap JSON. """ idx = self._build_variant_index() records = idx.get(from_slug, {}).get(to_slug, []) variants: List[Dict[str, Any]] = [] for rec in records: if rec["is_canonical"]: continue try: with open(rec["path"], "r", encoding="utf-8") as fh: data = json.load(fh) except (json.JSONDecodeError, IOError): continue ctrl = data.get("metadata", {}).get("control", {}) variants.append({ "variant_suffix": rec["variant_suffix"], "tier": self._get_tier_from_swap(data), "flip_position": self._get_flip_position_from_swap(data), "control_mode": ctrl.get("control_mode", ""), "concept_subsets_used": ctrl.get("concept_subsets_used", []), "replicate_id": ctrl.get("replicate_id"), }) return variants def get_run_variant_suffixes(self) -> List[str]: """Return sorted list of distinct variant suffixes across the entire run.""" idx = self._build_variant_index() suffixes: set = set() for targets in idx.values(): for records in targets.values(): for r in records: if r["variant_suffix"]: suffixes.add(r["variant_suffix"]) return sorted(suffixes) def _best_variant_path( self, from_slug: str, to_slug: str ) -> Optional[Path]: """Return the path of the best variant file for a pair (highest tier).""" idx = self._build_variant_index() records = idx.get(from_slug, {}).get(to_slug, []) variants = [r for r in records if not r["is_canonical"]] if not variants: return None best_path = None best_tier = -1 for rec in variants: try: with open(rec["path"], "r", encoding="utf-8") as fh: data = json.load(fh) tier = self._get_tier_from_swap(data) or 0 if tier > best_tier: best_tier = tier best_path = rec["path"] except (json.JSONDecodeError, IOError): continue return best_path # ------------------------------------------------------------------ # Matrix # ------------------------------------------------------------------ def get_matrix(self, variant: Optional[str] = None) -> Dict[str, Dict[str, Optional[int]]]: """Build tier matrix dynamically from individual swap JSON files. When *variant* is None (default), canonical files take priority and variant-only pairs aggregate using the best (highest) tier. When *variant* is set, only files matching that suffix are used. """ if variant in self._matrix_cache: return self._apply_slug_filter(self._matrix_cache[variant]) idx = self._build_variant_index() matrix: Dict[str, Dict[str, Optional[int]]] = {} for from_slug, targets in idx.items(): matrix[from_slug] = {} for to_slug, records in targets.items(): if variant is not None: matched = [r for r in records if r["variant_suffix"] == variant] if matched: try: with open(matched[0]["path"], "r", encoding="utf-8") as f: swap_data = json.load(f) matrix[from_slug][to_slug] = self._get_tier_from_swap(swap_data) except (json.JSONDecodeError, IOError): pass continue canonical = [r for r in records if r["is_canonical"]] variants = [r for r in records if not r["is_canonical"]] if canonical: try: with open(canonical[0]["path"], "r", encoding="utf-8") as f: swap_data = json.load(f) matrix[from_slug][to_slug] = self._get_tier_from_swap(swap_data) except (json.JSONDecodeError, IOError): pass elif variants: best_tier: Optional[float] = None for rec in variants: try: with open(rec["path"], "r", encoding="utf-8") as f: swap_data = json.load(f) tier = self._get_tier_from_swap(swap_data) if tier is not None and (best_tier is None or tier > best_tier): best_tier = tier except (json.JSONDecodeError, IOError): continue if best_tier is not None: matrix[from_slug][to_slug] = best_tier if variant is None: work_dir = self.swaps_dir / "work" if work_dir.exists(): for swap_file in work_dir.glob("*.json"): try: with open(swap_file, "r", encoding="utf-8") as f: swap_data = json.load(f) swap_id = swap_file.stem if "__to__" in swap_id: from_slug, to_slug = swap_id.split("__to__") if from_slug not in matrix: matrix[from_slug] = {} if to_slug not in matrix[from_slug]: matrix[from_slug][to_slug] = self._get_tier_from_swap(swap_data) except (json.JSONDecodeError, IOError): continue self._matrix_cache[variant] = matrix return self._apply_slug_filter(matrix) def _get_tier_from_swap(self, swap_data: Dict) -> Optional[float]: """Extract tier from swap data, computing if necessary. Tier semantics (domain-agnostic): 5 = PERFECT -- target answer detected in steered output via any of: a) full answer match (strict or fuzzy punctuation) b) steered first token is a substring of the target answer c) any word (len>=3) of the target answer appears in the steered output (e.g. "Harper" in "Nelle Harper...") 2 = SUPPRESSED_ONLY -- source gone, no target signal 1 = SOURCE_PERSISTS -- source answer still in output """ if 'classification' in swap_data: tier = swap_data['classification'].get('tier') if tier is not None: return float(tier) evaluation = swap_data.get('evaluation', {}) exact = evaluation.get('exact_match', {}) hit = exact.get('steered_has_to_capital') or exact.get('steered_has_to_answer') to_answer = evaluation.get('to_answer', '') steered_out = evaluation.get('raw', {}).get('steered_output', '') if not hit and to_answer and steered_out: to_norm = to_answer.replace('.', '').replace('-', ' ').lower() out_norm = steered_out.replace('.', '').replace('-', ' ').lower() if to_norm and to_norm in out_norm: hit = True if not hit: first_token = evaluation.get('first_token', {}) steered_first_tok = (first_token.get('steered', '') or '').strip() if len(steered_first_tok) >= 2 and to_answer: answer_norm = to_answer.replace('.', '').lower() tok_norm = steered_first_tok.lower() if tok_norm in answer_norm: hit = True if not hit and to_answer and steered_out: blacklist = set(self.get_domain_config().get('tier_word_blacklist', [])) out_lower = steered_out.lower() for word in to_answer.replace('.', '').split(): if len(word) >= 3 and word.lower() not in blacklist: pattern = r'\b' + re.escape(word.lower()) + r'\b' if re.search(pattern, out_lower): hit = True break if hit: return 5 if not exact.get('from_suppressed'): return 1 return 2 @staticmethod def _get_flip_position_from_swap(swap_data: Dict): """Extract logit flip position from swap data. Returns: int >= 0 -- generation position where the flip first occurs None -- trajectory was tracked but flip never occurred -1 -- no trajectory data available for this swap """ traj = swap_data.get('evaluation', {}).get('logit_trajectory') if traj is None or not isinstance(traj, dict) or 'summary' not in traj: return -1 return traj['summary'].get('flip_position') def get_flip_matrix(self, variant: Optional[str] = None) -> Dict[str, Dict[str, Optional[int]]]: """Build flip-position matrix from swap JSON files. Same structure as ``get_matrix()`` but values are the generation position where the logit flip first occurs (0, 1, 2 ...) or ``None`` when no flip is observed. When *variant* is None, canonical files take priority and variant-only pairs pick the earliest non-null flip. When *variant* is set, only files matching that suffix are used. """ if variant in self._flip_matrix_cache: return self._apply_slug_filter(self._flip_matrix_cache[variant]) idx = self._build_variant_index() matrix: Dict[str, Dict[str, Optional[int]]] = {} for from_slug, targets in idx.items(): matrix[from_slug] = {} for to_slug, records in targets.items(): if variant is not None: matched = [r for r in records if r["variant_suffix"] == variant] if matched: try: with open(matched[0]["path"], "r", encoding="utf-8") as f: swap_data = json.load(f) matrix[from_slug][to_slug] = self._get_flip_position_from_swap(swap_data) except (json.JSONDecodeError, IOError): pass continue canonical = [r for r in records if r["is_canonical"]] variants = [r for r in records if not r["is_canonical"]] if canonical: try: with open(canonical[0]["path"], "r", encoding="utf-8") as f: swap_data = json.load(f) matrix[from_slug][to_slug] = self._get_flip_position_from_swap(swap_data) except (json.JSONDecodeError, IOError): pass elif variants: best_flip = None for rec in variants: try: with open(rec["path"], "r", encoding="utf-8") as f: swap_data = json.load(f) fp = self._get_flip_position_from_swap(swap_data) if fp is not None and fp >= 0: if best_flip is None or fp < best_flip: best_flip = fp except (json.JSONDecodeError, IOError): continue matrix[from_slug][to_slug] = best_flip if variant is None: work_dir = self.swaps_dir / "work" if work_dir.exists(): for swap_file in work_dir.glob("*.json"): try: with open(swap_file, "r", encoding="utf-8") as f: swap_data = json.load(f) swap_id = swap_file.stem if "__to__" in swap_id: from_slug, to_slug = swap_id.split("__to__") if from_slug not in matrix: matrix[from_slug] = {} if to_slug not in matrix[from_slug]: matrix[from_slug][to_slug] = self._get_flip_position_from_swap(swap_data) except (json.JSONDecodeError, IOError): continue self._flip_matrix_cache[variant] = matrix return self._apply_slug_filter(matrix) # ------------------------------------------------------------------ # Regime & VsMax matrices # ------------------------------------------------------------------ @staticmethod def _get_regime_from_swap(swap_data: Dict) -> Optional[str]: """Extract logit-shift regime (A-E) from swap data.""" pos0 = swap_data.get('evaluation', {}).get('position_0_comparison', {}) return _classify_regime(pos0) @staticmethod def _get_vsmax_from_swap(swap_data: Dict) -> Optional[float]: """Extract best-trajectory vsMax from swap data. Returns the ``best_target_minus_max`` float or None. """ contrast = (swap_data.get('evaluation', {}) .get('logit_trajectory', {}) .get('contrast_groups', {})) if not isinstance(contrast, dict): return None grp = contrast.get('same_dataset', {}) if not isinstance(grp, dict): return None agg = grp.get('aggregate', {}) if not isinstance(agg, dict): return None return agg.get('best_target_minus_max') def _build_overlay_matrix( self, extractor, cache: Dict, variant: Optional[str], best_cmp, ) -> Dict[str, Dict]: """Generic matrix builder reused by regime and vsmax. ``extractor`` pulls a value from swap_data. ``best_cmp(current_best, candidate)`` returns True when candidate is preferred over current_best (used for variant-only fallback). """ if variant in cache: return self._apply_slug_filter(cache[variant]) idx = self._build_variant_index() matrix: Dict[str, Dict] = {} for from_slug, targets in idx.items(): matrix[from_slug] = {} for to_slug, records in targets.items(): if variant is not None: matched = [r for r in records if r["variant_suffix"] == variant] if matched: try: with open(matched[0]["path"], "r", encoding="utf-8") as f: swap_data = json.load(f) matrix[from_slug][to_slug] = extractor(swap_data) except (json.JSONDecodeError, IOError): pass continue canonical = [r for r in records if r["is_canonical"]] variants_list = [r for r in records if not r["is_canonical"]] if canonical: try: with open(canonical[0]["path"], "r", encoding="utf-8") as f: swap_data = json.load(f) matrix[from_slug][to_slug] = extractor(swap_data) except (json.JSONDecodeError, IOError): pass elif variants_list: best_val = None for rec in variants_list: try: with open(rec["path"], "r", encoding="utf-8") as f: swap_data = json.load(f) val = extractor(swap_data) if val is not None and (best_val is None or best_cmp(best_val, val)): best_val = val except (json.JSONDecodeError, IOError): continue if best_val is not None: matrix[from_slug][to_slug] = best_val if variant is None: work_dir = self.swaps_dir / "work" if work_dir.exists(): for swap_file in work_dir.glob("*.json"): try: with open(swap_file, "r", encoding="utf-8") as f: swap_data = json.load(f) swap_id = swap_file.stem if "__to__" in swap_id: from_slug, to_slug = swap_id.split("__to__") if from_slug not in matrix: matrix[from_slug] = {} if to_slug not in matrix[from_slug]: matrix[from_slug][to_slug] = extractor(swap_data) except (json.JSONDecodeError, IOError): continue cache[variant] = matrix return self._apply_slug_filter(matrix) def get_regime_matrix(self, variant: Optional[str] = None) -> Dict[str, Dict[str, Optional[str]]]: """Build regime matrix (A-E letters) from swap JSON files.""" _REGIME_RANK = {"A": 5, "B": 4, "C": 3, "D": 2, "E": 1} return self._build_overlay_matrix( self._get_regime_from_swap, self._regime_matrix_cache, variant, best_cmp=lambda cur, cand: _REGIME_RANK.get(cand, 0) > _REGIME_RANK.get(cur, 0), ) def get_vsmax_matrix(self, variant: Optional[str] = None) -> Dict[str, Dict[str, Optional[float]]]: """Build vsMax matrix from swap JSON files. Higher is better.""" return self._build_overlay_matrix( self._get_vsmax_from_swap, self._vsmax_matrix_cache, variant, best_cmp=lambda cur, cand: cand > cur, ) def get_states(self) -> List[Dict[str, Any]]: """Get entity list with metadata from by_source directories. Returns a list of dicts with backward-compatible keys (``state``, ``capital``, ``city``, ``abbr``) plus generic keys (``label``, ``answer``, ``fields``) so that frontends can work with any domain. """ if self._states_cache is not None: return self._states_cache dc = self.get_domain_config() entity_index = dc['entity_index'] primary_field = dc['primary_field'] answer_field = dc['answer_field'] analysis = self.get_analysis() archetypes = analysis.get('archetypes', {}) archetype_lookup: Dict[str, Dict] = {} for archetype, entity_list in archetypes.items(): for ed in entity_list: key = ed.get('state', ed.get(primary_field, '')) archetype_lookup[key] = { 'archetype': archetype, 'native_prob': ed.get('native_prob', 0), 'supernodes': ed.get('supernodes', 0), 'src_tier': ed.get('src_tier', 0), 'tgt_tier': ed.get('tgt_tier', 0), } entities: List[Dict[str, Any]] = [] by_source_dir = self.swaps_dir / "by_source" if by_source_dir.exists(): for source_dir in sorted(by_source_dir.iterdir()): if not source_dir.is_dir(): continue slug = source_dir.name if self.allowed_slugs is not None and slug not in self.allowed_slugs: continue edata = entity_index.get(slug, {}) if not edata: edata = self._get_entity_info_from_swap(source_dir) label = (edata.get(primary_field, '') or self._slug_to_label(slug)) answer = edata.get(answer_field, '') abbr = self._make_abbr(slug, edata.get('state', label)) # Backward-compatible keys state_val = edata.get('state', label) capital_val = edata.get('capital', answer) city_val = edata.get('city', edata.get(primary_field, label)) arch_key = state_val if dc['is_usa_states'] else label arch_data = archetype_lookup.get(arch_key, {}) neuronpedia_url = self._get_neuronpedia_url(slug) logit_flags = (self._get_entity_logit_flags(slug, answer) if self._find_state_dir(slug) else {}) entities.append({ 'slug': slug, 'label': label, 'answer': answer, 'state': state_val, 'capital': capital_val, 'city': city_val, 'abbr': abbr, 'archetype': arch_data.get('archetype', 'Unknown'), 'native_prob': arch_data.get('native_prob', 0), 'supernodes': arch_data.get('supernodes', 0), 'src_tier': arch_data.get('src_tier', 0), 'tgt_tier': arch_data.get('tgt_tier', 0), 'neuronpedia_url': neuronpedia_url, 'fields': {k: v for k, v in edata.items() if k != 'slug'}, **logit_flags, }) self._states_cache = entities return entities def _get_entity_info_from_swap(self, source_dir: Path) -> Dict[str, str]: """Extract entity fields from a sample swap file's source section. Tries canonical files first, falls back to variant files for variant-only runs. """ for swap_file in sorted(source_dir.glob("to_*.json")): try: with open(swap_file, 'r', encoding='utf-8') as f: swap_data = json.load(f) source = swap_data.get('source', {}) if source: return {k: v for k, v in source.items() if k not in ('prompt', 'concept', 'neuronpedia_url')} except (json.JSONDecodeError, IOError): continue return {} def _load_entity_logit_metadata(self, entity_dir: Path, answer: str) -> Dict[str, Any]: """Load top-logit metadata from graph.json. ``answer`` is the expected output token (capital, author, etc.) used only for highlighting; the rest is model-agnostic. """ graph_path = entity_dir / "00 Graph Generation" / "graph.json" if not graph_path.exists(): return {} try: with open(graph_path, 'r', encoding='utf-8') as f: graph = json.load(f) except (json.JSONDecodeError, IOError): return {} import re logits: List[Dict] = [] native_prob = 0.0 target_token = '' for node in graph.get('nodes', []): clerp = node.get('clerp', '') prob = node.get('token_prob') if clerp and prob is not None and clerp.startswith('Output '): match = re.search(r'Output\s+"([^"]*)"', clerp) if match: token = match.group(1).strip() if token: logits.append({ 'token': token, 'prob': prob, 'is_target': node.get('is_target_logit', False), }) if node.get('is_target_logit') and not target_token: native_prob = node.get('token_prob', 0) match = re.search(r'Output\s+"([^"]*)"', clerp) target_token = (match.group(1).strip() if match else clerp.strip().strip("'")) logits.sort(key=lambda x: x['prob'], reverse=True) top_logits = logits[:5] answer_lower = answer.lower() if answer else '' answer_in_logits = any( t and len(t) > 1 and t.lower() in answer_lower for t in (l.get('token', '') for l in top_logits) ) return { 'logits': top_logits, 'native_prob': native_prob, 'target_token': target_token, 'capital_is_top_logit': ( target_token.lower() in answer_lower if answer and target_token else False), 'capital_in_logits': answer_in_logits, } def _get_entity_logit_flags(self, slug: str, answer: str) -> Dict[str, Any]: """Return compact warning flags for the matrix entity list.""" entity_dir = self._find_state_dir(slug) if not entity_dir: return {} metadata = self._load_entity_logit_metadata(entity_dir, answer) return { 'capital_is_top_logit': metadata.get('capital_is_top_logit'), 'capital_in_logits': metadata.get('capital_in_logits'), } def get_analysis(self) -> Dict[str, Any]: """Load analysis summary data.""" if self._analysis_cache is not None: return self._analysis_cache analysis = {} analysis_dir = self.swaps_dir / "_analysis_v3" # Load tier summary tier_summary_path = analysis_dir / "tier_summary.json" if tier_summary_path.exists(): with open(tier_summary_path, 'r', encoding='utf-8') as f: analysis['tier_summary'] = json.load(f) # Load swap factor summary factor_path = analysis_dir / "swap_factor_summary.json" if factor_path.exists(): with open(factor_path, 'r', encoding='utf-8') as f: factor_data = json.load(f) analysis['correlations'] = factor_data.get('correlations', {}) analysis['archetypes'] = factor_data.get('archetypes', {}) analysis['insights'] = factor_data.get('insights', []) analysis['anomalies'] = factor_data.get('anomalies', {}) self._analysis_cache = analysis return analysis def get_swap_detail( self, from_slug: str, to_slug: str, variant: Optional[str] = None, ) -> Optional[Dict[str, Any]]: """Load detailed swap result with variant resolution. Resolution order: 1. Explicit *variant* requested -> load that variant file. 2. Canonical file exists -> load it. 3. Otherwise -> load the best variant for the pair. """ data = None swap_path = None if variant: swap_path = ( self.swaps_dir / "by_source" / from_slug / f"to_{to_slug}__{variant}.json" ) else: canonical = self.swaps_dir / "by_source" / from_slug / f"to_{to_slug}.json" if canonical.exists(): swap_path = canonical else: swap_path = self._best_variant_path(from_slug, to_slug) if swap_path is None or not swap_path.exists(): work_path = self.swaps_dir / "work" / f"{from_slug}__to__{to_slug}.json" if work_path.exists(): swap_path = work_path if swap_path is None or not swap_path.exists(): return None with open(swap_path, "r", encoding="utf-8") as f: data = json.load(f) # Enrich source/target with per-entity metadata data['source']['neuronpedia_url'] = self._get_neuronpedia_url(from_slug) data['target']['neuronpedia_url'] = self._get_neuronpedia_url(to_slug) for key, slug in (('source', from_slug), ('target', to_slug)): pct = self._get_error_node_influence(slug) if pct is not None: data[key]['error_node_influence_pct'] = pct # Inject derived metrics block (additive, null-safe) ctrl = data.get("metadata", {}).get("control", {}) traj_summary = data.get("evaluation", {}).get("logit_trajectory", {}).get("summary", {}) contrast = data.get("evaluation", {}).get("logit_trajectory", {}).get("contrast_groups", {}) pos0 = data.get("evaluation", {}).get("position_0_comparison", {}) regime = _classify_regime(pos0) regime_info = _REGIME_INFO.get(regime) if regime else None data["_derived"] = { "control_mode": ctrl.get("control_mode"), "fields_used": ctrl.get("concept_subsets_used"), "replicate_id": ctrl.get("replicate_id"), "flip_position": traj_summary.get("flip_position"), "gap_closure": traj_summary.get("gap_closure"), "initial_gap": traj_summary.get("initial_gap"), "best_gap": traj_summary.get("best_gap"), "regime": regime, "regime_label": regime_info["label"] if regime_info else None, "target_logit_delta": pos0.get("target_logit_delta"), "source_logit_delta": pos0.get("source_logit_delta"), "vs_max": None, "vs_topk": None, "rank_in_group": None, "vs_max_pos0": None, "vs_topk_pos0": None, "rank_in_group_pos0": None, } if isinstance(contrast, dict): grp = contrast.get("same_dataset", {}) if isinstance(grp, dict): agg = grp.get("aggregate", {}) if isinstance(grp, dict) else {} data["_derived"]["vs_max"] = agg.get("best_target_minus_max") data["_derived"]["vs_topk"] = agg.get("best_target_minus_topk") data["_derived"]["rank_in_group"] = agg.get("best_rank_within") data["_derived"]["vs_max_pos0"] = agg.get("initial_target_minus_max") data["_derived"]["vs_topk_pos0"] = agg.get("initial_target_minus_topk") data["_derived"]["rank_in_group_pos0"] = agg.get("initial_rank_within") data["_derived"]["contrast_members"] = grp.get("members", []) data["_derived"]["contrast_n"] = grp.get("n_members", 0) data["_derived"]["contrast_topk_k"] = grp.get("topk_k", 3) # Variant availability summary variants = self.get_variants_for_pair(from_slug, to_slug) data["_variants"] = variants _, loaded_suffix = _parse_to_slug(swap_path.stem) data["_loaded_variant"] = loaded_suffix or None return data def save_annotation(self, from_slug: str, to_slug: str, tier: Optional[int] = None, notes: Optional[str] = None) -> Dict[str, Any]: """ Save annotation (tier/notes) to the swap JSON file. Args: from_slug: Source state slug to_slug: Target state slug tier: New tier value (1-5), or None to keep existing notes: Annotation notes, or None to keep existing Returns: Updated swap data with new stats """ # Find the swap file swap_path = self.swaps_dir / "by_source" / from_slug / f"to_{to_slug}.json" if not swap_path.exists(): swap_path = self.swaps_dir / "work" / f"{from_slug}__to__{to_slug}.json" if not swap_path.exists(): raise FileNotFoundError(f"Swap file not found: {from_slug} -> {to_slug}") # Load existing data with open(swap_path, 'r', encoding='utf-8') as f: data = json.load(f) # Ensure classification section exists if 'classification' not in data: data['classification'] = {} # Update tier if provided if tier is not None: old_tier = data['classification'].get('tier') data['classification']['tier'] = tier data['classification']['manually_edited'] = True # Update notes if provided if notes is not None: data['classification']['notes'] = notes data['classification']['manually_edited'] = True # Save back to file with open(swap_path, 'w', encoding='utf-8') as f: json.dump(data, f, indent=2, ensure_ascii=False) # Clear caches so stats are recalculated self._matrix_cache = {} self._flip_matrix_cache = {} self._regime_matrix_cache = {} self._vsmax_matrix_cache = {} self._stats_cache = None # Return updated data with fresh stats return { 'swap': data, 'stats': self.get_stats(), 'matrix_update': { 'from': from_slug, 'to': to_slug, 'tier': data['classification'].get('tier'), 'manually_edited': True, } } def get_annotated_swaps(self) -> List[Dict[str, str]]: """Get list of all manually annotated swaps.""" annotated = [] by_source_dir = self.swaps_dir / "by_source" if by_source_dir.exists(): for source_dir in by_source_dir.iterdir(): if not source_dir.is_dir(): continue from_slug = source_dir.name for swap_file in source_dir.glob("to_*.json"): if _is_control_variant(swap_file.stem): continue try: with open(swap_file, 'r', encoding='utf-8') as f: data = json.load(f) if data.get('classification', {}).get('manually_edited'): to_slug, _ = _parse_to_slug(swap_file.stem) annotated.append({ 'from': from_slug, 'to': to_slug, 'tier': data['classification'].get('tier'), 'notes': data['classification'].get('notes', ''), }) except (json.JSONDecodeError, IOError): continue return annotated def get_stats(self) -> Dict[str, Any]: """Get aggregate statistics - computed dynamically from matrix.""" if self._stats_cache is not None: return self._stats_cache matrix = self.get_matrix() analysis = self.get_analysis() # Count tiers from actual matrix data tier_counts = { 'PERFECT': 0, # 5 'TARGET_STATE_CITY': 0, # 4 'TARGET_STATE_ONLY': 0, # 3 'SUPPRESSED_ONLY': 0, # 2 'SOURCE_PERSISTS': 0, # 1 } tier_names = {5: 'PERFECT', 4: 'TARGET_STATE_CITY', 3: 'TARGET_STATE_ONLY', 2: 'SUPPRESSED_ONLY', 1: 'SOURCE_PERSISTS'} total_swaps = 0 tier_sum = 0 perfect_count = 0 state_correct_count = 0 # T3+ suppressed_count = 0 # T2+ for from_slug, targets in matrix.items(): for to_slug, tier in targets.items(): if tier is not None and from_slug != to_slug: total_swaps += 1 tier_sum += tier tier_name = tier_names.get(tier, 'UNKNOWN') if tier_name in tier_counts: tier_counts[tier_name] += 1 if tier == 5: perfect_count += 1 if tier >= 3: state_correct_count += 1 if tier >= 2: suppressed_count += 1 # Calculate rates tier_rates = {} for name, count in tier_counts.items(): tier_rates[name] = count / total_swaps if total_swaps > 0 else 0 # Flip stats from trajectory data flip_matrix = self.get_flip_matrix() flip_tracked = 0 flip_at_0 = 0 flip_at_01 = 0 for from_slug, targets in flip_matrix.items(): for to_slug, fp in targets.items(): if fp == -1: continue flip_tracked += 1 if fp is not None and fp <= 1: flip_at_01 += 1 if fp is not None and fp == 0: flip_at_0 += 1 aggregate = { 'perfect_rate': perfect_count / total_swaps if total_swaps > 0 else 0, 'state_correct_rate': state_correct_count / total_swaps if total_swaps > 0 else 0, 'suppression_rate': suppressed_count / total_swaps if total_swaps > 0 else 0, 'avg_tier': tier_sum / total_swaps if total_swaps > 0 else 0, 'flip_at_01_rate': flip_at_01 / flip_tracked if flip_tracked > 0 else 0, 'flip_at_0_rate': flip_at_0 / flip_tracked if flip_tracked > 0 else 0, 'flip_tracked': flip_tracked, } stats = { 'total_swaps': total_swaps, 'tier_counts': tier_counts, 'tier_rates': tier_rates, 'aggregate': aggregate, 'correlations': analysis.get('correlations', {}), 'insights': analysis.get('insights', []), } self._stats_cache = stats return stats def get_methodology_stats(self) -> Dict[str, Any]: """Compute dynamic stats for the About / methodology panel. Scans the current run's swap files for regime distribution, aggregate vsMax, and target recovery rates. Results are cached. """ cache_key = "_methodology_stats" if hasattr(self, cache_key) and getattr(self, cache_key) is not None: return getattr(self, cache_key) stats = self.get_stats() dc = self.get_domain_config() total_swaps = stats.get("total_swaps", 0) aggregate = stats.get("aggregate", {}) regime_counts: Dict[str, int] = {} vs_max_values: list = [] recovery_count = 0 recovery_total = 0 by_source = self.swaps_dir / "by_source" if by_source.exists(): for swap_path in by_source.glob("*/to_*.json"): if _is_control_variant(swap_path.stem): continue try: with open(swap_path, "r", encoding="utf-8") as fh: d = json.load(fh) except (json.JSONDecodeError, IOError): continue pos0 = d.get("evaluation", {}).get("position_0_comparison", {}) regime = _classify_regime(pos0) if regime: regime_counts[regime] = regime_counts.get(regime, 0) + 1 contrast = ( d.get("evaluation", {}) .get("logit_trajectory", {}) .get("contrast_groups", {}) .get("same_dataset", {}) .get("aggregate", {}) ) vm = contrast.get("best_target_minus_max") if vm is not None: vs_max_values.append(vm) traj_sum = d.get("evaluation", {}).get("logit_trajectory", {}).get("summary", {}) baseline = d.get("evaluation", {}).get("baseline_logits", {}) if traj_sum and baseline: recovery_total += 1 tgt_baseline_logit = baseline.get("target", {}).get("logit") best_gap = traj_sum.get("best_gap") initial_gap = traj_sum.get("initial_gap") if tgt_baseline_logit is not None and best_gap is not None: recovery_count += 1 regime_total = sum(regime_counts.values()) or 1 result = { "display_name": dc.get("display_name", ""), "total_swaps": total_swaps, "entity_count": dc.get("entity_count", 0), "model_id": dc.get("model_id", ""), "concept_fields": dc.get("concept_fields", []), "answer_field": dc.get("answer_field", ""), "m_ablate": dc.get("m_ablate", -2), "m_amplify": dc.get("m_amplify", 20), "n_tokens": dc.get("n_tokens", 10), "temperature": dc.get("temperature", 0.3), "hit_rate": aggregate.get("perfect_rate", 0), "suppression_rate": aggregate.get("suppression_rate", 0), "flip_at_01_rate": aggregate.get("flip_at_01_rate", 0), "avg_tier": aggregate.get("avg_tier", 0), "vs_max_mean": ( sum(vs_max_values) / len(vs_max_values) if vs_max_values else None ), "vs_max_positive_pct": ( sum(1 for v in vs_max_values if v > 0) / len(vs_max_values) if vs_max_values else None ), "regime_counts": regime_counts, "regime_pcts": { k: round(v / regime_total * 100, 1) for k, v in regime_counts.items() }, "regime_labels": {k: v["label"] for k, v in _REGIME_INFO.items()}, } setattr(self, cache_key, result) return result def _get_gpu_batch_features_index(self) -> Dict[str, List]: """ Build and cache a swap_id -> feature_list index from GPU batch files. Scans _swaps/_work/_gpu_batch_*/features.json which use the format: {"global": [...], "per_prompt": {"swap_id": [...features], ...}} """ if self._gpu_batch_features_index is not None: return self._gpu_batch_features_index index: Dict[str, List] = {} work_dir = self.base_swaps_dir / "_work" if not work_dir.exists(): self._gpu_batch_features_index = index return index for batch_file in sorted(work_dir.glob("_gpu_batch_*/features.json")): try: with open(batch_file, 'r', encoding='utf-8') as f: data = json.load(f) for swap_id, feats in data.get('per_prompt', {}).items(): if isinstance(feats, list): index[swap_id] = feats except (json.JSONDecodeError, IOError): pass self._gpu_batch_features_index = index return index def _resolve_features_path( self, from_slug: str, to_slug: str, variant: Optional[str] = None, ) -> Optional[Path]: """Find the features.json path for a swap, handling variant suffixes. Resolution order: 1. Explicit variant -> ``work/{swap_id}__{variant}/features.json`` 2. Base path (labeled runs) -> ``work/{swap_id}/features.json`` 3. Best variant (mirroring ``get_swap_detail``) work dir. """ base_swap_id = f"{from_slug}__to__{to_slug}" work_dir = self.swaps_dir / "work" if variant: p = work_dir / f"{base_swap_id}__{variant}" / "features.json" return p if p.exists() else None # Try the base path first (works for labeled runs with no suffix) p = work_dir / base_swap_id / "features.json" if p.exists(): return p # Resolve via best variant (same logic as get_swap_detail) best_path = self._best_variant_path(from_slug, to_slug) if best_path is not None: _, variant_suffix = _parse_to_slug(best_path.stem) if variant_suffix: p = work_dir / f"{base_swap_id}__{variant_suffix}" / "features.json" if p.exists(): return p return None def get_swap_features( self, from_slug: str, to_slug: str, variant: Optional[str] = None, ) -> Optional[Dict]: """ Get intervention features for a swap. Tries locations in order: 1. Per-swap file via variant resolution (see ``_resolve_features_path``) 2. GPU batch index: _swaps/_work/_gpu_batch_N/features.json (per_prompt keyed) Returns structured data with ablated/amplified features grouped by layer, enriched with supernode_name from each entity's node_grouping.csv. Also returns source/target grouping lists with ablated/amplified flags. """ swap_id = f"{from_slug}__to__{to_slug}" raw_features: Optional[List] = None features_path = self._resolve_features_path(from_slug, to_slug, variant) if features_path is not None: try: with open(features_path, 'r', encoding='utf-8') as f: raw_features = json.load(f) except (json.JSONDecodeError, IOError): pass # Fall back to GPU batch index if raw_features is None: index = self._get_gpu_batch_features_index() raw_features = index.get(swap_id) if raw_features is None: return None np_model = self.get_domain_config().get('np_model', 'gemma-2-2b') def _load_grouping_map(slug: str) -> Dict: """Return {(layer, index): supernode_name} for an entity.""" entity_dir = self._find_state_dir(slug) if not entity_dir: return {} grouping_path = entity_dir / "02 Node Grouping" / "node_grouping.csv" if not grouping_path.exists(): return {} result = {} try: with open(grouping_path, 'r', encoding='utf-8') as f: reader = csv.DictReader(f) for row in reader: try: key = (int(row.get('layer', 0)), int(row.get('feature', 0))) result[key] = row.get('supernode_name', '') except (ValueError, TypeError): pass except IOError: pass return result def _load_all_groupings(slug: str) -> List[Dict[str, Any]]: """Return sorted groupings with total feature counts for an entity.""" entity_dir = self._find_state_dir(slug) if not entity_dir: return [] grouping_path = entity_dir / "02 Node Grouping" / "node_grouping.csv" if not grouping_path.exists(): return [] grouping_counts: Dict[str, int] = {} try: with open(grouping_path, 'r', encoding='utf-8') as f: reader = csv.DictReader(f) for row in reader: sn = row.get('supernode_name', '').strip() if sn: grouping_counts[sn] = grouping_counts.get(sn, 0) + 1 except IOError: pass return [ {'name': name, 'feature_count': grouping_counts[name]} for name in sorted(grouping_counts) ] source_map = _load_grouping_map(from_slug) target_map = _load_grouping_map(to_slug) all_source_groupings = _load_all_groupings(from_slug) all_target_groupings = _load_all_groupings(to_slug) # Group features by type and layer ablated = [] amplified = [] layer_counts: Dict = {'ablated': {}, 'amplified': {}} ablated_grouping_names: set = set() amplified_grouping_names: set = set() for feat in raw_features: layer = feat.get('layer', 0) index = feat.get('index', 0) M = feat.get('M', 0) stored_activation = feat.get('stored_activation') np_url = f"https://www.neuronpedia.org/{np_model}/{layer}-clt-hp/{index}" if M < 0: sn = source_map.get((layer, index), '') feature_info = { 'layer': layer, 'index': index, 'M': M, 'stored_activation': stored_activation, 'neuronpedia_url': np_url, 'supernode_name': sn, } ablated.append(feature_info) layer_counts['ablated'][layer] = layer_counts['ablated'].get(layer, 0) + 1 if sn: ablated_grouping_names.add(sn) else: sn = target_map.get((layer, index), '') feature_info = { 'layer': layer, 'index': index, 'M': M, 'stored_activation': stored_activation, 'neuronpedia_url': np_url, 'supernode_name': sn, } amplified.append(feature_info) layer_counts['amplified'][layer] = layer_counts['amplified'].get(layer, 0) + 1 if sn: amplified_grouping_names.add(sn) ablated.sort(key=lambda x: (x['layer'], x['index'])) amplified.sort(key=lambda x: (x['layer'], x['index'])) source_groupings = [ { 'name': grouping['name'], 'feature_count': grouping['feature_count'], 'ablated': grouping['name'] in ablated_grouping_names, } for grouping in all_source_groupings ] target_groupings = [ { 'name': grouping['name'], 'feature_count': grouping['feature_count'], 'amplified': grouping['name'] in amplified_grouping_names, } for grouping in all_target_groupings ] return { 'swap_id': swap_id, 'ablated': ablated, 'amplified': amplified, 'layer_counts': layer_counts, 'summary': { 'ablate_count': len(ablated), 'amplify_count': len(amplified), 'total_count': len(ablated) + len(amplified), }, 'source_groupings': source_groupings, 'target_groupings': target_groupings, } def get_state_profile(self, slug: str) -> Optional[Dict]: """Get comprehensive entity profile with stats (domain-agnostic).""" dc = self.get_domain_config() entity_index = dc['entity_index'] edata = entity_index.get(slug, {}) primary_field = dc['primary_field'] answer_field = dc['answer_field'] label = edata.get(primary_field, '') or self._slug_to_label(slug) answer = edata.get(answer_field, '') entity_dir = self._find_state_dir(slug) profile: Dict[str, Any] = { 'slug': slug, 'label': label, 'answer': answer, 'state': edata.get('state', label), 'city': edata.get('city', edata.get(primary_field, label)), 'capital': edata.get('capital', answer), 'fields': {k: v for k, v in edata.items() if k != 'slug'}, } if entity_dir: manifest = self._load_manifest(entity_dir) if manifest: np_data = manifest.get('neuronpedia', {}) profile['supernodes'] = np_data.get('supernodes', 0) profile['pinned_nodes'] = np_data.get('pinned_nodes', 0) profile['neuronpedia_url'] = np_data.get('url', '') gq = manifest.get('graph_quality', {}) pct = gq.get('error_node_influence_pct') if pct is not None: profile['error_node_influence_pct'] = pct if 'error_node_influence_pct' not in profile: pct = self._get_error_node_influence(slug) if pct is not None: profile['error_node_influence_pct'] = pct profile.update( self._load_entity_logit_metadata(entity_dir, answer)) features_path = (entity_dir / "00 Graph Generation" / "selected_features_with_nodes.json") if features_path.exists(): try: with open(features_path, 'r', encoding='utf-8') as f: features_data = json.load(f) features = features_data.get('features', []) profile['total_features'] = len(features) layer_counts: Dict[int, int] = {} for feat in features: layer = feat.get('layer', 0) layer_counts[layer] = layer_counts.get(layer, 0) + 1 profile['feature_layers'] = layer_counts except (json.JSONDecodeError, IOError): pass grouping_path = (entity_dir / "02 Node Grouping" / "node_grouping.csv") if grouping_path.exists(): try: match_terms: set = set() for cf in dc['concept_fields']: val = edata.get(cf, '') if val: match_terms.add(val.lower()) if label: match_terms.add(label.lower()) supernode_features = [] supernode_layer_counts: Dict[int, int] = {} with open(grouping_path, 'r', encoding='utf-8') as f: reader = csv.DictReader(f) for row in reader: sn = row.get('supernode_name', '').lower() if any(t in sn or sn in t for t in match_terms): layer = int(row.get('layer', 0)) supernode_features.append({ 'layer': layer, 'feature': row.get('feature', ''), 'supernode_name': row.get( 'supernode_name', ''), }) supernode_layer_counts[layer] = ( supernode_layer_counts.get(layer, 0) + 1) profile['supernode_features'] = supernode_features profile['supernode_feature_count'] = len( supernode_features) profile['supernode_layer_counts'] = supernode_layer_counts except (IOError, ValueError): pass # Attack/defense scores from matrix matrix = self.get_matrix() attack_tiers = [] if slug in matrix: for to_slug, tier in matrix[slug].items(): if tier is not None and to_slug != slug: attack_tiers.append(tier) profile['attack_avg'] = (sum(attack_tiers) / len(attack_tiers) if attack_tiers else 0) profile['attack_count'] = len(attack_tiers) attack_ok = len([t for t in attack_tiers if t >= 3]) profile['attack_success_rate'] = (attack_ok / len(attack_tiers) if attack_tiers else 0) defense_tiers = [] wrong_count = 0 for from_slug, targets in matrix.items(): if (from_slug != slug and slug in targets and targets[slug] is not None): tier = targets[slug] defense_tiers.append(tier) if tier == 2.5: wrong_count += 1 profile['defense_avg'] = (sum(defense_tiers) / len(defense_tiers) if defense_tiers else 0) profile['defense_count'] = len(defense_tiers) defense_ok = len([t for t in defense_tiers if t >= 3]) profile['defense_success_rate'] = (defense_ok / len(defense_tiers) if defense_tiers else 0) profile['wrong_state_rate'] = (wrong_count / len(defense_tiers) if defense_tiers else 0) profile['wrong_state_count'] = wrong_count profile['has_token_overlap'] = False if dc['is_usa_states']: overlap_slugs = [ 'colorado_colorado_springs', 'new_york_new_york_city', 'virginia_virginia_beach', 'idaho_idaho_falls', 'missouri_kansas_city', 'indiana_fort_wayne', ] profile['has_token_overlap'] = slug in overlap_slugs profile['swaps_as_target'] = self._get_swap_summaries_as_target(slug) profile['swaps_as_source'] = self._get_swap_summaries_as_source(slug) return profile def _get_swap_summaries_as_target(self, slug: str) -> List[Dict]: """Swaps where this entity is the target.""" matrix = self.get_matrix() entity_map = {s['slug']: s for s in self.get_states()} summaries = [] for from_slug, targets in matrix.items(): if from_slug != slug and slug in targets and targets[slug] is not None: tier = targets[slug] e = entity_map.get(from_slug, {}) summaries.append({ 'from_slug': from_slug, 'from_state': e.get('label', e.get('state', from_slug)), 'from_city': e.get('city', ''), 'tier': tier, }) summaries.sort(key=lambda x: -x['tier']) return summaries def _get_swap_summaries_as_source(self, slug: str) -> List[Dict]: """Swaps where this entity is the source.""" matrix = self.get_matrix() entity_map = {s['slug']: s for s in self.get_states()} summaries = [] if slug in matrix: for to_slug, tier in matrix[slug].items(): if to_slug != slug and tier is not None: e = entity_map.get(to_slug, {}) summaries.append({ 'to_slug': to_slug, 'to_state': e.get('label', e.get('state', to_slug)), 'to_city': e.get('city', ''), 'tier': tier, }) summaries.sort(key=lambda x: -x['tier']) return summaries def _get_error_node_influence(self, slug: str) -> Optional[float]: """Return error_node_influence_pct for an entity, with CSV fallback.""" entity_dir = self._find_state_dir(slug) if not entity_dir: return None manifest = self._load_manifest(entity_dir) if manifest: gq = manifest.get('graph_quality', {}) pct = gq.get('error_node_influence_pct') if pct is not None: return pct # Fallback: compute from graph_feature_static_metrics.csv csv_path = (entity_dir / "00 Graph Generation" / "graph_feature_static_metrics.csv") if not csv_path.exists(): return None try: total = 0.0 error = 0.0 with open(csv_path, 'r', encoding='utf-8') as f: reader = csv.DictReader(f) for row in reader: ni = float(row.get('node_influence', 0) or 0) total += ni feat = row.get('feature', '') if feat == '-1': error += ni if total > 0: return round(error / total * 100, 2) except (IOError, ValueError): pass return None def _find_state_dir(self, slug: str) -> Optional[Path]: """Find the state directory, handling case variations.""" state_dir = self.data_dir / slug if state_dir.exists(): return state_dir # Try case variations for candidate in self.data_dir.iterdir(): if candidate.is_dir() and candidate.name.lower().replace(' ', '_') == slug.lower(): return candidate return None def _load_manifest(self, state_dir: Path) -> Optional[Dict]: """Load manifest.json from state directory.""" manifest_path = state_dir / "manifest.json" if manifest_path.exists(): try: with open(manifest_path, 'r', encoding='utf-8') as f: return json.load(f) except (json.JSONDecodeError, IOError): pass return None @staticmethod def _add_query_params(url: str, params: Dict[str, str], overwrite: bool = False) -> str: """ Add query params to a URL, preserving existing params. Args: url: Base URL params: Params to add overwrite: If True, overwrite existing keys. If False, keep existing keys. """ try: parts = urlsplit(url) except Exception: return url existing = parse_qsl(parts.query, keep_blank_values=True) query_map = {} for k, v in existing: # Keep the first value for each key (Neuronpedia URLs typically do not repeat keys) if k not in query_map: query_map[k] = v for k, v in (params or {}).items(): if k in query_map and not overwrite: continue query_map[k] = v new_query = urlencode(query_map, doseq=False) return urlunsplit((parts.scheme, parts.netloc, parts.path, new_query, parts.fragment)) def get_complete_subgraph_url(self, slug: str) -> Optional[Dict[str, Any]]: """ Return the Neuronpedia "complete subgraph" URL using the uploaded subgraph ID. This uses neuronpedia.subgraph_id from manifest.json and appends it as the `subgraph` query parameter to the manifest neuronpedia.url. """ state_dir = self._find_state_dir(slug) if not state_dir: return None manifest = self._load_manifest(state_dir) if not manifest: return None np_data = manifest.get('neuronpedia', {}) or {} base_url = (np_data.get('url') or '').strip() subgraph_id = (np_data.get('subgraph_id') or '').strip() if not base_url or not subgraph_id: return None full_url = self._add_query_params(base_url, {"subgraph": subgraph_id}, overwrite=False) return { "url": full_url, "base_url": base_url, "subgraph_id": subgraph_id, "mode": "complete", } def _get_neuronpedia_url(self, slug: str) -> Optional[str]: """Get Neuronpedia URL for a state slug.""" # Try exact match first state_dir = self.data_dir / slug if state_dir.exists(): manifest = self._load_manifest(state_dir) if manifest: return manifest.get('neuronpedia', {}).get('url') # Try case variations for candidate in self.data_dir.iterdir(): if candidate.is_dir() and candidate.name.lower().replace(' ', '_') == slug.lower(): manifest = self._load_manifest(candidate) if manifest: return manifest.get('neuronpedia', {}).get('url') return None def _slug_to_state_name(self, slug: str) -> str: """Convert slug to state name.""" # arizona_tucson -> Arizona parts = slug.split('_') if len(parts) >= 1: state = parts[0].replace('_', ' ').title() # Handle two-word states if len(parts) >= 2 and parts[0].lower() in ['new', 'north', 'south', 'west', 'rhode']: state = f"{parts[0].title()} {parts[1].title()}" return state return slug.title() def _slug_to_city(self, slug: str) -> str: """Extract city from slug.""" # arizona_tucson -> Tucson parts = slug.split('_') state_parts = 1 if parts[0].lower() in ['new', 'north', 'south', 'west', 'rhode']: state_parts = 2 city_parts = parts[state_parts:] return ' '.join(p.title() for p in city_parts) def _state_to_abbr(self, state_name: str) -> str: """Convert state name to abbreviation.""" abbrs = { 'Alabama': 'AL', 'Alaska': 'AK', 'Arizona': 'AZ', 'Arkansas': 'AR', 'California': 'CA', 'Colorado': 'CO', 'Connecticut': 'CT', 'Delaware': 'DE', 'Florida': 'FL', 'Georgia': 'GA', 'Hawaii': 'HI', 'Idaho': 'ID', 'Illinois': 'IL', 'Indiana': 'IN', 'Iowa': 'IA', 'Kansas': 'KS', 'Kentucky': 'KY', 'Louisiana': 'LA', 'Maine': 'ME', 'Maryland': 'MD', 'Massachusetts': 'MA', 'Michigan': 'MI', 'Minnesota': 'MN', 'Mississippi': 'MS', 'Missouri': 'MO', 'Montana': 'MT', 'Nebraska': 'NE', 'Nevada': 'NV', 'New Hampshire': 'NH', 'New Jersey': 'NJ', 'New Mexico': 'NM', 'New York': 'NY', 'North Carolina': 'NC', 'North Dakota': 'ND', 'Ohio': 'OH', 'Oklahoma': 'OK', 'Oregon': 'OR', 'Pennsylvania': 'PA', 'Rhode Island': 'RI', 'South Carolina': 'SC', 'South Dakota': 'SD', 'Tennessee': 'TN', 'Texas': 'TX', 'Utah': 'UT', 'Vermont': 'VT', 'Virginia': 'VA', 'Washington': 'WA', 'West Virginia': 'WV', 'Wisconsin': 'WI', 'Wyoming': 'WY' } return abbrs.get(state_name, state_name[:2].upper()) def _get_state_capital(self, state_name: str) -> str: """Get the capital city for a state.""" capitals = { 'Alabama': 'Montgomery', 'Alaska': 'Juneau', 'Arizona': 'Phoenix', 'Arkansas': 'Little Rock', 'California': 'Sacramento', 'Colorado': 'Denver', 'Connecticut': 'Hartford', 'Delaware': 'Dover', 'Florida': 'Tallahassee', 'Georgia': 'Atlanta', 'Hawaii': 'Honolulu', 'Idaho': 'Boise', 'Illinois': 'Springfield', 'Indiana': 'Indianapolis', 'Iowa': 'Des Moines', 'Kansas': 'Topeka', 'Kentucky': 'Frankfort', 'Louisiana': 'Baton Rouge', 'Maine': 'Augusta', 'Maryland': 'Annapolis', 'Massachusetts': 'Boston', 'Michigan': 'Lansing', 'Minnesota': 'Saint Paul', 'Mississippi': 'Jackson', 'Missouri': 'Jefferson City', 'Montana': 'Helena', 'Nebraska': 'Lincoln', 'Nevada': 'Carson City', 'New Hampshire': 'Concord', 'New Jersey': 'Trenton', 'New Mexico': 'Santa Fe', 'New York': 'Albany', 'North Carolina': 'Raleigh', 'North Dakota': 'Bismarck', 'Ohio': 'Columbus', 'Oklahoma': 'Oklahoma City', 'Oregon': 'Salem', 'Pennsylvania': 'Harrisburg', 'Rhode Island': 'Providence', 'South Carolina': 'Columbia', 'South Dakota': 'Pierre', 'Tennessee': 'Nashville', 'Texas': 'Austin', 'Utah': 'Salt Lake City', 'Vermont': 'Montpelier', 'Virginia': 'Richmond', 'Washington': 'Olympia', 'West Virginia': 'Charleston', 'Wisconsin': 'Madison', 'Wyoming': 'Cheyenne' } return capitals.get(state_name, '') def get_simplified_subgraph_url(self, slug: str, max_features: int = 80, max_url_length: int = 4000) -> Optional[Dict]: """ Generate a simplified subgraph URL with important features pinned. Priority order for pinning: 1. Capital logit (layer 27) - always included 2. City embedding (layer -1, city token) - always included 3. State-related features - always included if space 4. Remaining features sorted by node_influence Args: slug: State slug (e.g., 'alabama_Birmingham') max_features: Maximum features to include (default 80) max_url_length: Maximum URL length (default 4000 chars) Returns: Dict with 'url', 'feature_count', 'supernode_count', 'url_length' or None if data not found """ import csv from urllib.parse import quote state_dir = self._find_state_dir(slug) if not state_dir: return None # Load manifest for base URL manifest = self._load_manifest(state_dir) if not manifest or 'neuronpedia' not in manifest: return None base_url = manifest['neuronpedia'].get('url', '') if not base_url: return None dc = self.get_domain_config() edata = dc['entity_index'].get(slug, {}) state_name = edata.get('state', self._slug_to_state_name(slug)) city_name = edata.get('city', self._slug_to_city(slug)) capital_name = edata.get('capital', edata.get(dc['answer_field'], '') or self._get_state_capital(state_name)) # Load graph.json to get the TOP logit by probability # The metrics CSV has incomplete/empty data for layer 27 graph_path = state_dir / "00 Graph Generation" / "graph.json" top_logit_node = None # Will store the full node dict for the top logit top_logit_prob = 0 if graph_path.exists(): try: with open(graph_path, 'r', encoding='utf-8') as f: graph_data = json.load(f) for node in graph_data.get('nodes', []): clerp = node.get('clerp', '') prob = node.get('token_prob', 0) node_id = node.get('node_id', '') # Output logits have format: "Output \" Token\" (p=0.123)" if clerp.startswith('Output ') and prob and prob > top_logit_prob: top_logit_prob = prob # Build a node dict matching our format top_logit_node = { 'node_id': node_id, # Use node_id directly from graph.json 'layer': '27', 'feature_id': node_id.split('_')[1] if '_' in node_id else '', 'ctx_idx': '8', 'token': clerp, 'influence': prob, # Use prob as influence 'supernode': 'output_logit', } except (json.JSONDecodeError, IOError): pass # Load graph metrics (has node_influence) metrics_path = state_dir / "00 Graph Generation" / "graph_feature_static_metrics.csv" if not metrics_path.exists(): return None # Load node grouping for supernode assignments # Key on {layer}_{feature} only, since a feature can appear at multiple positions # in the metrics file but node_grouping only has the peak position grouping_path = state_dir / "02 Node Grouping" / "node_grouping.csv" supernode_map = {} # {layer}_{feature} -> supernode_name if grouping_path.exists(): try: with open(grouping_path, 'r', encoding='utf-8') as f: reader = csv.DictReader(f) for row in reader: layer = row.get('layer', '0') feature = row.get('feature', '0') key = f"{layer}_{feature}" supernode_map[key] = row.get('supernode_name', '') except (IOError, csv.Error): pass # Parse metrics file embeddings = [] # Layer -1 (priority) output_logits = [] # Layer 27 (priority) priority_nodes = [] # State/city related other_nodes = [] # Everything else all_supernode_nodes = {} # supernode_name -> list[node] (features only; excludes embeddings/logits) try: with open(metrics_path, 'r', encoding='utf-8') as f: reader = csv.DictReader(f) for row in reader: layer = row.get('layer', '0') feature = row.get('feature', '0') ctx_idx = row.get('ctx_idx', '0') token = row.get('token', '').strip() feature_id = row.get('id', feature) # For embeddings # Skip residual connections (feature=-1 in layers > -1) if feature == '-1' and layer != '-1': continue # Parse node_influence for sorting try: influence = float(row.get('node_influence', 0) or 0) except (ValueError, TypeError): influence = 0 # Build node_id based on layer type # IMPORTANT: Use feature_id (the 'id' column) for ALL node IDs # The 'feature' column contains internal IDs that Neuronpedia doesn't recognize if layer == '-1': # Embedding: E_{id}_{ctx_idx} node_id = f"E_{feature_id}_{ctx_idx}" else: # Regular feature: {layer}_{id}_{ctx_idx} node_id = f"{layer}_{feature_id}_{ctx_idx}" # Get supernode from grouping or infer from token # Key on {layer}_{feature_id} without position to match all occurrences grouping_key = f"{layer}_{feature_id}" supernode = supernode_map.get(grouping_key, '') # For embeddings (layer -1), infer supernode from token if layer == '-1' and not supernode: clean_token = token.strip() if clean_token: # Use cleaned token as supernode name supernode = clean_token.replace('<', '').replace('>', '') # For output logits (layer 27), assign to output group if layer == '27' and not supernode: supernode = 'output_logit' node = { 'node_id': node_id, 'layer': layer, 'feature': feature, 'feature_id': feature_id, 'ctx_idx': ctx_idx, 'token': token, 'influence': influence, 'supernode': supernode, } # Track full membership for each supernode (features only) if supernode and layer not in ('-1', '27'): if supernode not in all_supernode_nodes: all_supernode_nodes[supernode] = [] all_supernode_nodes[supernode].append(node) # Categorize by priority if layer == '-1': # Exclude functional embeddings (bos, The, of, the, etc.) excluded_embeddings = {'bos', 'The', 'the', 'of', 'a', 'an', 'is'} clean_token = token.strip() if clean_token in excluded_embeddings or clean_token.replace('<', '').replace('>', '') in excluded_embeddings: continue # Skip this embedding embeddings.append(node) # Check if this is city or state embedding if city_name.lower() in token.lower(): node['priority'] = 'city' elif state_name.lower() in token.lower(): node['priority'] = 'state' elif layer == '27': output_logits.append(node) elif (city_name.lower() in token.lower() or state_name.lower() in token.lower() or (capital_name and capital_name.lower() in token.lower()) or 'capital' in supernode.lower() or state_name.lower() in supernode.lower() or city_name.lower() in supernode.lower() or (capital_name and capital_name.lower() in supernode.lower())): priority_nodes.append(node) else: other_nodes.append(node) except (IOError, csv.Error) as e: return None # Sort each category by node_influence (descending) embeddings.sort(key=lambda x: -x['influence']) priority_nodes.sort(key=lambda x: -x['influence']) other_nodes.sort(key=lambda x: -x['influence']) # Use the top logit from graph.json (highest probability) # This is more reliable than the metrics CSV which has incomplete layer 27 data top_logit = top_logit_node # Identify supernodes whose NAME contains the state or capital name. # For these, we try to include the ENTIRE supernode group (all member nodes), # as long as there is capacity (max_features). import re def _norm_name(s: str) -> str: s = (s or '').lower().replace('_', ' ') s = re.sub(r'[^a-z0-9 ]+', ' ', s) return ' '.join(s.split()) norm_state = _norm_name(state_name) norm_capital = _norm_name(capital_name) forced_supernodes = [] # [(priority, -max_influence, name)] if norm_state or norm_capital: for name, nodes in (all_supernode_nodes or {}).items(): if not name or len(nodes) < 2: continue norm_group = _norm_name(name) match_capital = bool(norm_capital) and (norm_capital in norm_group) match_state = bool(norm_state) and (norm_state in norm_group) if not (match_capital or match_state): continue max_inf = 0 try: max_inf = max((n.get('influence') or 0) for n in nodes) except ValueError: max_inf = 0 forced_priority = 0 if match_capital else 1 forced_supernodes.append((forced_priority, -max_inf, name)) forced_supernodes.sort() # Build forced nodes (include whole matching supernode groups if they fit) forced_nodes = [] forced_included = 0 forced_omitted = 0 forced_node_ids = set() reserved = len(embeddings) + (1 if top_logit else 0) remaining_forced = max_features - reserved if remaining_forced > 0 and forced_supernodes: for _, __, name in forced_supernodes: group_nodes = all_supernode_nodes.get(name, []) group_nodes_sorted = sorted(group_nodes, key=lambda x: -(x.get('influence') or 0)) unique_group = [n for n in group_nodes_sorted if n.get('node_id') and n['node_id'] not in forced_node_ids] if len(unique_group) <= remaining_forced: forced_nodes.extend(unique_group) forced_node_ids.update(n['node_id'] for n in unique_group) forced_included += 1 remaining_forced -= len(unique_group) else: forced_omitted += 1 # Build candidate nodes list (after reserving space for forced groups) # 1. Semantic embeddings (layer -1) - always included # 2. Top output logit (layer 27) - always included # 3. Forced groups (state/capital-named supernodes) - include whole group if it fits # 4. Priority nodes (state/city/capital related) # 5. Fill remaining with other high-influence nodes candidates = [] candidates.extend(priority_nodes) remaining_slots = max_features - len(embeddings) - (1 if top_logit else 0) - len(forced_nodes) - len(priority_nodes) if remaining_slots > 0: candidates.extend(other_nodes[:remaining_slots]) # First pass: only keep candidates that belong to supernodes with 2+ members temp_supernode_groups = {} for node in candidates: supernode = node.get('supernode', '') if not supernode: continue if supernode not in temp_supernode_groups: temp_supernode_groups[supernode] = [] temp_supernode_groups[supernode].append(node) valid_supernodes = {k for k, v in temp_supernode_groups.items() if len(v) >= 2} # Build final selected list (deduplicated by node_id, preserving priority order) selected = [] seen_node_ids = set() def _add_node(node: Dict[str, Any]) -> None: node_id = node.get('node_id', '') if not node_id or node_id in seen_node_ids: return seen_node_ids.add(node_id) selected.append(node) for node in embeddings: _add_node(node) if top_logit: _add_node(top_logit) for node in forced_nodes: _add_node(node) for node in candidates: if node.get('supernode', '') in valid_supernodes: _add_node(node) # Enforce max_features as a hard cap (forced groups are added before candidates, # so truncation only drops lower-priority tail nodes). if len(selected) > max_features: selected = selected[:max_features] # Build pinnedIds and supernodes pinned_ids = [] supernode_groups = {} for node in selected: pinned_ids.append(node['node_id']) # Embeddings (layer -1) and output logits (layer 27) are pinned but NOT grouped if node['layer'] in ('-1', '27'): continue supernode = node['supernode'] if supernode: if supernode not in supernode_groups: supernode_groups[supernode] = [] supernode_groups[supernode].append(node['node_id']) # Filter supernodes to only groups with 2+ members supernode_groups = {k: v for k, v in supernode_groups.items() if len(v) >= 2} # Build supernodes JSON array supernodes_array = [] for name, ids in sorted(supernode_groups.items()): supernodes_array.append([name] + ids) # Encode and build URL pinned_param = ','.join(pinned_ids) supernodes_json = json.dumps(supernodes_array, separators=(',', ':')) separator = '&' if '?' in base_url else '?' full_url = f"{base_url}{separator}pinnedIds={quote(pinned_param)}&supernodes={quote(supernodes_json)}" # If URL too long, reduce features while len(full_url) > max_url_length and len(selected) > 20: # Remove lowest influence non-priority nodes selected = selected[:-5] pinned_ids = [n['node_id'] for n in selected] supernode_groups = {} for node in selected: # Skip embeddings and logits for grouping if node['layer'] in ('-1', '27'): continue if node['supernode']: if node['supernode'] not in supernode_groups: supernode_groups[node['supernode']] = [] supernode_groups[node['supernode']].append(node['node_id']) # Filter to 2+ node groups only supernode_groups = {k: v for k, v in supernode_groups.items() if len(v) >= 2} supernodes_array = [[name] + ids for name, ids in sorted(supernode_groups.items())] pinned_param = ','.join(pinned_ids) supernodes_json = json.dumps(supernodes_array, separators=(',', ':')) full_url = f"{base_url}{separator}pinnedIds={quote(pinned_param)}&supernodes={quote(supernodes_json)}" return { 'url': full_url, 'base_url': base_url, 'feature_count': len(selected), 'supernode_count': len(supernode_groups), 'url_length': len(full_url), 'embeddings_count': len([n for n in selected if n['layer'] == '-1']), 'logits_count': len([n for n in selected if n['layer'] == '27']), 'priority_count': len([n for n in selected if n in priority_nodes]), 'forced_supernodes_included': forced_included, 'forced_supernodes_omitted': forced_omitted, }