Sync explorer_app.py and clip_utils.py from main repo

scripts/explorer_app.py

@@ -1115,8 +1115,9 @@ name_input.on_change('value', on_name_change)
 
 
 # ---------- Gemini auto-interp button ----------
-_N_GEMINI_IMAGES = 6
-_GEMINI_MODEL    = "gemini-2.5-flash"
+_N_GEMINI_IMAGES    = 6
+_GEMINI_MODEL       = "gemini-2.5-flash"
+_GEMINI_HM_ALPHA    = 0.25   # heatmap overlay opacity sent to Gemini
 
 def _resolve_img_path(stored_path):
     """Mirror the path resolution from auto_interp_vlm.py."""
@@ -1132,8 +1133,11 @@ def _resolve_img_path(stored_path):
     return None
 
 
-def _gemini_label_thread(feat, mei_paths, doc):
-    """Run in a worker thread: call Gemini and push the result back to the doc."""
+def _gemini_label_thread(feat, mei_items, doc):
+    """Run in a worker thread: call Gemini and push the result back to the doc.
+
+    mei_items: list of (path_str, heatmap_np_or_None) where heatmap is (H, W) float32.
+    """
     try:
         from google import genai
         from google.genai import types
@@ -1141,10 +1145,14 @@ def _gemini_label_thread(feat, mei_paths, doc):
         SYSTEM_PROMPT = (
             "You are labeling features of a Sparse Autoencoder (SAE) trained on a "
             "vision transformer. Each SAE feature is a sparse direction in activation "
-            "space that fires strongly on certain visual patterns."
+            "space that fires strongly on certain visual patterns. "
+            "Each image has a colour heatmap overlay highlighting the patches where "
+            "the feature activates most strongly."
         )
         USER_PROMPT = (
-            "The images below are the top maximally-activating images for one SAE feature. "
+            "The images below show the top maximally-activating images for one SAE feature, "
+            "with a heatmap overlay showing where in each image the feature fires most strongly. "
+            "Focus on the highlighted regions. "
             "In 2–5 words, give a precise label for the visual concept this feature detects. "
             "Be specific — prefer 'dog snout close-up' over 'dog', or 'brick wall texture' "
             "over 'texture'. "
@@ -1153,12 +1161,23 @@ def _gemini_label_thread(feat, mei_paths, doc):
 
         client = genai.Client(api_key=_gemini_api_key)
         parts = []
-        for p in mei_paths[:_N_GEMINI_IMAGES]:
-            resolved = _resolve_img_path(p)
+        for path, heatmap in mei_items[:_N_GEMINI_IMAGES]:
+            resolved = _resolve_img_path(path)
             if resolved is None:
                 continue
             try:
                 img = Image.open(resolved).convert("RGB").resize((224, 224), Image.BILINEAR)
+                if heatmap is not None:
+                    img_arr = np.array(img).astype(np.float32) / 255.0
+                    hm_up = cv2.resize(heatmap.astype(np.float32), (224, 224),
+                                       interpolation=cv2.INTER_CUBIC)
+                    hmax = hm_up.max()
+                    if hmax > 0:
+                        hm_up /= hmax
+                    overlay = ALPHA_JET(hm_up)
+                    ov_alpha = overlay[:, :, 3:4] * _GEMINI_HM_ALPHA
+                    blended = img_arr * (1 - ov_alpha) + overlay[:, :, :3] * ov_alpha
+                    img = Image.fromarray(np.clip(blended * 255, 0, 255).astype(np.uint8))
                 buf = io.BytesIO()
                 img.save(buf, format="JPEG", quality=85)
                 parts.append(types.Part.from_bytes(data=buf.getvalue(), mime_type="image/jpeg"))
@@ -1218,13 +1237,16 @@ def _on_gemini_click():
         return
 
     n_top_stored = top_img_idx.shape[1]
-    mei_paths = []
+    mei_items = []
     for j in range(n_top_stored):
         idx = top_img_idx[feat, j].item()
         if idx >= 0:
-            mei_paths.append(image_paths[idx])
+            hm = None
+            if top_heatmaps is not None:
+                hm = top_heatmaps[feat, j].float().numpy().reshape(patch_grid, patch_grid)
+            mei_items.append((image_paths[idx], hm))
 
-    if not mei_paths:
+    if not mei_items:
         gemini_status_div.text = "<span style='color:#c00'>No MEI paths found.</span>"
         return
 
@@ -1234,7 +1256,7 @@ def _on_gemini_click():
     doc = curdoc()
     t = threading.Thread(
         target=_gemini_label_thread,
-        args=(feat, mei_paths, doc),
+        args=(feat, mei_items, doc),
         daemon=True,
     )
     t.start()