Sync space code

scripts/explorer_app.py

@@ -4,29 +4,13 @@ Interactive SAE Feature Explorer - Bokeh Server App.
 Visualizes SAE features with:
   - UMAP scatter plot of features (activation-based and dictionary-based)
   - Click a feature to see its top-activating images with heatmap overlays
-  - 75th percentile images for distribution understanding
-  - Patch explorer: click patches of any image to find active features
+  - Patch explorer: click patches of any image to find active SAE features
+    (uses live GPU inference via the backbone + SAE loaded from --sae-path)
   - Feature naming: assign names to features, saved to JSON, searchable
+  - CLIP text search, Gemini auto-interp, DynaDiff brain steering panel
+  - Optional NSD brain MEI dataset (--brain-data) shown in the dataset dropdown
 
-All display is driven by pre-computed sidecars (_heatmaps.pt, _patch_acts.pt).
-No GPU or model weights are required at serve time.
-
-Launch:
-    bokeh serve explorer_app.py --port 5006 --allow-websocket-origin="*" \
-        --session-token-expiration 86400 \
-        --args \
-          --data ../../smart_init_stability_SAE/explorer_data_d32000_k160_val.pt \
-          --image-dir /scratch.global/lee02328/val \
-          --extra-image-dir /scratch.global/lee02328/coco/val2017 \
-          --primary-label "DINOv3 L24 Spatial (d=32K)" \
-          --compare-data ../../smart_init_stability_SAE/explorer_data_18.pt \
-          --compare-labels "DINOv3 L18 Spatial (d=20K)" \
-          --phi-dir /path/to/phis \
-          --brain-data /path/to/brain_meis_dinov3.pt \
-          --brain-thumbnails /path/to/nsd_thumbs
-
-Then SSH tunnel: ssh -L 5006:<node>:5006 <user>@<login-node>
-Open: http://localhost:5006/explorer_app
+Launch: see run_explorer.sh
 """
 
 import argparse
@@ -37,7 +21,6 @@ import base64
 import random
 import threading
 from collections import OrderedDict
-from functools import partial
 
 import cv2
 import numpy as np
@@ -56,7 +39,7 @@ from bokeh.layouts import column, row
 from bokeh.events import MouseMove
 from bokeh.models import (
     ColumnDataSource, HoverTool, Div, Select, TextInput, Button,
-    DataTable, TableColumn, NumberFormatter, IntEditor, NumberEditor,
+    DataTable, TableColumn, NumberFormatter, NumberEditor,
     Slider, Toggle, RadioButtonGroup, CustomJS,
 )
 from bokeh.plotting import figure
@@ -77,11 +60,6 @@ parser.add_argument("--inference-cache-size", type=int, default=64,
 parser.add_argument("--names-file", type=str, default=None,
                     help="Path to JSON file for saving feature names "
                          "(default: <data>_feature_names.json)")
-parser.add_argument("--compare-data", type=str, nargs="*", default=[],
-                    help="Additional explorer_data.pt files to show in cross-dataset "
-                         "comparison panel (e.g. layer 18, CLS SAE)")
-parser.add_argument("--compare-labels", type=str, nargs="*", default=[],
-                    help="Display labels for each --compare-data file")
 parser.add_argument("--primary-label", type=str, default="Primary",
                     help="Display label for the primary --data file")
 parser.add_argument("--clip-model", type=str, default="openai/clip-vit-large-patch14",
@@ -91,10 +69,7 @@ parser.add_argument("--google-api-key", type=str, default=None,
                     help="Google API key for Gemini auto-interp button "
                          "(default: GOOGLE_API_KEY env var)")
 parser.add_argument("--sae-url", type=str, default=None,
-                    help="Download URL for the primary dataset's SAE weights — "
-                         "shown as a link in the summary panel")
-parser.add_argument("--compare-sae-urls", type=str, nargs="*", default=[],
-                    help="Download URLs for each --compare-data dataset's SAE weights (in order)")
+                    help="Download URL for the SAE weights — shown as a link in the summary panel")
 parser.add_argument("--phi-dir", type=str, default=None,
                     help="Directory containing Phi_cv_*.npy, phi_c_*.npy, voxel_coords.npy "
                          "(brain-alignment data; enables cortical profile and brain leverage features)")
@@ -133,27 +108,29 @@ args = parser.parse_args()
 
 
 # ---------- Lazy CLIP model (loaded on first free-text query) ----------
-# _clip_handle[0] is None until the first out-of-vocab query is issued.
-_clip_handle = [None]   # (model, processor, device)
+_clip_handle = None   # (model, processor, device), set on first use
 
 def _get_clip():
     """Load CLIP once and cache it."""
-    if _clip_handle[0] is None:
+    global _clip_handle
+    if _clip_handle is None:
         _dev = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
         print(f"[CLIP] Loading {args.clip_model} on {_dev} (first free-text query)...")
         _m, _p = load_clip(_dev, model_name=args.clip_model)
-        _clip_handle[0] = (_m, _p, _dev)
+        _clip_handle = (_m, _p, _dev)
         print("[CLIP] Ready.")
-    return _clip_handle[0]
+    return _clip_handle
 
 
 # ---------- GPU backbone + SAE runner (optional, lazy-loaded) ----------
-_gpu_runner = [None]   # (forward_fn, sae, transform_fn, n_reg, extract_tokens_fn, backbone_name, device) or None
+# Tuple of (forward_fn, sae, transform_fn, n_reg, extract_tokens_fn, backbone_name, device)
+_gpu_runner = None
 
 def _get_gpu_runner():
-    """Load backbone + SAE on GPU once; return (forward_fn, sae, transform_fn, device) or None."""
-    if _gpu_runner[0] is not None:
-        return _gpu_runner[0]
+    """Load backbone + SAE on GPU once; return the runner tuple or None."""
+    global _gpu_runner
+    if _gpu_runner is not None:
+        return _gpu_runner
     if not args.sae_path:
         return None
     if not torch.cuda.is_available():
@@ -167,9 +144,9 @@ def _get_gpu_runner():
     print(f"[GPU runner] Loading {args.backbone} layer {args.layer} + SAE on {_dev} ...")
     _fwd, _d_hidden, _n_reg, _tfm = load_batched_backbone(args.backbone, args.layer, _dev)
     _sae = load_sae(args.sae_path, _d_hidden, d_model, args.top_k, _dev)
-    _gpu_runner[0] = (_fwd, _sae, _tfm, _n_reg, _et, args.backbone, _dev)
+    _gpu_runner = (_fwd, _sae, _tfm, _n_reg, _et, args.backbone, _dev)
     print("[GPU runner] Ready.")
-    return _gpu_runner[0]
+    return _gpu_runner
 
 
 def _run_gpu_inference(pil_img):
@@ -372,27 +349,11 @@ def _load_dataset_dict(path, label, sae_url=None):
         entry['heatmap_patch_grid'] = d['patch_grid']
         has_hm = 'no'
 
-    # Load pre-computed patch activations sidecar if present.
-    # Enables complete GPU-free patch exploration for any image covered by the file.
-    pa_sidecar = os.path.splitext(path)[0] + '_patch_acts.pt'
-    if os.path.exists(pa_sidecar):
-        print(f"  Loading pre-computed patch acts from {os.path.basename(pa_sidecar)} ...")
-        pa = torch.load(pa_sidecar, map_location='cpu', weights_only=True)
-        img_to_row = {int(idx): row for row, idx in enumerate(pa['img_indices'].tolist())}
-        entry['patch_acts'] = {
-            'feat_indices': pa['feat_indices'],  # (n_unique, n_patches, top_k) int16
-            'feat_values':  pa['feat_values'],   # (n_unique, n_patches, top_k) float16
-            'img_to_row':   img_to_row,
-        }
-        print(f"    patch_acts: {len(img_to_row)} images covered (GPU-free patch explorer)")
-    else:
-        entry['patch_acts'] = None
-
     entry['sae_url'] = sae_url
 
     print(f"  d={entry['d_model']}, n={entry['n_images']}, token={entry['token_type']}, "
           f"backbone={entry['backbone']}, clip={'yes' if (cs is not None or entry.get('clip_embeds') is not None) else 'no'}, "
-          f"heatmaps={has_hm}, patch_acts={'yes' if entry['patch_acts'] else 'no'}")
+          f"heatmaps={has_hm}")
     return entry
 
 
@@ -410,7 +371,7 @@ class _S:
     render_token: int = 0        # incremented on each feature selection; stale renders bail out
     search_filter = None         # set of feature indices matching the current name search, or None
     color_by: str = "Log Frequency"   # which field drives UMAP point colour
-    hf_push: object = None       # active Bokeh timeout handle for debounced HuggingFace upload
+    hf_push = None               # active Bokeh timeout handle for debounced HuggingFace upload
     patch_img = None             # image index currently loaded in the patch explorer
     patch_z = None               # cached (n_patches, d_model) float32 for the loaded image
 
@@ -423,16 +384,6 @@ def _ds():
 # Primary dataset — always loaded eagerly
 _all_datasets.append(_load_dataset_dict(args.data, args.primary_label, sae_url=args.sae_url))
 
-# Compare datasets — stored as lazy placeholders; loaded on first access
-for _ci, _cpath in enumerate(args.compare_data):
-    _clabel = (args.compare_labels[_ci]
-               if args.compare_labels and _ci < len(args.compare_labels)
-               else os.path.basename(_cpath))
-    _csae = (args.compare_sae_urls[_ci]
-             if args.compare_sae_urls and _ci < len(args.compare_sae_urls)
-             else None)
-    _all_datasets.append({'label': _clabel, 'path': _cpath, '_lazy': True, 'sae_url': _csae})
-
 def _load_brain_dataset_dict(path, label, thumb_dir):
     """Load a brain_meis.pt file and return a dataset entry dict.
 
@@ -497,7 +448,6 @@ def _load_brain_dataset_dict(path, label, thumb_dir):
         'feature_names':     {},
         'auto_interp_names': {},
         'sae_url':           None,
-        'patch_acts':        None,
     }
 
     # Load pre-computed heatmaps sidecar if present.
@@ -700,43 +650,23 @@ def _display_name(feat: int) -> str:
 
 
 def compute_patch_activations(img_idx):
-    """Return (n_patches, d_sae) float32 for the active dataset, or None.
+    """Return (n_patches, d_sae) float32 via GPU inference, or None if unavailable.
 
-    Priority order:
-      1. LRU cache
-      2. Pre-computed patch_acts lookup — complete activations for covered images
-      3. GPU live inference — full activations via backbone + SAE (requires --sae-path)
-    Uses a per-dataset LRU cache.
+    Results are cached in a per-dataset LRU cache keyed by image index.
     """
     ds    = _all_datasets[_S.active]
     cache = ds['inference_cache']
 
-    # 1. LRU cache
     if img_idx in cache:
         cache.move_to_end(img_idx)
         return cache[img_idx]
 
-    z_np = None
-
-    # 2. Try patch_acts lookup (complete activations for covered images)
-    pa = ds.get('patch_acts')
-    if pa is not None:
-        row = pa['img_to_row'].get(img_idx)
-        if row is not None:
-            fi = pa['feat_indices'][row].numpy()   # (n_patches, top_k) int16
-            fv = pa['feat_values'][row].float().numpy()  # (n_patches, top_k) float32
-            n_p = fi.shape[0]
-            z_np = np.zeros((n_p, ds['d_model']), dtype=np.float32)
-            z_np[np.arange(n_p)[:, None], fi.astype(np.int32)] = fv
-
-    # 3. GPU live inference
-    if z_np is None:
-        try:
-            pil = load_image(img_idx)
-            z_np = _run_gpu_inference(pil)
-        except Exception as _e:
-            print(f"[GPU runner] inference failed for img {img_idx}: {_e}")
-            z_np = None
+    try:
+        pil  = load_image(img_idx)
+        z_np = _run_gpu_inference(pil)
+    except Exception as _e:
+        print(f"[GPU runner] inference failed for img {img_idx}: {_e}")
+        z_np = None
 
     if z_np is not None:
         cache[img_idx] = z_np
@@ -1306,21 +1236,16 @@ def _on_dataset_switch(attr, old, new):
     # Update summary panel
     summary_div.text = _make_summary_html()
 
-    # Show/hide patch explorer depending on token type and data availability.
+    # Show/hide patch explorer depending on token type (spatial required) and GPU availability.
     ds = _all_datasets[idx]
-    has_heatmaps   = ds.get('top_heatmaps') is not None
-    has_patch_acts = ds.get('patch_acts') is not None
     can_explore = (
         ds.get('token_type', 'spatial') == 'spatial'
-        and (has_heatmaps or has_patch_acts)
+        and bool(args.sae_path)
     )
     patch_fig.visible = can_explore
     patch_info_div.visible = can_explore
     if not can_explore:
-        if ds.get('token_type') == 'cls':
-            reason = "CLS token — no patch grid"
-        else:
-            reason = "no pre-computed heatmaps or patch_acts for this model"
+        reason = "CLS token — no patch grid" if ds.get('token_type') == 'cls' else "no --sae-path provided"
         patch_info_div.text = (
             f'<p style="color:#888;font-style:italic">Patch explorer unavailable: {reason}.</p>')
         patch_info_div.visible = True
@@ -2223,8 +2148,6 @@ def _make_summary_html():
     backbone_label = ds.get('backbone', 'dinov3').upper()
     clip_label = "yes" if (ds['clip_scores'] is not None or ds.get('clip_embeds') is not None) else "no"
     hm_label   = "yes" if ds.get('top_heatmaps') is not None else "no"
-    pa = ds.get('patch_acts')
-    pa_label   = f"yes ({len(pa['img_to_row'])} images)" if pa is not None else "no — run --save-patch-acts"
     sae_url    = ds.get('sae_url')
     dl_row     = (f'<tr><td><b>SAE weights:</b></td>'
                   f'<td><a href="{sae_url}" download style="color:#1a6faf">⬇ Download</a></td></tr>'
@@ -2250,7 +2173,7 @@ summary_div = Div(text=_make_summary_html(), width=700)
 
 # ---------- Patch Explorer ----------
 # Click patches of an image to find the top active SAE features for that region.
-# Activations are served from pre-computed sidecars (no GPU required at serve time).
+# Activations are computed on-the-fly via GPU inference (backbone + SAE from --sae-path).
 
 _PATCH_FIG_PX = 400
 
@@ -2383,7 +2306,7 @@ def _do_load_patch_image():
     patch_info_div.text = (
         "<span style='color:#1a6faf'>&#x23F3; Computing patch activations"
         + (" (running GPU inference — first image may take ~10 s)…"
-           if _gpu_runner[0] is None and args.sae_path else "…")
+           if _gpu_runner is None and args.sae_path else "…")
         + "</span>"
     )
 
@@ -2410,18 +2333,14 @@ def _do_load_patch_image():
             if z_np is None:
                 patch_feat_table.visible = False
                 patch_info_div.text = (
-                    f"<b style='color:#888'>Image {img_idx} has no pre-computed patch activations "
-                    f"and no GPU runner is available.  Pass --sae-path to the explorer to enable "
-                    f"live GPU inference for any image.</b>"
+                    f"<b style='color:#888'>GPU inference unavailable for image {img_idx}. "
+                    f"Ensure --sae-path is set and the GPU runner loaded successfully.</b>"
                 )
                 return
 
             patch_feat_table.visible = True
-            _ds = _all_datasets[_S.active]
-            _pa = _ds.get('patch_acts')
-            source = "patch_acts" if (_pa is not None and img_idx in _pa['img_to_row']) else "GPU inference"
             patch_info_div.text = (
-                f"Image {img_idx} loaded ({source}). "
+                f"Image {img_idx} loaded. "
                 f"Drag to select a region, or click individual patches."
             )
 
@@ -2667,7 +2586,7 @@ summary_section     = _make_collapsible("SAE Summary",          summary_div)
 patch_section       = _make_collapsible("Patch Explorer",       patch_explorer_panel)
 clip_section        = _make_collapsible("CLIP Text Search",     clip_search_panel)
 
-_ds_select_row = ([dataset_select] if len(_all_datasets) > 1 and args.compare_data else [])
+_ds_select_row = ([dataset_select] if len(_all_datasets) > 1 else [])
 left_panel = column(*_ds_select_row, controls, umap_fig, feature_list_panel)
 
 middle_panel = column(