Sync explorer_app.py and clip_utils.py from main repo

scripts/explorer_app.py

@@ -100,11 +100,11 @@ parser.add_argument("--clip-model", type=str, default="openai/clip-vit-large-pat
 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-path", type=str, default=None,
-                    help="Path to SAE weights (.pth) for the primary dataset — "
-                         "enables the Download SAE weights button in the summary panel")
-parser.add_argument("--compare-sae-paths", type=str, nargs="*", default=[],
-                    help="SAE weight paths for each --compare-data dataset (in order)")
+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)")
 args = parser.parse_args()
 
 
@@ -125,7 +125,7 @@ def _get_clip():
 
 # ---------- Load all datasets into a unified list ----------
 
-def _load_dataset_dict(path, label, sae_path=None):
+def _load_dataset_dict(path, label, sae_url=None):
     """Load one explorer_data.pt file and return a unified dataset dict."""
     print(f"Loading [{label}] from {path} ...")
     d = torch.load(path, map_location='cpu', weights_only=False)
@@ -208,7 +208,7 @@ def _load_dataset_dict(path, label, sae_path=None):
     else:
         entry['patch_acts'] = None
 
-    entry['sae_path'] = sae_path
+    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 else 'no'}, "
@@ -220,17 +220,17 @@ _all_datasets = []
 _active = [0]   # index of the currently displayed dataset
 
 # Primary dataset — always loaded eagerly
-_all_datasets.append(_load_dataset_dict(args.data, args.primary_label, sae_path=args.sae_path))
+_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_paths[_ci]
-             if args.compare_sae_paths and _ci < len(args.compare_sae_paths)
+    _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_path': _csae})
+    _all_datasets.append({'label': _clabel, 'path': _cpath, '_lazy': True, 'sae_url': _csae})
 
 
 def _ensure_loaded(idx):
@@ -238,7 +238,7 @@ def _ensure_loaded(idx):
     ds = _all_datasets[idx]
     if ds.get('_lazy', False):
         print(f"[Lazy load] Loading '{ds['label']}' on first access ...")
-        _all_datasets[idx] = _load_dataset_dict(ds['path'], ds['label'], sae_path=ds.get('sae_path'))
+        _all_datasets[idx] = _load_dataset_dict(ds['path'], ds['label'], sae_url=ds.get('sae_url'))
 
 
 def _apply_dataset_globals(idx):
@@ -808,9 +808,8 @@ def _on_dataset_switch(attr, old, new):
     # Rebuild active-feature pool for random button
     _active_feats = [int(i) for i in range(d_model) if feature_frequency[i].item() > 0]
 
-    # Update summary panel and SAE download source
+    # Update summary panel
     summary_div.text = _make_summary_html()
-    _update_download_source()
 
     # Show/hide patch explorer depending on token type and data availability.
     ds = _all_datasets[idx]
@@ -1493,6 +1492,10 @@ def _make_summary_html():
     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>'
+                  if sae_url else '')
     return f"""
 <div style="background:#f0f4f8;padding:12px;border-radius:6px;margin-bottom:8px;">
 <h2 style="margin:0 0 8px 0">SAE Feature Explorer</h2>
@@ -1505,44 +1508,12 @@ def _make_summary_html():
 <tr><td><b>Dead:</b></td><td>{n_dead:,} ({100*n_dead/d_model:.1f}%)</td></tr>
 <tr><td><b>Images:</b></td><td>{n_images:,}</td></tr>
 <tr><td><b>Tokens/image:</b></td><td>{tok_label}</td></tr>
+{dl_row}
 </table>
 </div>"""
 
 summary_div = Div(text=_make_summary_html(), width=700)
 
-# --- SAE weights download button ---
-# Data is pre-loaded into the source so the JS download fires synchronously
-# with the user click (bypassing browser gesture-context restrictions).
-
-def _sae_b64(sae_path):
-    """Return (b64_str, filename) for a SAE weights file, or ('', '') if unavailable."""
-    if not sae_path or not os.path.exists(sae_path):
-        return '', ''
-    with open(sae_path, 'rb') as f:
-        return base64.b64encode(f.read()).decode('ascii'), os.path.basename(sae_path)
-
-_init_b64, _init_fname = _sae_b64(_all_datasets[0].get('sae_path'))
-_download_source = ColumnDataSource(data=dict(b64=[_init_b64], filename=[_init_fname]))
-
-sae_download_btn = Button(label="\u2b07 Download SAE weights", button_type="default", width=220)
-sae_download_btn.js_on_click(CustomJS(args=dict(src=_download_source), code="""
-    const b64 = src.data['b64'][0];
-    const fname = src.data['filename'][0];
-    if (!b64) { alert('No SAE weights available for this model.'); return; }
-    const bytes = Uint8Array.from(atob(b64), c => c.charCodeAt(0));
-    const blob = new Blob([bytes], {type: 'application/octet-stream'});
-    const url = URL.createObjectURL(blob);
-    const a = document.createElement('a');
-    a.href = url; a.download = fname; a.click();
-    URL.revokeObjectURL(url);
-"""))
-
-def _update_download_source():
-    """Reload SAE weights for the active dataset into _download_source."""
-    ds = _all_datasets[_active[0]]
-    b64, fname = _sae_b64(ds.get('sae_path'))
-    _download_source.data = dict(b64=[b64], filename=[fname])
-
 
 # ---------- Patch Explorer ----------
 # Click patches of an image to find the top active SAE features for that region.
@@ -1920,7 +1891,7 @@ patch_explorer_panel = column(
     patch_feat_table,
 )
 
-summary_section = _make_collapsible("SAE Summary",      column(summary_div, sae_download_btn))
+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)