extra_materials/visualize_probe_features.py

"""
visualize_probe_features.py — Multi-layer probe feature visualizer for LLaVA.

Uses a single forward pass per batch to cache residual-stream activations at ALL
target layers simultaneously (shared SAE d_sae), then runs the SAE per-layer on
the cached activations.  This is ~Nx faster than the old per-layer approach
(where N = number of layers).

For each batch:
  1. (Optional) Generate captions with the model (once, shared across layers).
  2. One forward pass caching activations at every target layer's hook point.
  3. For each layer: SAE → sparse features → update top-k heaps.

Generates ONE self-contained interactive HTML:
  Layers ▶ Top Features (id + probe weight) ▶ Top Image Patches / Text Tokens

Data modes (--data_mode):
  toilet  — only images from HF "pbcong/bathroom-toilet" that match
            --object_mode (toilet | bathroom).  Supports --caption_mode.
            Uses visualize_multilayer_sae_features.create_dataloader.
  cc3m    — full CC3M or COCO dataset via HF or a plain image folder.
            Uses visualize_multilayer_sae_features.create_dataloader.

Usage
-----
  # Toilet subset
  python training/visualize_probe_features.py \\
      --data_mode     toilet \\
      --sae_ckpt      training/multilayer_sae_ckpt/last.ckpt \\
      --probe_dir     training/multilayer_sae_ckpt \\
      --image_folder  /path/to/cc3m_images/train \\
      --object_mode   toilet \\
      --caption_mode  generated \\
      --layers 0 1 2 3 4 5 6 \\
      --output_dir    outputs/probe_features_toilet \\
      --device_id 0

  # Full CC3M (HF dataset)
  python training/visualize_probe_features.py \\
      --data_mode      cc3m \\
      --sae_ckpt       training/multilayer_sae_ckpt/last.ckpt \\
      --probe_dir      training/multilayer_sae_ckpt \\
      --hf_dataset     pixparse/cc3m-wds \\
      --local_val_path /path/to/cc3m_images/train \\
      --split          train \\
      --layers 0 1 2 3 4 5 6 \\
      --output_dir     outputs/probe_features_cc3m \\
      --device_id 0

  # Plain image folder
  python training/visualize_probe_features.py \\
      --data_mode  cc3m \\
      --data_dir   /path/to/images \\
      --sae_ckpt   training/multilayer_sae_ckpt/last.ckpt \\
      --probe_dir  training/multilayer_sae_ckpt \\
      --layers 0 1 2 3 4 5 6 \\
      --output_dir outputs/probe_features_folder \\
      --device_id 0
"""

import sys
import os
import json
import argparse
import pickle
import shutil
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import torch
import torch.distributed as dist
from tqdm import tqdm

sys.path.insert(0, str(Path(__file__).parent.parent))

# ── Lazy-import both source modules via importlib ─────────────────────────────
import importlib.util as _ilu


def _load_module(name: str, path: Path):
    spec = _ilu.spec_from_file_location(name, path)
    mod  = _ilu.module_from_spec(spec)
    sys.modules[name] = mod   # register before exec so pickle can resolve the module
    spec.loader.exec_module(mod)
    return mod


_scripts_dir = Path(__file__).parent
_vtf  = _load_module("visualize_multilayer_sae_features",    _scripts_dir / "visualize_multilayer_sae_features.py")

# Shared helpers from the multilayer SAE feature visualizer
setup_distributed        = _vtf.setup_distributed
cleanup_distributed      = _vtf.cleanup_distributed
TopKHeap                 = _vtf.TopKHeap
build_image_html         = _vtf.build_image_html
build_text_html          = _vtf.build_text_html
create_dataloader        = _vtf.create_dataloader
ImageRecord              = _vtf.ImageRecord
TextRecord               = _vtf.TextRecord
find_image_token_positions = _vtf.find_image_token_positions
expanded_to_original     = _vtf.expanded_to_original
original_to_expanded     = _vtf.original_to_expanded
N_IMAGE_PATCHES          = _vtf.N_IMAGE_PATCHES
IMAGE_TOKEN_ID           = _vtf.IMAGE_TOKEN_ID

from hallucination.extra_materials.mechanistic_interp.probe.probing import LinearProbe
from sae.SAE_Tools import load_sae_model, get_sae_activations, get_loader
from model.llava.hooked_llava import HookedSAELlavaConditionalGeneration
from sae.Training_Utils import str_to_torch_dtype
from transformers import LlavaProcessor


# ─────────────────────────────────────────────────────────────────────────────
# Probe loading
# ─────────────────────────────────────────────────────────────────────────────

def load_probe_top_features(
    probe_dir: str,
    layer: int,
    top_k: int,
    input_dim: int = 65536,
) -> Tuple[List[int], List[float]]:
    """Return (feature_ids, probe_weights) for the top-k probe features."""
    ckp_path = (
        Path(probe_dir)
        / f"probe_model.language_model.layers.{layer}.hook_resid_post.pt"
    )
    if not ckp_path.exists():
        raise FileNotFoundError(f"Probe checkpoint not found: {ckp_path}")

    ckpt = torch.load(ckp_path, map_location="cpu")
    probe = LinearProbe(input_dim=input_dim, num_outputs=1)
    probe.load_state_dict(ckpt)
    probe.eval()

    weights = probe.weights.squeeze()   # (input_dim,)
    top = torch.topk(weights, k=top_k)
    return top.indices.cpu().tolist(), weights[top.indices].cpu().tolist()


# ─────────────────────────────────────────────────────────────────────────────
# Multi-layer single pass (1 model fwd per batch instead of 1 per layer)
# ─────────────────────────────────────────────────────────────────────────────

@torch.no_grad()
def multi_layer_pass(
    model, sae, processor, dataloader, device,
    layer_probe_info: Dict[int, Tuple[List[int], List[float]]],
    args,
) -> Dict[int, dict]:
    """
    Single pass over the dataloader, hooking ALL target layers at once.

    For each batch:
      1. (Optional) Generate caption with original model.
      2. One forward pass caching residual-stream activations at every target layer.
      3. For each layer: SAE → sparse features → update top-k heaps.

    Returns {layer: {feature_ids, probe_weights, image_heaps, text_heaps,
                     hook_counts, hook_point}}.
    """
    sae_dtype = next(sae.parameters()).dtype
    sae.eval()
    sae.to(device)

    # Build hook-point mapping and per-layer data structures
    hook_points: Dict[int, str] = {}
    target_hooks: set = set()
    layer_data: Dict[int, dict] = {}

    for layer, (feat_ids, feat_wts) in layer_probe_info.items():
        hp = f"model.language_model.layers.{layer}.hook_resid_post"
        hook_points[layer] = hp
        target_hooks.add(hp)
        layer_data[layer] = {
            "image_heaps": {fi: TopKHeap(args.top_images) for fi in feat_ids},
            "text_heaps":  {fi: TopKHeap(args.top_texts)  for fi in feat_ids},
            "hook_counts": {fi: 0 for fi in feat_ids},
        }

    def _make_hook_fn(cache: dict):
        def hook_fn(act, hook):
            if hook.name in target_hooks:
                cache[hook.name] = act.detach().cpu()
        return hook_fn

    n_batches = (
        len(dataloader) if args.max_batches is None
        else min(args.max_batches, len(dataloader))
    )

    for batch_idx, batch in enumerate(tqdm(dataloader, total=n_batches, desc="Scanning")):
        if args.max_batches is not None and batch_idx >= args.max_batches:
            break

        B = batch["input_ids"].shape[0]
        model_inputs = {
            k: batch[k].to(device)
            for k in ("input_ids", "attention_mask", "pixel_values")
        }
        global_idxs = batch["global_idxs"]

        # ── Generated caption mode (once per batch, shared across layers) ────
        if args.caption_mode == "generated":
            gen_ids = model.generate(
                **model_inputs, do_sample=False, num_beams=1,
                use_cache=True, max_new_tokens=args.max_new_tokens,
            )
            full_texts = processor.batch_decode(gen_ids, skip_special_tokens=True)
            captions = [
                txt.split("ASSISTANT:")[-1].strip() if "ASSISTANT:" in txt
                else txt.strip()
                for txt in full_texts
            ]
            images = [
                dataloader.dataset.dataset[gi]["image"] for gi in global_idxs
            ]
            forced_texts = [
                f"USER: <image>\nDescribe this image. \nASSISTANT: {cap}"
                for cap in captions
            ]
            re_processed = processor(
                images=images, text=forced_texts,
                return_tensors="pt", padding=True,
            )
            model_inputs = {
                "input_ids":      re_processed["input_ids"].to(device),
                "attention_mask": re_processed["attention_mask"].to(device),
                "pixel_values":   re_processed["pixel_values"].to(device),
            }
            input_ids_cpu = re_processed["input_ids"]
            B = input_ids_cpu.shape[0]
        else:
            input_ids_cpu = batch["input_ids"]

        # ── Single forward pass, cache ALL layer activations ─────────────────
        cache: dict = {}
        model.run_with_hooks(
            model_inputs,
            fwd_hooks=[(lambda n: n in target_hooks, _make_hook_fn(cache))],
        )

        img_positions = find_image_token_positions(input_ids_cpu)

        # ── Process each layer from the cache ────────────────────────────────
        for layer, (feat_ids, _) in layer_probe_info.items():
            hp = hook_points[layer]
            if hp not in cache:
                continue

            acts = cache[hp]                              # (B, T, D)
            _, T, D = acts.shape
            flat = acts.reshape(B * T, 1, D).to(device, dtype=sae_dtype)
            loader = get_loader(flat, batch_size=min(args.sae_batch, B * T))
            idxs, vals = get_sae_activations(sae, loader, device, no_tqdm=True)
            idxs = idxs.squeeze(1).reshape(B, T, -1)
            vals = vals.squeeze(1).reshape(B, T, -1)

            ld = layer_data[layer]
            image_heaps = ld["image_heaps"]
            text_heaps  = ld["text_heaps"]
            hook_counts = ld["hook_counts"]

            for fi in feat_ids:
                mask = (idxs == fi)
                feat_acts = (vals * mask.float()).sum(dim=-1)  # (B, T)
                above = feat_acts > args.threshold
                hook_counts[fi] += above.sum().item()

                for b in range(B):
                    ip = img_positions[b].item()
                    global_idx = global_idxs[b]
                    acts_b = feat_acts[b]                      # (T,)

                    # — Image patches —
                    if ip >= 0:
                        img_acts = acts_b[ip : ip + N_IMAGE_PATCHES]
                        n_img = (img_acts > args.threshold).sum().item()
                        if n_img > 0:
                            k_img = min(args.top_images, n_img)
                            topk_v, topk_i = img_acts.topk(k_img)
                            for v, pidx in zip(topk_v, topk_i):
                                if v.item() < args.threshold:
                                    break
                                image_heaps[fi].push(
                                    v.item(),
                                    ImageRecord(v.item(), global_idx,
                                                pidx.item(), hp),
                                )

                    # — Text tokens —
                    text_ranges: List[Tuple[int, int]] = []
                    if ip >= 0:
                        if ip > 0:
                            text_ranges.append((0, ip))
                        if ip + N_IMAGE_PATCHES < T:
                            text_ranges.append((ip + N_IMAGE_PATCHES, T))
                    else:
                        text_ranges.append((0, T))

                    for rng_start, rng_end in text_ranges:
                        txt_acts = acts_b[rng_start:rng_end]
                        n_txt = (txt_acts > args.threshold).sum().item()
                        if n_txt == 0:
                            continue
                        k_txt = min(args.top_texts, n_txt)
                        topk_v, topk_i = txt_acts.topk(k_txt)

                        for v, rel_i in zip(topk_v, topk_i):
                            if v.item() < args.threshold:
                                break
                            exp_pos = rng_start + rel_i.item()
                            orig_pos = expanded_to_original(exp_pos, ip)
                            if orig_pos < 0:
                                continue
                            if (orig_pos < input_ids_cpu.shape[1]
                                    and input_ids_cpu[b, orig_pos].item() == IMAGE_TOKEN_ID):
                                continue

                            seq_len = input_ids_cpu.shape[1]
                            ctx_s = max(0, orig_pos - args.buffer)
                            ctx_e = min(seq_len, orig_pos + args.buffer + 1)
                            if ip >= 0:
                                if orig_pos < ip:
                                    ctx_e = min(ctx_e, ip)
                                elif orig_pos > ip:
                                    ctx_s = max(ctx_s, ip + 1)
                            ctx_ids = input_ids_cpu[b, ctx_s:ctx_e].tolist()

                            ctx_acts_list: List[float] = []
                            for op in range(ctx_s, ctx_e):
                                ep = original_to_expanded(op, ip)
                                if ep < 0 or ep >= T:
                                    ctx_acts_list.append(0.0)
                                else:
                                    ctx_acts_list.append(feat_acts[b, ep].item())

                            text_heaps[fi].push(
                                v.item(),
                                TextRecord(
                                    v.item(), global_idx, hp,
                                    ctx_ids, ctx_acts_list,
                                    orig_pos - ctx_s,
                                ),
                            )

            del flat, idxs, vals

        del cache

    # Build return dict
    results: Dict[int, dict] = {}
    for layer, (feat_ids, feat_wts) in layer_probe_info.items():
        ld = layer_data[layer]
        results[layer] = {
            "feature_ids":   feat_ids,
            "probe_weights": feat_wts,
            "image_heaps":   ld["image_heaps"],
            "text_heaps":    ld["text_heaps"],
            "hook_counts":   ld["hook_counts"],
            "hook_point":    hook_points[layer],
        }
    return results


# ─────────────────────────────────────────────────────────────────────────────
# Interactive HTML builder
# ─────────────────────────────────────────────────────────────────────────────

_CSS = """
<style>
  * { box-sizing: border-box; }
  body {
    font-family: 'Segoe UI', Arial, sans-serif;
    background: #f0f2f5;
    margin: 0; padding: 0;
  }
  h1 {
    background: #2c3e50; color: #fff;
    margin: 0; padding: 18px 28px;
    font-size: 1.25rem;
  }
  .container { max-width: 1280px; margin: auto; padding: 20px; }

  /* ── Layer accordion ───────────────────────────────────────────────── */
  .layer-block {
    background: #fff;
    border: 1px solid #cdd3da;
    border-radius: 8px;
    margin-bottom: 14px;
    overflow: hidden;
  }
  .layer-header {
    display: flex; align-items: center; gap: 12px;
    background: #2c3e50; color: #fff;
    padding: 12px 18px; cursor: pointer;
    user-select: none;
  }
  .layer-header:hover { background: #3d5166; }
  .layer-header .arrow { transition: transform .2s; font-size: 0.85rem; }
  .layer-header.open .arrow { transform: rotate(90deg); }
  .layer-header .layer-title { font-size: 1rem; font-weight: 600; }
  .layer-header .layer-meta { font-size: 0.8rem; color: #aec6e3; margin-left: auto; }
  .layer-body { display: none; padding: 14px; }
  .layer-body.open { display: block; }

  /* ── Feature list ──────────────────────────────────────────────────── */
  .feature-list { display: flex; flex-direction: column; gap: 10px; }
  .feature-block { border: 1px solid #dde1e7; border-radius: 6px; overflow: hidden; }
  .feature-header {
    display: flex; align-items: center; gap: 10px;
    background: #f7f8fa; padding: 9px 14px;
    cursor: pointer; user-select: none;
  }
  .feature-header:hover { background: #edf0f4; }
  .feature-header .arrow { font-size: 0.75rem; transition: transform .2s; }
  .feature-header.open .arrow { transform: rotate(90deg); }
  .feat-id {
    font-weight: 700; font-size: 0.95rem; color: #2c3e50;
    background: #dce9f7; border-radius: 4px; padding: 2px 8px;
  }
  .feat-rank { font-size: 0.78rem; color: #888; }
  .feat-weight {
    margin-left: auto; font-size: 0.82rem; color: #155a8a;
    background: #e8f4fd; border-radius: 4px; padding: 2px 8px;
  }
  .feat-acts { font-size: 0.78rem; color: #555; }
  .feature-body { display: none; padding: 12px 14px; background: #fefefe; }
  .feature-body.open { display: block; }

  /* ── Image / text sections ─────────────────────────────────────────── */
  .section-title {
    font-size: 0.88rem; font-weight: 600; color: #34495e;
    margin: 10px 0 6px;
    padding-bottom: 4px; border-bottom: 1px solid #e2e6ea;
  }
  .image-card {
    border: 1px solid #ddd; border-radius: 5px;
    padding: 8px; margin: 6px 0; background: #fff;
  }
  .image-card img { max-height: 240px; }
  .text-card {
    border: 1px solid #ddd; border-radius: 5px;
    padding: 7px 10px; margin: 5px 0; background: #fff;
    display: flex; align-items: flex-start; gap: 8px;
  }
  .text-rank { font-weight: bold; color: #888; min-width: 30px;
               padding-top: 2px; font-size: 0.82rem; }
  .text-context { font-family: monospace; font-size: 0.88rem;
                  line-height: 1.6; word-break: break-word; }
  .meta { font-size: 0.75rem; color: #777; margin-top: 3px; }
</style>
"""

_JS = """
<script>
function toggle(bodyId, headerEl) {
  const body = document.getElementById(bodyId);
  const open = body.classList.toggle('open');
  headerEl.classList.toggle('open', open);
}
</script>
"""


def _feature_html(
    layer_idx: int,
    rank: int,
    feat_id: int,
    probe_weight: float,
    image_records,
    text_records,
    total_act: int,
    dataset,
    processor,
) -> str:
    body_id = f"feat_L{layer_idx}_F{feat_id}"

    img_html = build_image_html(image_records, dataset)
    txt_html = build_text_html(text_records, processor)

    # Replace raw <h2> from helpers with styled section titles
    img_html = img_html.replace(
        "<h2>Top Activating Image Patches</h2>",
        '<div class="section-title">Top Activating Image Patches</div>', 1,
    )
    txt_html = txt_html.replace(
        "<h2>Top Activating Text Tokens</h2>",
        '<div class="section-title">Top Activating Text Tokens</div>', 1,
    )

    return f"""
<div class="feature-block">
  <div class="feature-header" onclick="toggle('{body_id}', this)">
    <span class="arrow">&#9658;</span>
    <span class="feat-rank">#{rank}</span>
    <span class="feat-id">Feature {feat_id}</span>
    <span class="feat-weight">probe w = {probe_weight:+.4f}</span>
    <span class="feat-acts">&nbsp;activations: {total_act}</span>
  </div>
  <div class="feature-body" id="{body_id}">
    {img_html}
    {txt_html}
  </div>
</div>
"""


def build_interactive_html(
    layer_results: Dict[int, dict],
    dataset,
    processor,
    probe_dir: str,
    data_mode: str,
) -> str:
    layer_blocks: List[str] = []

    for layer_idx in sorted(layer_results.keys()):
        res          = layer_results[layer_idx]
        feat_ids     = res["feature_ids"]
        probe_wts    = res["probe_weights"]
        image_heaps  = res["image_heaps"]
        text_heaps   = res["text_heaps"]
        hook_counts  = res["hook_counts"]
        hook_point   = res["hook_point"]

        layer_body_id = f"layer_{layer_idx}_body"

        feature_blocks: List[str] = []
        for rank, (fi, pw) in enumerate(zip(feat_ids, probe_wts), start=1):
            img_recs = image_heaps[fi].sorted_records() if fi in image_heaps else []
            txt_recs = text_heaps[fi].sorted_records()  if fi in text_heaps  else []
            total    = hook_counts.get(fi, 0)
            feature_blocks.append(
                _feature_html(layer_idx, rank, fi, pw, img_recs, txt_recs,
                               total, dataset, processor)
            )

        n_feats    = len(feat_ids)
        total_acts = sum(hook_counts.get(fi, 0) for fi in feat_ids)

        layer_blocks.append(f"""
<div class="layer-block">
  <div class="layer-header" onclick="toggle('{layer_body_id}', this)">
    <span class="arrow">&#9658;</span>
    <span class="layer-title">Layer {layer_idx}</span>
    <span class="layer-meta">
      {hook_point} &nbsp;|&nbsp;
      {n_feats} features &nbsp;|&nbsp;
      {total_acts} total activations
    </span>
  </div>
  <div class="layer-body" id="{layer_body_id}">
    <div class="feature-list">{"".join(feature_blocks)}</div>
  </div>
</div>
""")

    return f"""<!DOCTYPE html>
<html lang="en">
<head>
  <meta charset="UTF-8">
  <title>Probe Feature Visualizer — Multi-Layer</title>
  {_CSS}
</head>
<body>
  <h1>Probe Feature Visualizer &mdash; Multi-Layer SAE Activations</h1>
  <div class="container">
    <p style="color:#555;font-size:0.88rem;">
      Probe dir: <code>{probe_dir}</code> &nbsp;|&nbsp; data mode: <code>{data_mode}</code><br>
      Click a <strong>Layer</strong> to expand its top probe features.
      Click a <strong>Feature</strong> to see top-activating image patches and text tokens.
    </p>
    {"".join(layer_blocks)}
  </div>
  {_JS}
</body>
</html>"""


# ─────────────────────────────────────────────────────────────────────────────
# Main
# ─────────────────────────────────────────────────────────────────────────────

def main():
    ap = argparse.ArgumentParser(
        description="Multi-layer probe feature visualizer for LLaVA.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )

    # ── Data mode ────────────────────────────────────────────────────────────
    ap.add_argument(
        "--data_mode", required=True, choices=["toilet", "cc3m", "coco", "folder"],
        help=(
            "toilet: pbcong/bathroom-toilet positives + CC3M negatives (needs --image_folder). "
            "cc3m:   full CC3M via --hf_dataset + --local_val_path. "
            "coco:   COCO via --hf_dataset + --local_val_path. "
            "folder: plain image folder via --data_dir."
        ),
    )

    # ── Model / SAE ──────────────────────────────────────────────────────────
    ap.add_argument("--sae_ckpt",   required=True)
    ap.add_argument("--model_name", default="llava-hf/llava-1.5-7b-hf")
    ap.add_argument("--device_id",  type=int, default=0)
    ap.add_argument("--dtype",      default="float16")

    # ── Probe ────────────────────────────────────────────────────────────────
    ap.add_argument("--probe_dir",       required=True,
                    help="Dir containing probe_model.language_model.layers.*.pt files.")
    ap.add_argument("--probe_input_dim", type=int, default=65536,
                    help="SAE width (LinearProbe input_dim).")
    ap.add_argument("--layers",          type=int, nargs="+", default=list(range(7)),
                    help="LLM layers to process.")
    ap.add_argument("--top_probe_k",     type=int, default=10,
                    help="Top-k probe-weight features to visualize per layer.")

    # ── Data — toilet ────────────────────────────────────────────────────────
    ap.add_argument("--image_folder",  default=None,
                    help="[toilet] Local CC3M image folder (positives + negative pool).")
    ap.add_argument("--object_mode",   default="toilet", choices=["toilet", "bathroom", "both"],
                    help="[toilet] Positive class: toilet==1, bathroom==1, or both (toilet==1 OR bathroom==1).")
    ap.add_argument("--num_negatives", type=int, default=10000,
                    help="[toilet] Number of random CC3M negatives to include.")

    # ── Caption mode ─────────────────────────────────────────────────────────
    ap.add_argument(
        "--caption_mode", default="generated", choices=["generated", "caption"],
        help=(
            "generated: model generates caption first; teacher-forced pass collects acts. "
            "caption:   use stored caption (CC3M txt / COCO sentences / "
            "pbcong/bathroom-toilet caption field)."
        ),
    )
    ap.add_argument("--max_new_tokens", type=int, default=128,
                    help="Max tokens to generate per image (only with --caption_mode generated).")

    # ── Data — CC3M / COCO ───────────────────────────────────────────────────
    ap.add_argument("--hf_dataset",    default=None,
                    help="HF dataset path (CC3M / COCO / also for toilet-negative captions).")
    ap.add_argument("--local_val_path", default=None,
                    help="[cc3m/coco] Local image root for the HF dataset split.")
    ap.add_argument("--data_dir",       default=None,
                    help="[folder] Plain image folder.")
    ap.add_argument("--split",          default="train",
                    help="HF dataset split.")

    # ── Common data ──────────────────────────────────────────────────────────
    ap.add_argument("--num_workers", type=int, default=4)

    # ── Processing ───────────────────────────────────────────────────────────
    ap.add_argument("--batch_size",  type=int,   default=4)
    ap.add_argument("--sae_batch",   type=int,   default=4096)
    ap.add_argument("--threshold",   type=float, default=1e-3)
    ap.add_argument("--max_batches", type=int,   default=None)

    # ── Visualisation ────────────────────────────────────────────────────────
    ap.add_argument("--output_dir",  default="outputs/probe_features")
    ap.add_argument("--top_images",  type=int, default=10)
    ap.add_argument("--top_texts",   type=int, default=10)
    ap.add_argument("--buffer",      type=int, default=10)

    args = ap.parse_args()

    # ── Validate data-mode args ───────────────────────────────────────────────
    if args.data_mode == "toilet" and not args.image_folder:
        ap.error("--data_mode toilet requires --image_folder.")
    if args.data_mode in ("cc3m", "coco") and not args.hf_dataset:
        ap.error(f"--data_mode {args.data_mode} requires --hf_dataset (+ --local_val_path).")
    if args.data_mode == "folder" and not args.data_dir:
        ap.error("--data_mode folder requires --data_dir.")

    # ── Distributed setup ────────────────────────────────────────────────────
    rank, world_size, local_rank = setup_distributed()
    is_distributed = world_size > 1

    device = torch.device(
        f"cuda:{local_rank}" if is_distributed
        else f"cuda:{args.device_id}" if torch.cuda.is_available() else "cpu"
    )
    dtype = str_to_torch_dtype(args.dtype)

    if rank == 0:
        print(f"Device: {device}  |  world_size: {world_size}  |  data_mode: {args.data_mode}")

    # ── Load SAE once ────────────────────────────────────────────────────────
    sae = load_sae_model(args.sae_ckpt, model_type="llava", hook_type="text", device=device)
    sae.eval()

    # ── Load LLaVA model + processor once ────────────────────────────────────
    model = HookedSAELlavaConditionalGeneration.from_pretrained(args.model_name)
    model.to(device, dtype=dtype)
    model.eval()
    processor = LlavaProcessor.from_pretrained(args.model_name)

    # ── Build probe feature dict ──────────────────────────────────────────────
    probe_info: Dict[int, Tuple[List[int], List[float]]] = {}
    for layer in args.layers:
        feat_ids, feat_wts = load_probe_top_features(
            args.probe_dir, layer, args.top_probe_k, args.probe_input_dim,
        )
        probe_info[layer] = (feat_ids, feat_wts)
        if rank == 0:
            print(f"  Layer {layer:2d} top features: {feat_ids}")

    # ── Build dataloader ────────────────────────────────────────────────────────
    dataset, dataloader = create_dataloader(args, processor, rank, world_size)
    if rank == 0:
        print(f"Dataset: {len(dataset)} samples, {len(dataloader)} batches/rank.")

    # ── Single multi-layer forward pass ────────────────────────────────────────
    raw_results = multi_layer_pass(
        model, sae, processor, dataloader, device, probe_info, args,
    )

    # ── DDP gather (all layers) ──────────────────────────────────────────────
    layer_results: Dict[int, dict] = {}

    for layer in args.layers:
        res = raw_results[layer]
        feat_ids    = res["feature_ids"]
        image_heaps = res["image_heaps"]
        text_heaps  = res["text_heaps"]
        hook_counts = res["hook_counts"]

        if is_distributed:
            output_dir = Path(args.output_dir)
            tmp_dir    = output_dir / f".ddp_tmp_L{layer}"
            if rank == 0:
                tmp_dir.mkdir(parents=True, exist_ok=True)
            dist.barrier()

            with open(tmp_dir / f"rank_{rank}.pkl", "wb") as f:
                pickle.dump({
                    "image_heaps": {fi: image_heaps[fi].sorted_records() for fi in feat_ids},
                    "text_heaps":  {fi: text_heaps[fi].sorted_records()  for fi in feat_ids},
                    "hook_counts": hook_counts,
                }, f)
            dist.barrier()

            if rank == 0:
                merged_img: Dict[int, TopKHeap] = {fi: TopKHeap(args.top_images) for fi in feat_ids}
                merged_txt: Dict[int, TopKHeap] = {fi: TopKHeap(args.top_texts)  for fi in feat_ids}
                merged_cnt: Dict[int, int]       = {fi: 0 for fi in feat_ids}

                for r in range(world_size):
                    with open(tmp_dir / f"rank_{r}.pkl", "rb") as f:
                        data = pickle.load(f)
                    for fi in feat_ids:
                        for rec in data["image_heaps"][fi]:
                            merged_img[fi].push(rec.value, rec)
                        for rec in data["text_heaps"][fi]:
                            merged_txt[fi].push(rec.value, rec)
                        merged_cnt[fi] += data["hook_counts"][fi]

                image_heaps = merged_img
                text_heaps  = merged_txt
                hook_counts = merged_cnt
                shutil.rmtree(tmp_dir)

        if rank == 0:
            layer_results[layer] = {
                "feature_ids":   feat_ids,
                "probe_weights": res["probe_weights"],
                "image_heaps":   image_heaps,
                "text_heaps":    text_heaps,
                "hook_counts":   hook_counts,
                "hook_point":    res["hook_point"],
            }
            for fi in feat_ids:
                print(f"  feature {fi}: {hook_counts.get(fi, 0)} activations")

    # ── Generate HTML (rank 0 only) ──────────────────────────────────────────
    if rank == 0:
        output_dir = Path(args.output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)

        print("\nBuilding interactive HTML …")
        html = build_interactive_html(
            layer_results, dataset, processor, args.probe_dir, args.data_mode,
        )
        out_path = output_dir / "probe_features.html"
        out_path.write_text(html, encoding="utf-8")
        print(f"Done. Open: {out_path}")

        summary = {}
        for layer, res in layer_results.items():
            summary[f"layer_{layer}"] = {
                str(fi): {
                    "probe_weight":      pw,
                    "total_activations": res["hook_counts"].get(fi, 0),
                }
                for fi, pw in zip(res["feature_ids"], res["probe_weights"])
            }
        with open(output_dir / "summary.json", "w") as f:
            json.dump(summary, f, indent=2)

    cleanup_distributed()


if __name__ == "__main__":
    main()