clean initial state

.gitattributes

@@ -0,0 +1,41 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.jpeg filter=lfs diff=lfs merge=lfs -text
+*.docx filter=lfs diff=lfs merge=lfs -text
+*.csv filter=lfs diff=lfs merge=lfs -text
+*.log filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,3 @@
+__pycache__/
+*.pyc
+logs/

README.md

@@ -0,0 +1,41 @@
+---
+title: AttrLLM Medical Attribution
+emoji: 🏥
+colorFrom: blue
+colorTo: indigo
+sdk: gradio
+sdk_version: "5.0.0"
+python_version: "3.10"
+app_file: app.py
+pinned: false
+---
+
+# AttrLLM — Medical Image Attribution
+
+Interactive visualization of **feature-level attributions** for medical images,
+computed via game-theoretic methods (Shapley values, Banzhaf indices).
+
+## What's in this Space
+
+Three precomputed attribution tabs — no GPU required:
+
+| Tab | Dataset | Examples | Methods |
+|---|---|---|---|
+| **Medical Benchmark** | PubMedVision | 5 curated (CT, MRI, Microscopy) | BiomedCLIP + LLaVA-Med (UnSAM) |
+| **MIMIC-CXR** | MIMIC-CXR-JPG | 10 chest X-rays (10 pathologies) | BiomedCLIP + LLaVA-Med (UnSAM) |
+| **MS-COCO** | MS-COCO | 5 examples | CLIP cross-modal + interactive view |
+
+Each example shows:
+- **Segment overlay** — image regions colored by Shapley attribution strength
+- **Token bar chart** — which caption words the model finds most diagnostic
+- **Interactive cross-modal view** — click a region to see linked tokens, and vice versa
+- **Side-by-side method comparison** — BiomedCLIP vs LLaVA-Med
+
+## Attribution Methods
+
+- **BiomedCLIP cross-modal**: UnSAM image segments × PubMedBERT tokens, scored via BiomedCLIP cosine similarity, Shapley values via ProxySPEX
+- **LLaVA-Med (UnSAM)**: UnSAM segments, scored by log P(caption | masked image) using LLaVA-Med-7B, Shapley values via ProxySPEX
+
+## Citation
+
+Part of the AttrLLM capstone project (NCSA Delta HPC, UIUC).

app.py

@@ -0,0 +1,5 @@
+from visualization.app import launch_app
+
+
+if __name__ == "__main__":
+    launch_app()

attribution/attribution.py

@@ -0,0 +1,1262 @@
+import math
+import os
+import time
+import logging
+from dataclasses import asdict, dataclass
+from functools import lru_cache
+from typing import Any, Dict, List, Sequence, Tuple
+
+from attribution.masker import MaskingInfo, get_masker, mask_text
+from attribution.proxyspex import run_proxyspex
+from attribution.utils import (
+    AttributionResult,
+    banzhaf_interactions,
+    convert_to_heatmap_data,
+    influence_interactions,
+    mobius_to_banzhaf,
+    mobius_to_fourier,
+    mobius_to_influence,
+    mobius_to_shapley,
+    shapley_interactions,
+)
+from loader.data import get_example_by_id
+from loader.perplexity import score_continuation, score_continuation_batch
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+try:
+    from loader.models import BaseModel
+except ImportError:  # pragma: no cover - fallback type
+    BaseModel = Any  # type: ignore
+
+
+SUPPORTED_METHODS = {"shapley", "banzhaf", "influence"}
+SUPPORTED_SCALARIZERS = {
+    "semantic_similarity",
+    "logprob",
+    "jointprob",
+    "geomean_jointprob",
+    "half_simlog",
+}
+MAX_MOBIUS_TERMS = 1000
+DEFAULT_SIM_MODEL = "all-MiniLM-L6-v2"
+# Pre-downloaded local copy (see scripts/download_st_model.py)
+_LOCAL_ST_MODEL_DIR = os.path.join(
+    os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
+    "models", "all_MiniLM_L6_v2",
+)
+MIN_LOGPROB_FALLBACK = -1e6
+EXP_CLAMP_MIN = -200.0
+EXP_CLAMP_MAX = 50.0
+
+_ST_MODEL = None
+
+
+def _get_st_model():
+    global _ST_MODEL
+    if _ST_MODEL is None:
+        try:
+            from sentence_transformers import SentenceTransformer
+        except ImportError as exc:  # pragma: no cover
+            raise RuntimeError(
+                "sentence-transformers not installed; required for semantic_similarity scalarizer."
+            ) from exc
+        # Prefer the local snapshot if it exists.
+        if os.path.isdir(_LOCAL_ST_MODEL_DIR):
+            model_path = _LOCAL_ST_MODEL_DIR
+            logger.info("Loading SentenceTransformer from local dir: %s", model_path)
+        else:
+            model_path = DEFAULT_SIM_MODEL
+            logger.info("Local ST model not found; loading from HuggingFace hub: %s", model_path)
+        _ST_MODEL = SentenceTransformer(model_path)
+    return _ST_MODEL
+
+
+@dataclass
+class AttributorConfig:
+    """Configuration for the attribution pipeline."""
+
+    mask_level: str = "sentence"
+    method: str = "shapley"
+    interaction_order: int = 3
+    top_k_interactions: int = 10
+    random_seed: int = 0
+    scalarizer: str = "logprob"  # "semantic_similarity" | "logprob" | "jointprob" | "geomean_jointprob" | "half_simlog"
+    embedding_model: str | None = None
+
+    def __post_init__(self) -> None:
+        normalized = self.method.lower()
+        if normalized not in SUPPORTED_METHODS:
+            raise ValueError(f"Unsupported attribution method: {self.method}")
+        self.method = normalized
+        self.scalarizer = (self.scalarizer or "logprob").lower()
+        if self.scalarizer not in SUPPORTED_SCALARIZERS:
+            raise ValueError(f"Unsupported scalarizer: {self.scalarizer}")
+
+
+class Attributor:
+    """Encapsulates the full text attribution workflow."""
+
+    def __init__(self, model: BaseModel, config: AttributorConfig | None = None) -> None:
+        self.model = model
+        self.config = config or AttributorConfig()
+
+    def _mask_context(self, context: str) -> Tuple[List[str], MaskingInfo]:
+        """Split the context into maskable features and capture span metadata."""
+        masker_fn = get_masker(self.config.mask_level)
+        features, info = mask_text(context, masker_fn)
+        info.mask_level = self.config.mask_level
+        
+        return features, info
+
+    def _build_set_function(
+        self,
+        features: Sequence[str],
+        answer: str,
+        reference_answer: str,
+        prompt: str,
+    ):
+        """
+        Build the cooperative game f(S) with a configurable scalarizer.
+        """
+        max_new_tokens_score = int(os.getenv("ATTRLLM_MAX_NEW_TOKENS_SCORE", "32"))
+        batch_size = int(os.getenv("VLLM_SCORE_BATCH_SIZE",
+                                   os.getenv("ATTRLLM_VLLM_BATCH_SIZE", "128")))
+        n_features = len(features)
+        scalarizer = self.config.scalarizer
+        prompt_cache: Dict[Tuple[str, str, int], Dict[str, Any]] = {}
+        error_cache: Dict[Tuple[int, ...], str | None] = {}
+        min_logprob_seen: Dict[str, float | None] = {"value": None}
+        answer = (answer or "").strip()
+        reference_answer = (reference_answer or "").strip()
+
+        def _safe_exp(log_value: float) -> float:
+            clamped = max(EXP_CLAMP_MIN, min(EXP_CLAMP_MAX, log_value))
+            return float(math.exp(clamped))
+
+        def _sanitize_score_value(
+            score_value: float,
+            *,
+            masked_context: str,
+            coalition_size: int,
+            raw_score: Any | None = None,
+        ) -> float:
+            if math.isfinite(score_value):
+                return float(score_value)
+            logger.warning(
+                "Non-finite scalarizer score; falling back to 0. "
+                "scalarizer=%s coalition_size=%d masked_len=%d raw_score=%s",
+                scalarizer,
+                coalition_size,
+                len(masked_context),
+                raw_score if raw_score is not None else score_value,
+            )
+            return 0.0
+
+        def _build_prompt(masked_context: str) -> str:
+            masked_context = masked_context or ""
+            p = prompt or ""
+            if p and masked_context:
+                return f"{p}\nContext: {masked_context}"
+            if p:
+                return p
+            return masked_context
+
+        full_context = " ".join(features).strip()
+        full_prompt = _build_prompt(full_context)
+        y_full = ""
+        y_len_tokens = 0
+        try:
+            if full_prompt:
+                outputs = self.model.infer(
+                    [full_prompt],
+                    max_new_tokens=max_new_tokens_score,
+                    temperature=0.0,
+                    top_p=1.0,
+                    logprobs=0,
+                )
+                if outputs:
+                    y_full = (outputs[0].text or "").strip()
+                    token_ids = outputs[0].token_ids or []
+                    y_len_tokens = len(token_ids) or len(y_full.split())
+        except Exception as exc:
+            logger.warning("Failed to generate unmasked output y: %s", exc)
+            y_full = ""
+            y_len_tokens = 0
+
+        if reference_answer:
+            scoring_target_source = "reference_answer"
+            scoring_target_text = reference_answer
+        elif answer:
+            scoring_target_source = "answer_input"
+            scoring_target_text = answer
+        else:
+            scoring_target_source = "model_output"
+            scoring_target_text = y_full
+
+        scoring_target_text = (scoring_target_text or "").strip()
+        if scoring_target_source == "model_output":
+            scoring_target_len_tokens = y_len_tokens
+        else:
+            scoring_target_len_tokens = len(scoring_target_text.split())
+
+        def _score_logprob(masked_context: str) -> Dict[str, Any]:
+            if not scoring_target_text:
+                return {
+                    "logprob": MIN_LOGPROB_FALLBACK,
+                    "total_logprob": MIN_LOGPROB_FALLBACK,
+                    "avg_nll": 0.0,
+                    "ppl": 1.0,
+                    "target_len": scoring_target_len_tokens,
+                    "generated": scoring_target_text,
+                    "error": "empty scoring target",
+                }
+            prompt_text = _build_prompt(masked_context)
+            cache_key = (prompt_text, scoring_target_text, max_new_tokens_score)
+            if cache_key in prompt_cache:
+                score = prompt_cache[cache_key]
+            else:
+                score = score_continuation(
+                    self.model,
+                    prefix=prompt_text,
+                    continuation=scoring_target_text,
+                    max_new_tokens=max_new_tokens_score,
+                )
+                prompt_cache[cache_key] = score
+            avg_logprob = float(score.get("avg_logprob", float("-inf")))
+            target_len = int(score.get("target_len", 0) or 0)
+            if target_len <= 0:
+                target_len = scoring_target_len_tokens
+            seq_logprobs = score.get("sequence_logprobs", [])
+            if seq_logprobs:
+                total_logprob = float(seq_logprobs[0])
+            else:
+                total_logprob = avg_logprob * target_len if target_len else float("-inf")
+            error: str | None = None
+            if not math.isfinite(total_logprob):
+                total_logprob = MIN_LOGPROB_FALLBACK
+                error = "non-finite total_logprob"
+            if target_len > 0 and math.isfinite(total_logprob):
+                avg_nll = -total_logprob / target_len
+                ppl = _safe_exp(avg_nll)
+            else:
+                avg_nll = 0.0
+                ppl = 1.0
+            current_min = min_logprob_seen["value"]
+            if math.isfinite(total_logprob) and (current_min is None or total_logprob < current_min):
+                min_logprob_seen["value"] = total_logprob
+            return {
+                "logprob": total_logprob,
+                "total_logprob": total_logprob,
+                "avg_nll": avg_nll,
+                "ppl": ppl,
+                "target_len": target_len,
+                "generated": scoring_target_text,
+                "error": error,
+            }
+
+        def _logsumexp(a: float, b: float) -> float:
+            if not math.isfinite(a) and not math.isfinite(b):
+                return MIN_LOGPROB_FALLBACK
+            if not math.isfinite(a):
+                return b
+            if not math.isfinite(b):
+                return a
+            m = max(a, b)
+            return m + math.log(math.exp(a - m) + math.exp(b - m))
+
+        def _scalarize(masked_context: str) -> Dict[str, Any]:
+            if scalarizer == "semantic_similarity":
+                prompt_text = _build_prompt(masked_context)
+                try:
+                    outputs = self.model.infer(
+                        [prompt_text],
+                        max_new_tokens=16,
+                        temperature=0.0,
+                        top_p=1.0,
+                        logprobs=0,
+                    )
+                    y_s = (outputs[0].text or "").strip() if outputs else ""
+                except Exception as exc:
+                    logger.warning("Generation failed for semantic_similarity: %s", exc)
+                    y_s = ""
+
+                if not scoring_target_text or not y_s.strip():
+                    logger.warning(
+                        "semantic_similarity fallback to 0. scalarizer=%s masked_len=%d y_len=%d",
+                        scalarizer,
+                        len(masked_context),
+                        len(y_s),
+                    )
+                    return {
+                        "score": 0.0,
+                        "total_logprob": 0.0,
+                        "avg_nll": 0.0,
+                        "ppl": 1.0,
+                        "target_len": scoring_target_len_tokens,
+                        "generated": y_s,
+                        "error": "empty scoring target or y_s",
+                    }
+
+                try:
+                    model = _get_st_model()
+                    emb = model.encode([scoring_target_text, y_s])
+                    emb_arr = np.asarray(emb, dtype=np.float64)
+                    norm_target = float(np.linalg.norm(emb_arr[0])) if emb_arr.shape[0] > 0 else 0.0
+                    norm_masked = float(np.linalg.norm(emb_arr[1])) if emb_arr.shape[0] > 1 else 0.0
+                    if (
+                        emb_arr.shape[0] < 2
+                        or not np.isfinite(norm_target)
+                        or not np.isfinite(norm_masked)
+                        or norm_target <= 0.0
+                        or norm_masked <= 0.0
+                    ):
+                        sim = 0.0
+                        error = "zero-norm or invalid embedding"
+                    else:
+                        sim = float(model.similarity(emb_arr, emb_arr)[0, 1])
+                        error = None
+                    if not np.isfinite(sim):
+                        logger.warning(
+                            "semantic_similarity non-finite score; fallback to 0. "
+                            "masked_len=%d y_len=%d raw_score=%s",
+                            len(masked_context),
+                            len(y_s),
+                            sim,
+                        )
+                        sim = 0.0
+                        error = "non-finite semantic similarity"
+                except Exception as exc:
+                    logger.warning("semantic_similarity failed: %s", exc)
+                    sim = 0.0
+                    error = "semantic_similarity exception"
+
+                return {
+                    "score": sim,
+                    "total_logprob": sim,
+                    "avg_nll": 0.0,
+                    "ppl": 1.0,
+                    "target_len": scoring_target_len_tokens,
+                    "generated": y_s,
+                    "error": error,
+                }
+
+            logprob_details = _score_logprob(masked_context)
+            logprob = float(logprob_details.get("logprob", MIN_LOGPROB_FALLBACK))
+            if not math.isfinite(logprob):
+                logger.warning(
+                    "logprob fallback to %s. scalarizer=%s masked_len=%d raw_score=%s",
+                    MIN_LOGPROB_FALLBACK,
+                    scalarizer,
+                    len(masked_context),
+                    logprob,
+                )
+                logprob = MIN_LOGPROB_FALLBACK
+                logprob_details["error"] = "non-finite logprob"
+
+            if scalarizer == "logprob":
+                score_value = logprob
+            elif scalarizer == "jointprob":
+                score_value = _safe_exp(logprob)
+            elif scalarizer == "geomean_jointprob":
+                denom = max(scoring_target_len_tokens, 1)
+                score_value = _safe_exp(logprob / denom)
+            elif scalarizer == "half_simlog":
+                c_val = min_logprob_seen["value"]
+                if c_val is None or not math.isfinite(c_val):
+                    c_val = logprob
+                score_value = _logsumexp(logprob, c_val)
+            else:
+                raise ValueError(f"Unsupported scalarizer: {scalarizer}")
+
+            logprob_details.update({"score": float(score_value)})
+            return logprob_details
+
+        @lru_cache(maxsize=4096)
+        def cached(selection_tuple: Tuple[int, ...]) -> Tuple[float, float, float, float, int]:
+            masked_context = " ".join(
+                features[idx] for idx, active in enumerate(selection_tuple) if active
+            ).strip()
+            details = _scalarize(masked_context)
+            score_value = float(details.get("score", 0.0))
+            total_logprob = float(details.get("total_logprob", score_value))
+            avg_nll = float(details.get("avg_nll", 0.0))
+            ppl = float(details.get("ppl", 1.0))
+            target_len = int(details.get("target_len", 0) or 0)
+            generated_len = len(details.get("generated", "") or "")
+            coalition_size = sum(int(x) for x in selection_tuple)
+            error_cache[selection_tuple] = details.get("error")
+            if error_cache[selection_tuple]:
+                logger.warning(
+                    "Scalarizer fallback: scalarizer=%s coalition_size=%d masked_len=%d y_len=%d error=%s raw_score=%s",
+                    scalarizer,
+                    coalition_size,
+                    len(masked_context),
+                    generated_len,
+                    error_cache[selection_tuple],
+                    details.get("score"),
+                )
+
+            score_value = _sanitize_score_value(
+                score_value,
+                masked_context=masked_context,
+                coalition_size=coalition_size,
+                raw_score=details.get("score"),
+            )
+            if not math.isfinite(total_logprob):
+                total_logprob = 0.0
+                avg_nll = 0.0
+                ppl = 1.0
+            return score_value, total_logprob, avg_nll, ppl, target_len
+
+        def set_function(selection) -> float:
+            if len(selection) != n_features:
+                raise ValueError("Selection length must match number of features.")
+            selection_tuple = tuple(int(x) for x in selection)
+            return cached(selection_tuple)[0]
+
+        def score_details(selection) -> Dict[str, Any]:
+            if len(selection) != n_features:
+                raise ValueError("Selection length must match number of features.")
+            selection_tuple = tuple(int(x) for x in selection)
+            score_value, total_logprob, avg_nll, ppl, target_len = cached(selection_tuple)
+            return {
+                "score": score_value,
+                "total_logprob": total_logprob,
+                "logprob": total_logprob,
+                "avg_nll": avg_nll,
+                "ppl": ppl,
+                "target_len": target_len,
+                "scalarizer": scalarizer,
+                "error": error_cache.get(selection_tuple),
+            }
+
+        set_function.score_details = score_details  # type: ignore[attr-defined]
+
+        def _batch_semantic_similarity(
+            selection_tuples: List[Tuple[int, ...]],
+            masked_contexts: List[str],
+        ) -> List[float]:
+            """
+            Batch semantic similarity computation:
+            1. Generate y_s for all masked contexts in batches (using LLM)
+            2. Encode all texts (target + all y_s) in one batch with SentenceTransformer
+            3. Compute cosine similarities
+            """
+            n_samples = len(selection_tuples)
+            if n_samples == 0:
+                return []
+
+            # Fallback if target is empty
+            if not scoring_target_text:
+                return [0.0] * n_samples
+
+            # Step 1: Generate y_s for all masked contexts in batches
+            prompts = [_build_prompt(mc) for mc in masked_contexts]
+            y_s_list: List[str] = []
+
+            for start in range(0, len(prompts), batch_size):
+                chunk = prompts[start : start + batch_size]
+                try:
+                    outputs = self.model.infer(
+                        chunk,
+                        max_new_tokens=16,
+                        temperature=0.0,
+                        top_p=1.0,
+                        logprobs=0,
+                    )
+                    for out in outputs:
+                        y_s_list.append((out.text or "").strip() if out else "")
+                except Exception as exc:
+                    logger.warning("Batch generation failed for semantic_similarity: %s", exc)
+                    # Fill with empty strings for this batch
+                    y_s_list.extend([""] * len(chunk))
+
+            # Step 2: Get the sentence transformer model
+            try:
+                st_model = _get_st_model()
+            except Exception as exc:
+                logger.warning("Failed to load SentenceTransformer: %s", exc)
+                return [0.0] * n_samples
+
+            # Step 3: Encode all texts in one batch
+            # texts_to_encode = [scoring_target_text] + y_s_list (index 0 is the target)
+            texts_to_encode = [scoring_target_text] + y_s_list
+            try:
+                embeddings = st_model.encode(
+                    texts_to_encode,
+                    show_progress_bar=False,
+                    batch_size=batch_size,
+                )
+                emb_arr = np.asarray(embeddings, dtype=np.float64)
+                target_emb = emb_arr[0:1]  # Shape: (1, dim)
+                y_s_embs = emb_arr[1:]      # Shape: (n_samples, dim)
+            except Exception as exc:
+                logger.warning("Batch encoding failed for semantic_similarity: %s", exc)
+                return [0.0] * n_samples
+
+            target_norm = float(np.linalg.norm(target_emb[0])) if target_emb.shape[0] > 0 else 0.0
+            y_s_norms = np.linalg.norm(y_s_embs, axis=1) if y_s_embs.size else np.array([])
+            if not np.isfinite(target_norm) or target_norm <= 0.0:
+                logger.warning(
+                    "semantic_similarity fallback to 0 in batch due to invalid target embedding norm. "
+                    "scalarizer=%s",
+                    scalarizer,
+                )
+                for sel_tuple in selection_tuples:
+                    error_cache[sel_tuple] = "invalid target embedding"
+                return [0.0] * n_samples
+
+            # Step 4: Compute similarities
+            results: List[float] = []
+            try:
+                # Compute cosine similarity between target and each y_s
+                similarities = st_model.similarity(target_emb, y_s_embs)  # Shape: (1, n_samples)
+                for i, sel_tuple in enumerate(selection_tuples):
+                    y_s = y_s_list[i]
+                    if not y_s.strip():
+                        sim = 0.0
+                        error_cache[sel_tuple] = "empty y_s"
+                    elif (
+                        i >= len(y_s_norms)
+                        or not np.isfinite(float(y_s_norms[i]))
+                        or float(y_s_norms[i]) <= 0.0
+                    ):
+                        sim = 0.0
+                        error_cache[sel_tuple] = "invalid masked embedding"
+                    else:
+                        sim = float(similarities[0, i])
+                        if not np.isfinite(sim):
+                            sim = 0.0
+                            error_cache[sel_tuple] = "non-finite semantic similarity"
+                        else:
+                            error_cache[sel_tuple] = None
+                    coalition_size = sum(int(x) for x in sel_tuple)
+                    sim = _sanitize_score_value(
+                        sim,
+                        masked_context=masked_contexts[i],
+                        coalition_size=coalition_size,
+                        raw_score=sim,
+                    )
+                    results.append(sim)
+            except Exception as exc:
+                logger.warning("Similarity computation failed: %s", exc)
+                return [0.0] * n_samples
+
+            return results
+
+        def batch_set_function(selections: Sequence[Sequence[int]]) -> List[float]:
+            selection_tuples = [tuple(int(x) for x in sel) for sel in selections]
+            masked_contexts = [
+                " ".join(features[idx] for idx, active in enumerate(sel) if active).strip()
+                for sel in selection_tuples
+            ]
+
+            # Handle semantic_similarity with batching
+            if scalarizer == "semantic_similarity":
+                return _batch_semantic_similarity(selection_tuples, masked_contexts)
+
+            logprob_scalarizers = {
+                "logprob",
+                "jointprob",
+                "geomean_jointprob",
+                "half_simlog",
+            }
+            if scalarizer not in logprob_scalarizers:
+                return [cached(sel)[0] for sel in selection_tuples]
+
+            if not scoring_target_text:
+                fallback_logprob = MIN_LOGPROB_FALLBACK
+                outputs: List[float] = []
+                for sel_tuple, mc in zip(selection_tuples, masked_contexts):
+                    if scalarizer == "logprob":
+                        score_value = fallback_logprob
+                    elif scalarizer == "jointprob":
+                        score_value = _safe_exp(fallback_logprob)
+                    elif scalarizer == "geomean_jointprob":
+                        denom = max(scoring_target_len_tokens, 1)
+                        score_value = _safe_exp(fallback_logprob / denom)
+                    elif scalarizer == "half_simlog":
+                        score_value = _logsumexp(fallback_logprob, fallback_logprob)
+                    else:
+                        score_value = fallback_logprob
+                    coalition_size = sum(int(x) for x in sel_tuple)
+                    outputs.append(
+                        _sanitize_score_value(
+                            float(score_value),
+                            masked_context=mc,
+                            coalition_size=coalition_size,
+                            raw_score=score_value,
+                        )
+                    )
+                return outputs
+
+            unique_map: Dict[Tuple[str, str, int], List[int]] = {}
+            for idx, mc in enumerate(masked_contexts):
+                prompt_text = _build_prompt(mc)
+                key = (prompt_text, scoring_target_text, max_new_tokens_score)
+                unique_map.setdefault(key, []).append(idx)
+
+            to_fetch_keys = []
+            to_fetch_contexts = []
+            for key in unique_map.keys():
+                if key not in prompt_cache:
+                    to_fetch_keys.append(key)
+                    to_fetch_contexts.append(key[0])
+
+            logger.info(
+                "[batch_set_function] total=%d unique_prompts=%d new_to_fetch=%d cached=%d batch_size=%d",
+                len(selection_tuples),
+                len(unique_map),
+                len(to_fetch_contexts),
+                len(prompt_cache),
+                batch_size,
+            )
+
+            if to_fetch_contexts:
+                batch_scores: List[Dict[str, Any]] = []
+                for start in range(0, len(to_fetch_contexts), batch_size):
+                    chunk = to_fetch_contexts[start : start + batch_size]
+                    logger.info(
+                        "[batch_set_function] scoring chunk %d..%d of %d via vLLM",
+                        start, start + len(chunk), len(to_fetch_contexts),
+                    )
+                    scored = score_continuation_batch(
+                        self.model,
+                        prefixes=chunk,
+                        continuation=scoring_target_text,
+                        max_new_tokens=max_new_tokens_score,
+                        batch_size=batch_size,
+                    )
+                    batch_scores.extend(scored)
+                for key, score in zip(to_fetch_keys, batch_scores):
+                    prompt_cache[key] = score
+
+            outputs: List[float] = []
+            for sel_tuple, mc in zip(selection_tuples, masked_contexts):
+                prompt_text = _build_prompt(mc)
+                key = (prompt_text, scoring_target_text, max_new_tokens_score)
+                score = prompt_cache.get(key)
+                if not score:
+                    outputs.append(cached(sel_tuple)[0])
+                    continue
+                avg_logprob = float(score.get("avg_logprob", float("-inf")))
+                target_len = int(score.get("target_len", 0) or 0)
+                if target_len <= 0:
+                    target_len = scoring_target_len_tokens
+                seq_logprobs = score.get("sequence_logprobs", [])
+                if seq_logprobs:
+                    total_logprob = float(seq_logprobs[0])
+                else:
+                    total_logprob = avg_logprob * target_len if target_len else float("-inf")
+                if not math.isfinite(total_logprob):
+                    total_logprob = MIN_LOGPROB_FALLBACK
+                current_min = min_logprob_seen["value"]
+                if math.isfinite(total_logprob) and (current_min is None or total_logprob < current_min):
+                    min_logprob_seen["value"] = total_logprob
+
+                if scalarizer == "logprob":
+                    score_value = total_logprob
+                elif scalarizer == "jointprob":
+                    score_value = _safe_exp(total_logprob)
+                elif scalarizer == "geomean_jointprob":
+                    denom = max(scoring_target_len_tokens, 1)
+                    score_value = _safe_exp(total_logprob / denom)
+                elif scalarizer == "half_simlog":
+                    c_val = min_logprob_seen["value"]
+                    if c_val is None or not math.isfinite(c_val):
+                        c_val = total_logprob
+                    score_value = _logsumexp(total_logprob, c_val)
+                else:
+                    score_value = total_logprob
+                coalition_size = sum(int(x) for x in sel_tuple)
+                outputs.append(
+                    _sanitize_score_value(
+                        float(score_value),
+                        masked_context=mc,
+                        coalition_size=coalition_size,
+                        raw_score=score_value,
+                    )
+                )
+            return outputs
+
+        set_function.batch = batch_set_function  # type: ignore[attr-defined]
+
+        set_function.y_full = y_full  # type: ignore[attr-defined]
+        set_function.y_len_tokens = y_len_tokens  # type: ignore[attr-defined]
+        set_function.scoring_target_source = scoring_target_source  # type: ignore[attr-defined]
+        set_function.scoring_target_text = scoring_target_text  # type: ignore[attr-defined]
+        set_function.answer_received = answer  # type: ignore[attr-defined]
+        set_function.reference_answer_received = reference_answer  # type: ignore[attr-defined]
+        set_function.min_logprob_seen = min_logprob_seen  # type: ignore[attr-defined]
+        return set_function
+
+    def _compute_mobius(self, context: str, answer: str, reference_answer: str, prompt: str):
+        """Compute Mobius coefficients for the masked context game."""
+        features, masking_info = self._mask_context(context)
+        set_function = self._build_set_function(features, answer, reference_answer, prompt)
+        mobius = run_proxyspex(
+            set_function,
+            features,
+            max_order=self.config.interaction_order,
+            seed=self.config.random_seed,
+        )
+        return mobius, features, masking_info, set_function
+
+    def _index_to_tuple(self, idx: int) -> Tuple[int, ...]:
+        """Return a sparse coalition tuple with a single active index."""
+        return (idx,)
+
+    def _tuple_to_indices(self, loc_tuple: Tuple[int, ...]) -> List[int]:
+        """Return active feature indices from a sparse coalition tuple."""
+        return sorted(loc_tuple)
+
+    @staticmethod
+    def _assert_tuple_int_key(key: Any, object_name: str) -> None:
+        assert isinstance(
+            key, tuple
+        ), f"{object_name} key must be tuple; got type={type(key).__name__} value={key!r}"
+        assert all(
+            isinstance(i, int) for i in key
+        ), f"{object_name} key entries must be int; got value={key!r}"
+
+    def _validate_mobius_keys(self, mobius_obj: Any, object_name: str) -> Dict[Tuple[int, ...], float]:
+        if not isinstance(mobius_obj, dict):
+            message = (
+                f"{object_name} must be dict with tuple[int,...] keys; "
+                f"got type={type(mobius_obj).__name__}"
+            )
+            print(message)
+            raise TypeError(message)
+        for key in mobius_obj.keys():
+            try:
+                self._assert_tuple_int_key(key, object_name)
+            except AssertionError as exc:
+                message = (
+                    f"{object_name} contains invalid key. "
+                    f"type={type(key).__name__} value={key!r} error={exc}"
+                )
+                print(message)
+                raise AssertionError(message) from exc
+        return mobius_obj
+
+    def _validate_interaction_pairs(
+        self,
+        pairs_obj: Any,
+        object_name: str,
+    ) -> List[Tuple[Tuple[int, ...], float]]:
+        if not isinstance(pairs_obj, list):
+            message = (
+                f"{object_name} must be list of ((int,...), value) pairs; "
+                f"got type={type(pairs_obj).__name__}"
+            )
+            print(message)
+            raise TypeError(message)
+
+        validated: List[Tuple[Tuple[int, ...], float]] = []
+        for idx, item in enumerate(pairs_obj):
+            if not (isinstance(item, tuple) and len(item) == 2):
+                message = (
+                    f"{object_name}[{idx}] must be tuple(len=2); "
+                    f"got type={type(item).__name__} value={item!r}"
+                )
+                print(message)
+                raise TypeError(message)
+            key, value = item
+            try:
+                self._assert_tuple_int_key(key, f"{object_name}[{idx}]")
+            except AssertionError as exc:
+                message = (
+                    f"{object_name}[{idx}] has invalid interaction key. "
+                    f"type={type(key).__name__} value={key!r} error={exc}"
+                )
+                print(message)
+                raise AssertionError(message) from exc
+            validated.append((key, float(value)))
+        return validated
+
+    def _build_singletons(
+        self, values: Dict[Tuple[int, ...], float], features: Sequence[str], method: str
+    ) -> Dict[Tuple[int, ...], AttributionResult]:
+        """Wrap singleton values into AttributionResult objects."""
+        n = len(features)
+        results: Dict[Tuple[int, ...], AttributionResult] = {}
+        for idx in range(n):
+            loc = self._index_to_tuple(idx)
+            value = float(values.get(loc, 0.0))
+            results[loc] = AttributionResult(
+                feature=features[idx],
+                value=value,
+                method=method,
+                interaction_order=1,
+            )
+        return results
+
+    def _build_interactions(
+        self,
+        interactions: List[Tuple[Tuple[int, ...], float]],
+        features: Sequence[str],
+        method: str,
+    ) -> List[Tuple[Tuple[int, ...], AttributionResult]]:
+        """Wrap interaction scores into AttributionResult objects."""
+        payload: List[Tuple[Tuple[int, ...], AttributionResult]] = []
+        for loc, value in interactions:
+            indices = self._tuple_to_indices(loc)
+            if len(indices) <= 1:
+                continue
+            feature_names = tuple(features[i] for i in indices)
+            payload.append(
+                (
+                    loc,
+                    AttributionResult(
+                        feature=feature_names,
+                        value=float(value),
+                        method=method,
+                        interaction_order=len(indices),
+                    ),
+                )
+            )
+        return payload
+
+    def _trim_mobius(
+        self, mobius_dict: Dict[Tuple[int, ...], float]
+    ) -> Dict[Tuple[int, ...], float]:
+        """Keep only the largest Mobius coefficients by magnitude."""
+        if len(mobius_dict) <= MAX_MOBIUS_TERMS:
+            return mobius_dict
+        sorted_items = sorted(mobius_dict.items(), key=lambda x: abs(x[1]), reverse=True)
+        return dict(sorted_items[:MAX_MOBIUS_TERMS])
+
+    def _mobius_to_values_and_interactions(
+        self, mobius_dict: Dict[Tuple[int, ...], float], features: Sequence[str]
+    ):
+        """
+        Convert Mobius coefficients into singleton values and pairwise interactions.
+        """
+        mobius_dict = self._trim_mobius(mobius_dict)
+        shapley_values = self._build_singletons(
+            mobius_to_shapley(mobius_dict),
+            features,
+            "shapley",
+        )
+        banzhaf_values = self._build_singletons(
+            mobius_to_banzhaf(mobius_dict),
+            features,
+            "banzhaf",
+        )
+        influence_values = self._build_singletons(
+            mobius_to_influence(mobius_dict),
+            features,
+            "influence",
+        )
+        shapley_pairs = shapley_interactions(
+            mobius_dict,
+            order=2,
+            top_k=self.config.top_k_interactions,
+        )
+        banzhaf_pairs = banzhaf_interactions(
+            mobius_dict,
+            order=2,
+            top_k=self.config.top_k_interactions,
+        )
+        influence_pairs = influence_interactions(
+            mobius_dict,
+            order=2,
+            top_k=self.config.top_k_interactions,
+        )
+        
+        shapley_ints = self._build_interactions(shapley_pairs, features, "shapley")
+        banzhaf_ints = self._build_interactions(banzhaf_pairs, features, "banzhaf")
+        influence_ints = self._build_interactions(influence_pairs, features, "influence")
+
+        return {
+            "shapley": {
+                "values": shapley_values,
+                "interactions": shapley_ints,
+            },
+            "banzhaf": {
+                "values": banzhaf_values,
+                "interactions": banzhaf_ints,
+            },
+            "influence": {
+                "values": influence_values,
+                "interactions": influence_ints,
+            },
+        }
+
+    def _log_influence_debug(
+        self,
+        mobius_dict: Dict[Tuple[int, ...], float],
+        features: Sequence[str],
+        values: Dict[Tuple[int, ...], AttributionResult],
+        interactions: List[Tuple[Tuple[int, ...], AttributionResult]],
+    ) -> None:
+        """Print influence internals for debugging in terminal logs."""
+        debug_flag = os.getenv("ATTRLLM_INFLUENCE_DEBUG", "1").strip().lower()
+        if debug_flag in {"0", "false", "no", "off"}:
+            return
+
+        try:
+            top_n = int(os.getenv("ATTRLLM_INFLUENCE_DEBUG_TOPN", "10"))
+        except ValueError:
+            top_n = 10
+        top_n = max(1, top_n)
+
+        fourier_dict = mobius_to_fourier(mobius_dict)
+        top_mobius = sorted(
+            mobius_dict.items(),
+            key=lambda kv: abs(float(np.real(kv[1]))),
+            reverse=True,
+        )[:top_n]
+        top_fourier = sorted(
+            fourier_dict.items(),
+            key=lambda kv: abs(float(np.real(kv[1]))),
+            reverse=True,
+        )[:top_n]
+
+        singleton_rows: List[Tuple[int, str, float]] = []
+        for loc, result in values.items():
+            idxs = self._tuple_to_indices(loc)
+            if len(idxs) != 1:
+                continue
+            idx = idxs[0]
+            if 0 <= idx < len(features):
+                singleton_rows.append((idx, str(features[idx]), float(result.value)))
+        singleton_rows.sort(key=lambda row: abs(row[2]), reverse=True)
+        singleton_rows = singleton_rows[:top_n]
+
+        interaction_rows: List[Tuple[Tuple[int, ...], str, float]] = []
+        for loc, result in interactions:
+            idxs = self._tuple_to_indices(loc)
+            if len(idxs) != 2:
+                continue
+            labels = [str(features[i]) for i in idxs if 0 <= i < len(features)]
+            interaction_rows.append((tuple(idxs), " × ".join(labels), float(result.value)))
+        interaction_rows.sort(key=lambda row: abs(row[2]), reverse=True)
+        interaction_rows = interaction_rows[:top_n]
+
+        lines = [
+            "[influence-debug] "
+            f"feature_count={len(features)} mobius_terms={len(mobius_dict)} fourier_terms={len(fourier_dict)} top_n={top_n}",
+            "[influence-debug] top_mobius_terms="
+            + "; ".join([f"{loc}:{float(np.real(val)):.6f}" for loc, val in top_mobius]),
+            "[influence-debug] top_fourier_terms="
+            + "; ".join([f"{loc}:{float(np.real(val)):.6f}" for loc, val in top_fourier]),
+            "[influence-debug] top_singleton_influence="
+            + "; ".join([f"idx={idx} token={token!r} value={val:.6f}" for idx, token, val in singleton_rows]),
+            "[influence-debug] top_pair_influence="
+            + "; ".join([f"idx={idxs} pair={pair!r} value={val:.6f}" for idxs, pair, val in interaction_rows]),
+        ]
+        debug_text = "\n".join(lines)
+        print(debug_text, flush=True)
+        logger.info("%s", debug_text)
+
+    def _format_feature_payload(
+        self, singletons: Dict[Tuple[int, ...], AttributionResult], features: Sequence[str]
+    ) -> List[Dict[str, Any]]:
+        """Serialize singleton attributions for API output."""
+        payload: List[Dict[str, Any]] = []
+        n = len(features)
+        for idx in range(n):
+            loc = self._index_to_tuple(idx)
+            result = singletons.get(loc)
+            value = 0.0 if result is None else result.value
+            payload.append(
+                {
+                    "index": idx,
+                    "feature": features[idx],
+                    "value": value,
+                    "method": result.method if result else self.config.method,
+                }
+            )
+        return payload
+
+    def _format_interaction_payload(
+        self, interactions: List[Tuple[Tuple[int, ...], AttributionResult]]
+    ) -> List[Dict[str, Any]]:
+        """Serialize interaction attributions for API output."""
+        payload: List[Dict[str, Any]] = []
+        for loc, result in interactions:
+            indices = self._tuple_to_indices(loc)
+            feature_list = list(result.feature) if isinstance(result.feature, tuple) else [result.feature]
+            feature_label = " × ".join(feature_list)
+
+            payload.append(
+                {
+                    "indices": indices,
+                    "features": feature_label,
+                    "feature_list": feature_list,
+                    "value": result.value,
+                    "method": result.method,
+                    "order": result.interaction_order,
+                }
+            )
+        return payload
+    
+    def compute_attributions(
+        self,
+        context: str,
+        answer: str,
+        reference_answer: str = "",
+        prompt: str = "",
+        debug: bool = False,
+        debug_coalitions: Any | None = None,
+    ) -> Dict[str, Any]:
+        """Compute singleton values, interactions, and heatmap-ready output."""
+        mobius, features, masking_info, set_function = self._compute_mobius(
+            context,
+            answer,
+            reference_answer,
+            prompt,
+        )
+        aggregated = self._mobius_to_values_and_interactions(mobius, features)
+        method_payload = aggregated[self.config.method]
+        if self.config.method == "influence":
+            self._log_influence_debug(
+                mobius,
+                features,
+                method_payload["values"],
+                method_payload["interactions"],
+            )
+        feature_items = self._format_feature_payload(method_payload["values"], features)
+        interaction_items = self._format_interaction_payload(method_payload["interactions"])
+        interactions_only = method_payload["interactions"]
+        heatmap_attrs = (
+            [item[1] for item in interactions_only]
+            if interactions_only
+            else list(method_payload["values"].values())
+        )
+        heatmap = convert_to_heatmap_data(heatmap_attrs, feature_order=features)
+
+        debug_scores: List[Dict[str, Any]] = []
+        score_full = None
+        score_empty = None
+        score_full_details: Dict[str, Any] | None = None
+        score_empty_details: Dict[str, Any] | None = None
+        try:
+            if hasattr(set_function, "score_details"):
+                all_sel = [1] * len(features)
+                empty_sel = [0] * len(features)
+                score_full_details = set_function.score_details(all_sel)
+                score_empty_details = set_function.score_details(empty_sel)
+                score_full = score_full_details.get("score") if score_full_details else None
+                score_empty = score_empty_details.get("score") if score_empty_details else None
+        except Exception as exc:
+            logger.warning("score_full/empty computation failed: %s", exc)
+        if debug and hasattr(set_function, "score_details"):
+            n_features = len(features)
+            selections: List[Tuple[str, Tuple[int, ...], List[int]]] = []
+            selections.append(("ALL", tuple(1 for _ in range(n_features)), list(range(n_features))))
+            selections.append(("EMPTY", tuple(0 for _ in range(n_features)), []))
+
+            coalitions: List[List[int]] = []
+            if isinstance(debug_coalitions, list):
+                if debug_coalitions and all(isinstance(x, int) for x in debug_coalitions):
+                    coalitions = [[int(x)] for x in debug_coalitions]
+                else:
+                    for entry in debug_coalitions:
+                        if isinstance(entry, (list, tuple)):
+                            coalitions.append([int(x) for x in entry])
+            if coalitions:
+                for indices in coalitions:
+                    indices = [idx for idx in indices if 0 <= idx < n_features]
+                    mask = [0] * n_features
+                    for idx in indices:
+                        mask[idx] = 1
+                    selections.append((f"COALITION {indices}", tuple(mask), indices))
+            else:
+                for idx in range(min(5, n_features)):
+                    mask = [0] * n_features
+                    mask[idx] = 1
+                    selections.append((f"SINGLETON {idx}", tuple(mask), [idx]))
+
+            for label, selection, indices in selections:
+                details = set_function.score_details(selection)
+                record = {
+                    "label": label,
+                    "indices": indices,
+                    **details,
+                }
+                debug_scores.append(record)
+                print(
+                    f"[debug_scores] {label} indices={indices} "
+                    f"score={details['score']:.4f} "
+                    f"total_logprob={details['total_logprob']:.4f} "
+                    f"avg_nll={details['avg_nll']:.4f} ppl={details['ppl']:.4f} "
+                    f"target_len={details['target_len']}"
+                )
+
+        y_full = getattr(set_function, "y_full", "")
+        y_len_tokens = getattr(set_function, "y_len_tokens", 0)
+        scoring_target_source = getattr(set_function, "scoring_target_source", "model_output")
+        scoring_target_text = getattr(set_function, "scoring_target_text", y_full)
+        reference_answer_received = getattr(set_function, "reference_answer_received", "")
+        answer_received = getattr(set_function, "answer_received", "")
+
+        def _summarize_text(value: Any, preview_chars: int = 200) -> Dict[str, Any]:
+            text = str(value or "")
+            preview = text[:preview_chars]
+            if len(text) > preview_chars:
+                preview = preview + "..."
+            return {
+                "length": len(text),
+                "preview": preview,
+            }
+
+        min_logprob_seen = None
+        if hasattr(set_function, "min_logprob_seen"):
+            try:
+                min_logprob_seen = set_function.min_logprob_seen.get("value")  # type: ignore[attr-defined]
+            except Exception:
+                min_logprob_seen = None
+
+        return {
+            "method": self.config.method,
+            "features": feature_items,
+            "interactions": interaction_items,
+            "masking": asdict(masking_info),
+            "heatmap": heatmap,
+            "debug_scores": debug_scores or None,
+            "scalarizer_used": self.config.scalarizer,
+            "embedding_model": self.config.embedding_model,
+            "y_full": y_full,
+            "y_len_tokens": y_len_tokens,
+            "scoring_target_source": scoring_target_source,
+            "scoring_target_text": scoring_target_text,
+            "reference_answer_received": _summarize_text(reference_answer_received),
+            "answer_received": _summarize_text(answer_received),
+            "score_full": score_full,
+            "score_empty": score_empty,
+            "logprob_full": score_full_details.get("logprob") if score_full_details else None,
+            "logprob_empty": score_empty_details.get("logprob") if score_empty_details else None,
+            "avg_nll_full": score_full_details.get("avg_nll") if score_full_details else None,
+            "avg_nll_empty": score_empty_details.get("avg_nll") if score_empty_details else None,
+            "ppl_full": score_full_details.get("ppl") if score_full_details else None,
+            "ppl_empty": score_empty_details.get("ppl") if score_empty_details else None,
+            "min_logprob_seen": min_logprob_seen,
+        }
+
+    def compute_interactions(
+        self,
+        context: str,
+        answer: str,
+        order: int = 2,
+        prompt: str = "",
+    ) -> List[Tuple[Tuple[int, ...], float]]:
+        """Compute interaction-only payload for a requested order."""
+        mobius, _, _, _ = self._compute_mobius(context, answer, "", prompt)
+        mobius = self._validate_mobius_keys(self._trim_mobius(mobius), "mobius")
+        if self.config.method == "shapley":
+            pairs = shapley_interactions(mobius, order=order, top_k=self.config.top_k_interactions)
+        elif self.config.method == "banzhaf":
+            pairs = banzhaf_interactions(mobius, order=order, top_k=self.config.top_k_interactions)
+        elif self.config.method == "influence":
+            pairs = influence_interactions(mobius, order=order, top_k=self.config.top_k_interactions)
+        else:  # defensive; AttributorConfig should already validate this
+            raise ValueError(f"Unsupported attribution method: {self.config.method}")
+        validated_pairs = self._validate_interaction_pairs(pairs, "interactions")
+        if self.config.method == "influence":
+            try:
+                top_n = int(os.getenv("ATTRLLM_INFLUENCE_DEBUG_TOPN", "10"))
+            except ValueError:
+                top_n = 10
+            top_n = max(1, top_n)
+            top_pairs = sorted(validated_pairs, key=lambda kv: abs(float(kv[1])), reverse=True)[:top_n]
+            debug_line = (
+                "[influence-debug][interactions] "
+                f"order={order} mobius_terms={len(mobius)} returned_pairs={len(validated_pairs)} "
+                f"top_pairs={top_pairs}"
+            )
+            print(debug_line, flush=True)
+            logger.info("%s", debug_line)
+        return validated_pairs
+
+    def explain_sample(
+        self,
+        dataset: str,
+        ex_id: str,
+        method: str | None = None,
+    ) -> Dict[str, Any]:
+        """Fetch a dataset example and return attribution results."""
+        if method and method.lower() != self.config.method:
+            self.config.method = method.lower()
+        example = get_example_by_id(dataset, ex_id)
+        context = example.get("context")
+        answer = example.get("answer")
+        if context is None or answer is None:
+            raise KeyError("Dataset example must include 'context' and 'answer' fields.")
+        attributions = self.compute_attributions(context, answer)
+        return {
+            "dataset": dataset,
+            "example_id": ex_id,
+            "question": example.get("question"),
+            "attribution": attributions,
+        }
+
+
+def compute_attributions(
+    model: BaseModel,
+    context: str,
+    answer: str,
+    method: str,
+    mask_level: str = "sentence",
+    interaction_order: int = 3,
+    prompt: str | None = None,
+    debug: bool = False,
+    debug_coalitions: Any | None = None,
+    scalarizer: str = "logprob",
+    embedding_model: str | None = None,
+    reference_answer: str | None = None,
+) -> Dict[str, Any]:
+    """Convenience wrapper for full attribution payloads."""
+    config = AttributorConfig(
+        mask_level=mask_level,
+        method=method,
+        interaction_order=interaction_order,
+        scalarizer=scalarizer,
+        embedding_model=embedding_model,
+    )
+    _ = prompt  # prompt is reserved for future template control
+    return Attributor(model, config).compute_attributions(
+        context,
+        answer,
+        reference_answer=reference_answer or "",
+        prompt=prompt or "",
+        debug=debug,
+        debug_coalitions=debug_coalitions,
+    )
+
+
+def compute_interactions(
+    model: BaseModel,
+    context: str,
+    answer: str,
+    order: int = 2,
+    method: str = "shapley",
+    mask_level: str = "sentence",
+    prompt: str | None = None,
+    scalarizer: str = "logprob",
+    embedding_model: str | None = None,
+) -> List[Tuple[Tuple[int, ...], float]]:
+    """Convenience wrapper for interaction-only queries."""
+    config = AttributorConfig(
+        mask_level=mask_level,
+        method=method,
+        interaction_order=order,
+        scalarizer=scalarizer,
+        embedding_model=embedding_model,
+    )
+    return Attributor(model, config).compute_interactions(context, answer, order=order, prompt=prompt or "")
+
+
+def explain_sample(
+    dataset: str,
+    ex_id: str,
+    model: BaseModel,
+    method: str = "shapley",
+    mask_level: str = "sentence",
+) -> Dict[str, Any]:
+    """Fetch a dataset example and run attribution."""
+    config = AttributorConfig(mask_level=mask_level, method=method)
+    return Attributor(model, config).explain_sample(dataset, ex_id)

attribution/image_masker.py

@@ -0,0 +1,182 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import List, Sequence, Tuple
+
+import numpy as np
+from PIL import Image, ImageFilter
+
+
+@dataclass(frozen=True)
+class ImageRegion:
+    index: int
+    bbox: Tuple[int, int, int, int]
+    label: str
+    mask: np.ndarray | None = None
+
+
+def _ensure_rgb(image: Image.Image) -> Image.Image:
+    if image.mode == "RGB":
+        return image
+    return image.convert("RGB")
+
+
+def _region_mask(region: ImageRegion, shape: Tuple[int, int]) -> np.ndarray:
+    if region.mask is not None:
+        return region.mask
+    height, width = shape
+    x0, y0, x1, y1 = region.bbox
+    mask = np.zeros((height, width), dtype=bool)
+    x0 = max(0, min(width, x0))
+    x1 = max(0, min(width, x1))
+    y0 = max(0, min(height, y0))
+    y1 = max(0, min(height, y1))
+    if x1 > x0 and y1 > y0:
+        mask[y0:y1, x0:x1] = True
+    return mask
+
+
+def _mask_to_bbox(mask: np.ndarray) -> Tuple[int, int, int, int]:
+    ys, xs = np.where(mask)
+    if ys.size == 0 or xs.size == 0:
+        return (0, 0, 0, 0)
+    y0, y1 = int(ys.min()), int(ys.max()) + 1
+    x0, x1 = int(xs.min()), int(xs.max()) + 1
+    return (x0, y0, x1, y1)
+
+
+def _patch_grid_edges(size: int, bins: int) -> List[int]:
+    bins = max(1, int(bins))
+    edges = np.linspace(0, size, bins + 1)
+    return [int(round(val)) for val in edges]
+
+
+def supports_superpixel() -> bool:
+    try:
+        from skimage.segmentation import slic  # type: ignore  # noqa: F401
+    except Exception:
+        return False
+    return True
+
+
+def build_regions(
+    image: Image.Image,
+    mask_level: str,
+    grid_size: int | None = None,
+) -> List[ImageRegion]:
+    """Build image regions as either a patch grid or superpixel segments."""
+    image = _ensure_rgb(image)
+    width, height = image.size
+    mask_level = (mask_level or "patch").lower()
+
+    if mask_level == "patch":
+        grid = int(grid_size or 8)
+        x_edges = _patch_grid_edges(width, grid)
+        y_edges = _patch_grid_edges(height, grid)
+        regions: List[ImageRegion] = []
+        idx = 0
+        for r in range(len(y_edges) - 1):
+            for c in range(len(x_edges) - 1):
+                x0, x1 = x_edges[c], x_edges[c + 1]
+                y0, y1 = y_edges[r], y_edges[r + 1]
+                label = f"patch_{r}_{c}"
+                regions.append(ImageRegion(idx, (x0, y0, x1, y1), label))
+                idx += 1
+        return regions
+
+    if mask_level == "superpixel":
+        if not supports_superpixel():
+            raise ValueError(
+                "Superpixel masking requires scikit-image. Install it or use mask_level='patch'."
+            )
+        from skimage.segmentation import slic  # type: ignore
+
+        np_img = np.asarray(image)
+        n_segments = max(25, int((width * height) / 4096))
+        segments = slic(np_img, n_segments=n_segments, compactness=10.0, start_label=0)
+        regions = []
+        idx = 0
+        for seg_id in np.unique(segments):
+            mask = segments == seg_id
+            if not mask.any():
+                continue
+            bbox = _mask_to_bbox(mask)
+            label = f"superpixel_{int(seg_id)}"
+            regions.append(ImageRegion(idx, bbox, label, mask=mask))
+            idx += 1
+        return regions
+
+    raise ValueError(f"Unsupported mask_level: {mask_level}")
+
+
+def apply_image_mask(
+    image: Image.Image,
+    regions: Sequence[ImageRegion],
+    coalition_mask: Sequence[int],
+    mode: str = "blur",
+) -> Image.Image:
+    """
+    Apply a coalition mask to an image by hiding inactive regions.
+    """
+    image = _ensure_rgb(image)
+    arr = np.asarray(image).copy()
+    height, width = arr.shape[:2]
+    mode = (mode or "blur").lower()
+
+    blur_arr = None
+    mean_color = None
+    if mode == "blur":
+        blur_arr = np.asarray(image.filter(ImageFilter.GaussianBlur(radius=12)))
+    elif mode == "mean":
+        mean_color = arr.mean(axis=(0, 1)).astype(np.uint8)
+    elif mode == "gray":
+        mean_color = np.array([128, 128, 128], dtype=np.uint8)
+    else:
+        raise ValueError(f"Unsupported masking mode: {mode}")
+
+    if len(coalition_mask) != len(regions):
+        raise ValueError("Coalition mask length must match number of regions.")
+
+    for idx, region in enumerate(regions):
+        if int(coalition_mask[idx]):
+            continue
+        mask = _region_mask(region, (height, width))
+        if not mask.any():
+            continue
+        if mode == "blur":
+            arr[mask] = blur_arr[mask]  # type: ignore[index]
+        else:
+            arr[mask] = mean_color  # type: ignore[assignment]
+
+    return Image.fromarray(arr)
+
+
+def render_overlay(
+    image: Image.Image,
+    regions: Sequence[ImageRegion],
+    values: Sequence[float],
+    *,
+    alpha: float = 0.6,
+) -> Image.Image:
+    """Render a transparent overlay image from region attribution values."""
+    image = _ensure_rgb(image)
+    width, height = image.size
+    overlay = np.zeros((height, width, 4), dtype=np.uint8)
+    max_abs = max((abs(v) for v in values), default=0.0) or 1.0
+
+    for region, value in zip(regions, values):
+        mask = _region_mask(region, (height, width))
+        if not mask.any():
+            continue
+        norm = min(1.0, abs(value) / max_abs)
+        if norm <= 0:
+            continue
+        if value >= 0:
+            rgb = np.array([1, 109, 1], dtype=np.uint8)
+        else:
+            rgb = np.array([221, 19, 19], dtype=np.uint8)
+        opacity = int(round(255 * alpha * norm))
+        overlay[mask, :3] = rgb
+        overlay[mask, 3] = opacity
+
+    return Image.fromarray(overlay, mode="RGBA")

attribution/masker.py

@@ -0,0 +1,149 @@
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+from typing import Callable, List, Sequence, Tuple
+
+try:
+    import spacy
+except ImportError:  # pragma: no cover - optional dependency
+    spacy = None
+
+try:
+    import nltk
+    from nltk.tokenize import sent_tokenize
+except ImportError:  # pragma: no cover - optional dependency
+    nltk = None
+    sent_tokenize = None
+
+
+MaskerCallable = Callable[[str], List[Tuple[str, Tuple[int, int]]]]
+
+_SPACY_NLP = None
+
+
+@dataclass
+class MaskingInfo:
+    """Metadata describing how a piece of text was masked."""
+
+    mask_level: str
+    mask_indices: List[int]
+    feature_spans: List[Tuple[int, int]]
+
+
+def _load_spacy_model():
+    global _SPACY_NLP
+    if spacy is None:
+        return None
+    if _SPACY_NLP is not None:
+        return _SPACY_NLP
+    try:
+        _SPACY_NLP = spacy.load("en_core_web_sm")
+    except OSError:  # model might not be installed
+        _SPACY_NLP = spacy.blank("en")
+        if not _SPACY_NLP.has_pipe("sentencizer"):
+            _SPACY_NLP.add_pipe("sentencizer")
+    return _SPACY_NLP
+
+
+def _word_masker(text: str) -> List[Tuple[str, Tuple[int, int]]]:
+    pattern = re.compile(r"\S+")
+    spans = []
+    for match in pattern.finditer(text):
+        start, end = match.span()
+        spans.append((text[start:end], (start, end)))
+    return spans
+
+
+def _sentence_masker_spacy(text: str) -> List[Tuple[str, Tuple[int, int]]]:
+    nlp = _load_spacy_model()
+    if nlp is None:
+        return []
+    doc = nlp(text)
+    spans = []
+    for sent in doc.sents:
+        spans.append((sent.text.strip(), (sent.start_char, sent.end_char)))
+    return spans
+
+
+def _sentence_masker(text: str) -> List[Tuple[str, Tuple[int, int]]]:
+    spans = _sentence_masker_spacy(text)
+    if spans:
+        return spans
+    if sent_tokenize is not None:
+        sentences = sent_tokenize(text)
+    else:
+        regex = re.compile(r"[^.!?]+[.!?]?\s*", re.MULTILINE)
+        sentences = [m.group().strip() for m in regex.finditer(text)]
+    spans = []
+    cursor = 0
+    for sentence in sentences:
+        if not sentence:
+            continue
+        idx = text.find(sentence, cursor)
+        if idx == -1:
+            idx = cursor
+        cursor = idx + len(sentence)
+        spans.append((sentence.strip(), (idx, idx + len(sentence))))
+    return spans
+
+
+def _paragraph_masker(text: str) -> List[Tuple[str, Tuple[int, int]]]:
+    paragraphs: Sequence[str] = re.split(r"\n\s*\n", text.strip())
+    spans: List[Tuple[str, Tuple[int, int]]] = []
+    cursor = 0
+    for para in paragraphs:
+        normalized = para.strip()
+        if not normalized:
+            continue
+        idx = text.find(para, cursor)
+        if idx == -1:
+            idx = text.find(normalized, cursor)
+            if idx == -1:
+                continue
+            para = normalized
+        cursor = idx + len(para)
+        spans.append((normalized, (idx, idx + len(para))))
+    if not spans and text.strip():
+        spans.append((text.strip(), (0, len(text))))
+    return spans
+
+
+def get_masker(level: str) -> MaskerCallable:
+    """
+    Return a masker callable for the requested granularity level.
+
+    Parameters
+    ----------
+    level:
+        One of ``"word"``, ``"sentence"``, or ``"paragraph"``.
+    """
+    normalized = (level or "").strip().lower()
+    if normalized == "word":
+        return _word_masker
+    if normalized == "sentence":
+        return _sentence_masker
+    if normalized == "paragraph":
+        return _paragraph_masker
+    raise ValueError(f"Unsupported mask level: {level}")
+
+
+def mask_text(text: str, masker: MaskerCallable) -> Tuple[List[str], MaskingInfo]:
+    """
+    Chunk text into features according to the provided masker.
+
+    Returns
+    -------
+    features:
+        The textual surface form of each chunk.
+    masking_info:
+        Metadata describing the spans and mask indices.
+    """
+    spans_with_text = masker(text)
+    features = [entry[0] for entry in spans_with_text]
+    feature_spans = [entry[1] for entry in spans_with_text]
+    mask_indices = list(range(len(features)))
+    level = getattr(masker, "__name__", "custom")
+    info = MaskingInfo(mask_level=level, mask_indices=mask_indices, feature_spans=feature_spans)
+    return features, info
+

attribution/mm_game.py

@@ -0,0 +1,922 @@
+from __future__ import annotations
+
+import base64
+import io
+import math
+from dataclasses import asdict, dataclass
+from typing import Any, Dict, List, Sequence, Tuple
+
+from PIL import Image
+
+from attribution.image_masker import ImageRegion, apply_image_mask, build_regions, render_overlay
+from attribution.masker import MaskingInfo, get_masker, mask_text
+from attribution.proxyspex import run_proxyspex
+from attribution.utils import banzhaf_interactions, mobius_to_banzhaf, mobius_to_shapley, shapley_interactions
+
+try:
+    from loader.models import BaseModel
+except ImportError:  # pragma: no cover
+    BaseModel = Any  # type: ignore
+
+
+SUPPORTED_METHODS = {"shapley", "banzhaf"}
+
+
+def _join_prompt_answer(prompt: str, answer: str) -> str:
+    if not prompt:
+        return answer
+    if not answer:
+        return prompt
+    if prompt[-1].isspace() or answer[0].isspace():
+        return prompt + answer
+    return prompt + " " + answer
+
+
+def _inject_image_placeholder(prompt: str, placeholder: str) -> str:
+    prompt = prompt or ""
+    placeholder = placeholder or "<image>"
+    known_tokens = (placeholder, "<image>", "<|image_pad|>", "<img>")
+    for token in known_tokens:
+        if token and token in prompt:
+            return prompt
+    return f"{placeholder}\n{prompt}".strip()
+
+
+def _resolve_image_placeholder(tokenizer: Any, fallback: str) -> str:
+    for attr in ("image_placeholder", "image_token", "image_pad_token"):
+        placeholder = getattr(tokenizer, attr, None)
+        if isinstance(placeholder, str) and placeholder:
+            return placeholder
+    additional_tokens = list(getattr(tokenizer, "additional_special_tokens", None) or [])
+    for token in additional_tokens:
+        if "image_pad" in token:
+            return token
+    for token in additional_tokens:
+        if "image" in token:
+            return token
+    return fallback or "<image>"
+
+
+def _decode_image_b64(image_b64: str) -> Image.Image:
+    if not image_b64:
+        raise ValueError("image_b64 must be provided")
+    if "," in image_b64:
+        image_b64 = image_b64.split(",", 1)[1]
+    raw = base64.b64decode(image_b64)
+    return Image.open(io.BytesIO(raw)).convert("RGB")
+
+
+def _encode_image_b64(image: Image.Image) -> str:
+    buffer = io.BytesIO()
+    image.save(buffer, format="PNG")
+    return base64.b64encode(buffer.getvalue()).decode("utf-8")
+
+
+def _score_from_prompt_logprobs(
+    tokenizer: Any,
+    base_prompt: str,
+    full_prompt: str,
+    prompt_logprobs: Any,
+) -> Dict[str, Any]:
+    if prompt_logprobs is None:
+        raise RuntimeError("vLLM did not return prompt_logprobs for multimodal scoring.")
+
+    base_ids = tokenizer(base_prompt, add_special_tokens=False)["input_ids"]
+    full_ids = tokenizer(full_prompt, add_special_tokens=False)["input_ids"]
+    if not full_ids or len(full_ids) <= len(base_ids):
+        raise ValueError("Prompt did not include target tokens.")
+
+    if len(prompt_logprobs) == len(full_ids) + 1:
+        prompt_logprobs = prompt_logprobs[1:]
+    elif len(prompt_logprobs) != len(full_ids):
+        length = min(len(prompt_logprobs), len(full_ids))
+        prompt_logprobs = prompt_logprobs[:length]
+        full_ids = full_ids[:length]
+        if len(full_ids) <= len(base_ids):
+            raise ValueError("Prompt did not include target tokens.")
+
+    ref_ids = full_ids[len(base_ids) :]
+    ref_len = len(ref_ids)
+
+    def _to_float(value: Any) -> float:
+        return float(getattr(value, "logprob", value))
+
+    seq_lp = 0.0
+    token_logprobs: List[float] = []
+    for offset, token_id in enumerate(ref_ids):
+        pos = len(base_ids) + offset
+        if pos >= len(prompt_logprobs):
+            token_lp = -20.0
+        else:
+            cand = prompt_logprobs[pos] or {}
+            if token_id in cand:
+                token_lp = _to_float(cand[token_id])
+            elif cand:
+                token_lp = min(_to_float(v) for v in cand.values()) - 5.0
+            else:
+                token_lp = -20.0
+        seq_lp += token_lp
+        token_logprobs.append(token_lp)
+
+    avg_lp = seq_lp / ref_len if ref_len else float("-inf")
+    ppl = math.exp(-avg_lp) if ref_len else float("inf")
+    return {
+        "avg_logprob": avg_lp,
+        "perplexity": ppl,
+        "target_len": ref_len,
+        "sequence_logprobs": [seq_lp],
+        "token_logprobs": [token_logprobs],
+    }
+
+
+def _build_prompt_candidates(
+    tokenizer: Any,
+    prompt_text: str,
+    answer: str,
+    placeholder: str,
+) -> List[Tuple[str, str]]:
+    candidate_prompts: List[Tuple[str, str]] = []
+    tokenizer_placeholder = _resolve_image_placeholder(tokenizer, placeholder)
+    prompt_with_image = _inject_image_placeholder(prompt_text, tokenizer_placeholder)
+
+    if hasattr(tokenizer, "apply_chat_template"):
+        try:
+            content: List[Dict[str, Any]] = [{"type": "image"}]
+            if prompt_text:
+                content.append({"type": "text", "text": prompt_text})
+            messages = [{"role": "user", "content": content}]
+            base_prompt = tokenizer.apply_chat_template(
+                messages,
+                tokenize=False,
+                add_generation_prompt=True,
+            )
+            full_prompt = tokenizer.apply_chat_template(
+                messages + [{"role": "assistant", "content": answer}],
+                tokenize=False,
+                add_generation_prompt=False,
+            )
+            candidate_prompts.append((base_prompt, full_prompt))
+        except Exception:
+            pass
+        try:
+            messages = [{"role": "user", "content": prompt_with_image}]
+            base_prompt = tokenizer.apply_chat_template(
+                messages,
+                tokenize=False,
+                add_generation_prompt=True,
+            )
+            full_prompt = tokenizer.apply_chat_template(
+                messages + [{"role": "assistant", "content": answer}],
+                tokenize=False,
+                add_generation_prompt=False,
+            )
+            candidate_prompts.append((base_prompt, full_prompt))
+        except Exception:
+            pass
+
+    candidate_prompts.append((prompt_with_image, _join_prompt_answer(prompt_with_image, answer)))
+    return candidate_prompts
+
+
+def _build_prompt_candidates_batch(
+    tokenizer: Any,
+    prompts: Sequence[str],
+    answer: str,
+    placeholder: str,
+) -> List[Tuple[List[str], List[str]]]:
+    prompt_texts = [prompt or "" for prompt in prompts]
+    tokenizer_placeholder = _resolve_image_placeholder(tokenizer, placeholder)
+    prompt_with_image_list = [
+        _inject_image_placeholder(prompt_text, tokenizer_placeholder)
+        for prompt_text in prompt_texts
+    ]
+    candidates: List[Tuple[List[str], List[str]]] = []
+
+    if hasattr(tokenizer, "apply_chat_template"):
+        try:
+            base_prompts: List[str] = []
+            full_prompts: List[str] = []
+            for prompt_text in prompt_texts:
+                content: List[Dict[str, Any]] = [{"type": "image"}]
+                if prompt_text:
+                    content.append({"type": "text", "text": prompt_text})
+                messages = [{"role": "user", "content": content}]
+                base_prompts.append(
+                    tokenizer.apply_chat_template(
+                        messages,
+                        tokenize=False,
+                        add_generation_prompt=True,
+                    )
+                )
+                full_prompts.append(
+                    tokenizer.apply_chat_template(
+                        messages + [{"role": "assistant", "content": answer}],
+                        tokenize=False,
+                        add_generation_prompt=False,
+                    )
+                )
+            candidates.append((base_prompts, full_prompts))
+        except Exception:
+            pass
+        try:
+            base_prompts = []
+            full_prompts = []
+            for prompt_with_image in prompt_with_image_list:
+                messages = [{"role": "user", "content": prompt_with_image}]
+                base_prompts.append(
+                    tokenizer.apply_chat_template(
+                        messages,
+                        tokenize=False,
+                        add_generation_prompt=True,
+                    )
+                )
+                full_prompts.append(
+                    tokenizer.apply_chat_template(
+                        messages + [{"role": "assistant", "content": answer}],
+                        tokenize=False,
+                        add_generation_prompt=False,
+                    )
+                )
+            candidates.append((base_prompts, full_prompts))
+        except Exception:
+            pass
+
+    base_prompts = prompt_with_image_list
+    full_prompts = [_join_prompt_answer(prompt_with_image, answer) for prompt_with_image in base_prompts]
+    candidates.append((base_prompts, full_prompts))
+    return candidates
+
+
+def _score_multimodal_vllm(model: BaseModel, prompt: str, answer: str, image: Image.Image) -> Dict[str, Any]:
+    info = getattr(model, "info", None)
+    if callable(info):
+        info = info()
+    is_multimodal = bool(getattr(info, "is_multimodal", False))
+    model_name = getattr(info, "name", type(model).__name__)
+    if not is_multimodal:
+        raise ValueError(
+            f"Model '{model_name}' is not a multimodal model. "
+            "Select a VLM from the multimodal registry."
+        )
+    if not hasattr(model, "llm"):
+        raise RuntimeError("Multimodal scoring requires a VLLMModel with .llm available.")
+    llm = model.llm
+    tokenizer = llm.get_tokenizer()
+
+    placeholder = getattr(info, "image_placeholder", None) or "<image>"
+    prompt_text = prompt or ""
+    candidate_prompts = _build_prompt_candidates(tokenizer, prompt_text, answer, placeholder)
+
+    from vllm import SamplingParams
+
+    sp = SamplingParams(
+        max_tokens=1,
+        temperature=0.0,
+        top_p=1.0,
+        logprobs=5,
+        prompt_logprobs=1,
+    )
+
+    def _generate_with_fallback(full_prompt: str) -> Any:
+        mm_key = getattr(info, "mm_data_key", None) or "image"
+        attempts = [
+            (f"{mm_key}=[image]", {mm_key: [image]}),
+            (f"{mm_key}=image", {mm_key: image}),
+        ]
+        errors: List[str] = []
+        last_exc: Exception | None = None
+        for label, mm_data in attempts:
+            try:
+                requests = [{"prompt": full_prompt, "multi_modal_data": mm_data}]
+                return llm.generate(requests, sp)
+            except Exception as exc:
+                last_exc = exc
+                errors.append(f"{label}: {type(exc).__name__}: {exc}")
+        detail = "; ".join(errors) if errors else "Unknown error"
+        raise ValueError(
+            "vLLM failed to process the multimodal prompt. "
+            f"Attempts: {detail}"
+        ) from last_exc
+
+    last_exc: Exception | None = None
+    for base_prompt, full_prompt in candidate_prompts:
+        try:
+            results = _generate_with_fallback(full_prompt)
+            req_out = results[0]
+            prompt_logprobs = getattr(req_out, "prompt_logprobs", None)
+            return _score_from_prompt_logprobs(tokenizer, base_prompt, full_prompt, prompt_logprobs)
+        except Exception as exc:
+            last_exc = exc
+            continue
+
+    detail = f"vLLM failed to score the multimodal prompt. Last error: {last_exc}"
+    raise ValueError(detail)
+
+
+def _score_multimodal_vllm_batch(
+    model: BaseModel,
+    prompts: Sequence[str],
+    answer: str,
+    images: Sequence[Image.Image],
+) -> List[float]:
+    info = getattr(model, "info", None)
+    if callable(info):
+        info = info()
+    is_multimodal = bool(getattr(info, "is_multimodal", False))
+    model_name = getattr(info, "name", type(model).__name__)
+    if not is_multimodal:
+        raise ValueError(
+            f"Model '{model_name}' is not a multimodal model. "
+            "Select a VLM from the multimodal registry."
+        )
+    if not hasattr(model, "llm"):
+        raise RuntimeError("Multimodal scoring requires a VLLMModel with .llm available.")
+
+    if len(prompts) != len(images):
+        raise ValueError("prompts and images must have the same length.")
+    if not prompts:
+        return []
+
+    llm = model.llm
+    tokenizer = llm.get_tokenizer()
+    placeholder = getattr(info, "image_placeholder", None) or "<image>"
+
+    candidate_batches = _build_prompt_candidates_batch(tokenizer, prompts, answer, placeholder)
+
+    from vllm import SamplingParams
+
+    sp = SamplingParams(
+        max_tokens=1,
+        temperature=0.0,
+        top_p=1.0,
+        logprobs=5,
+        prompt_logprobs=1,
+    )
+
+    def _generate_with_fallback_batch(full_prompts: Sequence[str]) -> Any:
+        mm_key = getattr(info, "mm_data_key", None) or "image"
+        attempts = [
+            (f"{mm_key}=[image]", lambda img: {mm_key: [img]}),
+            (f"{mm_key}=image", lambda img: {mm_key: img}),
+        ]
+        errors: List[str] = []
+        last_exc: Exception | None = None
+        for label, builder in attempts:
+            try:
+                requests = [
+                    {"prompt": prompt, "multi_modal_data": builder(image)}
+                    for prompt, image in zip(full_prompts, images)
+                ]
+                return llm.generate(requests, sp)
+            except Exception as exc:
+                last_exc = exc
+                errors.append(f"{label}: {type(exc).__name__}: {exc}")
+        detail = "; ".join(errors) if errors else "Unknown error"
+        raise ValueError(
+            "vLLM failed to process the multimodal prompt. "
+            f"Attempts: {detail}"
+        ) from last_exc
+
+    last_exc: Exception | None = None
+    for base_prompts, full_prompts in candidate_batches:
+        try:
+            results = _generate_with_fallback_batch(full_prompts)
+            scores: List[float] = []
+            for base_prompt, full_prompt, req_out in zip(base_prompts, full_prompts, results):
+                prompt_logprobs = getattr(req_out, "prompt_logprobs", None)
+                stats = _score_from_prompt_logprobs(
+                    tokenizer,
+                    base_prompt,
+                    full_prompt,
+                    prompt_logprobs,
+                )
+                avg_logprob = stats.get("avg_logprob")
+                if avg_logprob is None or not math.isfinite(avg_logprob):
+                    raise ValueError("Multimodal scoring returned a non-finite logprob.")
+                scores.append(float(avg_logprob))
+            return scores
+        except Exception as exc:
+            last_exc = exc
+            continue
+
+    detail = f"vLLM failed to score the multimodal prompt. Last error: {last_exc}"
+    raise ValueError(detail)
+
+
+@dataclass
+class ImageAttributorConfig:
+    mask_level: str = "patch"
+    grid_size: int | None = None
+    method: str = "shapley"
+    interaction_order: int = 2
+    top_k_interactions: int = 15
+    random_seed: int = 0
+    mask_mode: str = "blur"
+    batch_size: int = 8
+
+    def __post_init__(self) -> None:
+        method = (self.method or "shapley").lower()
+        if method not in SUPPORTED_METHODS:
+            raise ValueError(f"Unsupported attribution method: {self.method}")
+        self.method = method
+
+
+@dataclass
+class MMAttributorConfig:
+    mask_level_image: str = "patch"
+    mask_level_text: str = "word"
+    grid_size: int | None = None
+    method: str = "shapley"
+    interaction_order: int = 2
+    top_k_interactions: int = 15
+    random_seed: int = 0
+    mask_mode: str = "blur"
+    batch_size: int = 8
+
+    def __post_init__(self) -> None:
+        method = (self.method or "shapley").lower()
+        if method not in SUPPORTED_METHODS:
+            raise ValueError(f"Unsupported attribution method: {self.method}")
+        self.method = method
+
+
+def _format_interactions(
+    interactions: List[Tuple[Tuple[int, ...], float]],
+    labels: Sequence[str],
+    method: str,
+) -> List[Dict[str, Any]]:
+    payload: List[Dict[str, Any]] = []
+    for loc, value in interactions:
+        indices = sorted(loc)
+        feature_list = [labels[idx] for idx in indices if idx < len(labels)]
+        payload.append(
+            {
+                "indices": indices,
+                "features": " × ".join(feature_list),
+                "feature_list": feature_list,
+                "value": float(value),
+                "method": method,
+                "order": len(indices),
+            }
+        )
+    return payload
+
+
+class ImageAttributor:
+    def __init__(self, model: BaseModel, config: ImageAttributorConfig | None = None) -> None:
+        self.model = model
+        self.config = config or ImageAttributorConfig()
+
+    def _build_set_function(
+        self,
+        image: Image.Image,
+        prompt: str,
+        answer: str,
+        regions: Sequence[ImageRegion],
+    ):
+        if not answer:
+            raise ValueError("Answer must be non-empty for multimodal scoring.")
+        n_features = len(regions)
+        cache: Dict[Tuple[int, ...], float] = {}
+
+        def score_active(active_tuple: Tuple[int, ...]) -> float:
+            if active_tuple in cache:
+                return cache[active_tuple]
+            coalition_mask = [0] * n_features
+            for idx in active_tuple:
+                if 0 <= idx < n_features:
+                    coalition_mask[idx] = 1
+            masked_image = apply_image_mask(
+                image,
+                regions,
+                coalition_mask,
+                mode=self.config.mask_mode,
+            )
+            score = _score_multimodal_vllm(self.model, prompt, answer, masked_image)
+            avg_logprob = score.get("avg_logprob")
+            if avg_logprob is None or not math.isfinite(avg_logprob):
+                raise ValueError("Multimodal scoring returned a non-finite logprob.")
+            cache[active_tuple] = float(avg_logprob)
+            return cache[active_tuple]
+
+        def set_function(selection: Sequence[int]) -> float:
+            if len(selection) != n_features:
+                raise ValueError("Selection length must match number of regions.")
+            active = tuple(idx for idx, flag in enumerate(selection) if int(flag))
+            return score_active(active)
+
+        def batch_set_function(selections: Sequence[Sequence[int]]) -> List[float]:
+            num_rows = len(selections)
+            if num_rows == 0:
+                return []
+            if len(selections[0]) != n_features:
+                raise ValueError("Selection length must match number of regions.")
+            results: List[float] = [0.0] * num_rows
+            pending: List[Tuple[int, ...]] = []
+            pending_indices: List[int] = []
+            for row_idx, row in enumerate(selections):
+                active = tuple(idx for idx, flag in enumerate(row) if int(flag))
+                if active in cache:
+                    results[row_idx] = cache[active]
+                else:
+                    pending.append(active)
+                    pending_indices.append(row_idx)
+
+            batch_size = max(1, int(self.config.batch_size or 1))
+            for start in range(0, len(pending), batch_size):
+                batch_actives = pending[start : start + batch_size]
+                batch_indices = pending_indices[start : start + batch_size]
+                masked_images: List[Image.Image] = []
+                batch_prompts: List[str] = []
+                for active_tuple in batch_actives:
+                    coalition_mask = [0] * n_features
+                    for idx in active_tuple:
+                        if 0 <= idx < n_features:
+                            coalition_mask[idx] = 1
+                    masked_images.append(
+                        apply_image_mask(
+                            image,
+                            regions,
+                            coalition_mask,
+                            mode=self.config.mask_mode,
+                        )
+                    )
+                    batch_prompts.append(prompt)
+                try:
+                    batch_scores = _score_multimodal_vllm_batch(
+                        self.model,
+                        batch_prompts,
+                        answer,
+                        masked_images,
+                    )
+                except Exception:
+                    batch_scores = []
+                    for batch_prompt, masked_image in zip(batch_prompts, masked_images):
+                        score = _score_multimodal_vllm(
+                            self.model,
+                            batch_prompt,
+                            answer,
+                            masked_image,
+                        )
+                        avg_logprob = score.get("avg_logprob")
+                        if avg_logprob is None or not math.isfinite(avg_logprob):
+                            raise ValueError("Multimodal scoring returned a non-finite logprob.")
+                        batch_scores.append(float(avg_logprob))
+
+                for row_idx, active_tuple, score in zip(
+                    batch_indices,
+                    batch_actives,
+                    batch_scores,
+                ):
+                    if score is None or not math.isfinite(score):
+                        raise ValueError("Multimodal scoring returned a non-finite logprob.")
+                    cache[active_tuple] = float(score)
+                    results[row_idx] = float(score)
+
+            return results
+
+        set_function.batch = batch_set_function  # type: ignore[attr-defined]
+
+        return set_function
+
+    def compute_attributions(self, image_b64: str, prompt: str, answer: str) -> Dict[str, Any]:
+        image = _decode_image_b64(image_b64)
+        regions = build_regions(image, self.config.mask_level, self.config.grid_size)
+        labels = [region.label for region in regions]
+
+        set_function = self._build_set_function(image, prompt, answer, regions)
+        mobius = run_proxyspex(
+            set_function,
+            labels,
+            max_order=self.config.interaction_order,
+            seed=self.config.random_seed,
+        )
+
+        if self.config.method == "shapley":
+            values = mobius_to_shapley(mobius)
+            interactions = shapley_interactions(
+                mobius,
+                order=self.config.interaction_order,
+                top_k=self.config.top_k_interactions,
+            )
+        else:
+            values = mobius_to_banzhaf(mobius)
+            interactions = banzhaf_interactions(
+                mobius,
+                order=self.config.interaction_order,
+                top_k=self.config.top_k_interactions,
+            )
+
+        value_list = []
+        raw_values = []
+        for idx in range(len(regions)):
+            value = float(values.get((idx,), 0.0))
+            raw_values.append(value)
+            value_list.append({"index": idx, "value": value, "method": self.config.method})
+
+        overlay = render_overlay(image, regions, raw_values)
+        image_payload = {
+            "image_b64": _encode_image_b64(image),
+            "overlay_b64": _encode_image_b64(overlay),
+            "width": image.size[0],
+            "height": image.size[1],
+            "mask_level": self.config.mask_level,
+            "grid_size": self.config.grid_size,
+        }
+
+        region_payload = [
+            {"index": region.index, "label": region.label, "bbox": list(region.bbox)}
+            for region in regions
+        ]
+
+        return {
+            "method": self.config.method,
+            "regions": region_payload,
+            "values": value_list,
+            "interactions": _format_interactions(interactions, labels, self.config.method),
+            "image": image_payload,
+        }
+
+    def compute_attributions_with_regions(
+        self,
+        image_b64: str,
+        prompt: str,
+        answer: str,
+        regions: List[ImageRegion],
+    ) -> Dict[str, Any]:
+        """Like compute_attributions() but accepts pre-built regions (e.g. UnSAM segments).
+
+        Bypasses build_regions() — all other logic is identical.
+        """
+        image = _decode_image_b64(image_b64)
+        labels = [region.label for region in regions]
+
+        set_function = self._build_set_function(image, prompt, answer, regions)
+        mobius = run_proxyspex(
+            set_function,
+            labels,
+            max_order=self.config.interaction_order,
+            seed=self.config.random_seed,
+        )
+
+        if self.config.method == "shapley":
+            values = mobius_to_shapley(mobius)
+            interactions = shapley_interactions(
+                mobius,
+                order=self.config.interaction_order,
+                top_k=self.config.top_k_interactions,
+            )
+        else:
+            values = mobius_to_banzhaf(mobius)
+            interactions = banzhaf_interactions(
+                mobius,
+                order=self.config.interaction_order,
+                top_k=self.config.top_k_interactions,
+            )
+
+        value_list = []
+        raw_values = []
+        for idx in range(len(regions)):
+            value = float(values.get((idx,), 0.0))
+            raw_values.append(value)
+            value_list.append({"index": idx, "value": value, "method": self.config.method})
+
+        overlay = render_overlay(image, regions, raw_values)
+        image_payload = {
+            "image_b64": _encode_image_b64(image),
+            "overlay_b64": _encode_image_b64(overlay),
+            "width": image.size[0],
+            "height": image.size[1],
+            "mask_level": "unsam",
+            "grid_size": len(regions),
+        }
+
+        region_payload = [
+            {"index": region.index, "label": region.label, "bbox": list(region.bbox)}
+            for region in regions
+        ]
+
+        return {
+            "method": self.config.method,
+            "regions": region_payload,
+            "values": value_list,
+            "interactions": _format_interactions(interactions, labels, self.config.method),
+            "image": image_payload,
+        }
+
+
+class MMAttributor:
+    def __init__(self, model: BaseModel, config: MMAttributorConfig | None = None) -> None:
+        self.model = model
+        self.config = config or MMAttributorConfig()
+
+    def _mask_text(self, text: str) -> Tuple[List[str], MaskingInfo]:
+        masker_fn = get_masker(self.config.mask_level_text)
+        features, info = mask_text(text or "", masker_fn)
+        info.mask_level = self.config.mask_level_text
+        return features, info
+
+    def _build_set_function(
+        self,
+        image: Image.Image,
+        text_features: Sequence[str],
+        answer: str,
+        regions: Sequence[ImageRegion],
+    ):
+        if not answer:
+            raise ValueError("Answer must be non-empty for multimodal scoring.")
+        n_img = len(regions)
+        n_txt = len(text_features)
+        n_total = n_img + n_txt
+        cache: Dict[Tuple[int, ...], float] = {}
+
+        def score_active(active_tuple: Tuple[int, ...]) -> float:
+            if active_tuple in cache:
+                return cache[active_tuple]
+            mask_img = [0] * n_img
+            mask_txt = [0] * n_txt
+            for idx in active_tuple:
+                if idx < n_img:
+                    mask_img[idx] = 1
+                else:
+                    mask_txt[idx - n_img] = 1
+            masked_image = apply_image_mask(
+                image,
+                regions,
+                mask_img,
+                mode=self.config.mask_mode,
+            )
+            masked_text = " ".join(
+                text_features[i] for i, flag in enumerate(mask_txt) if flag
+            ).strip()
+            score = _score_multimodal_vllm(self.model, masked_text, answer, masked_image)
+            avg_logprob = score.get("avg_logprob")
+            if avg_logprob is None or not math.isfinite(avg_logprob):
+                raise ValueError("Multimodal scoring returned a non-finite logprob.")
+            cache[active_tuple] = float(avg_logprob)
+            return cache[active_tuple]
+
+        def set_function(selection: Sequence[int]) -> float:
+            if len(selection) != n_total:
+                raise ValueError("Selection length must match number of features.")
+            active = tuple(idx for idx, flag in enumerate(selection) if int(flag))
+            return score_active(active)
+
+        def batch_set_function(selections: Sequence[Sequence[int]]) -> List[float]:
+            num_rows = len(selections)
+            if num_rows == 0:
+                return []
+            if len(selections[0]) != n_total:
+                raise ValueError("Selection length must match number of features.")
+            results: List[float] = [0.0] * num_rows
+            pending: List[Tuple[int, ...]] = []
+            pending_indices: List[int] = []
+            for row_idx, row in enumerate(selections):
+                active = tuple(idx for idx, flag in enumerate(row) if int(flag))
+                if active in cache:
+                    results[row_idx] = cache[active]
+                else:
+                    pending.append(active)
+                    pending_indices.append(row_idx)
+
+            batch_size = max(1, int(self.config.batch_size or 1))
+            for start in range(0, len(pending), batch_size):
+                batch_actives = pending[start : start + batch_size]
+                batch_indices = pending_indices[start : start + batch_size]
+                masked_images: List[Image.Image] = []
+                batch_prompts: List[str] = []
+                for active_tuple in batch_actives:
+                    mask_img = [0] * n_img
+                    mask_txt = [0] * n_txt
+                    for idx in active_tuple:
+                        if idx < n_img:
+                            mask_img[idx] = 1
+                        else:
+                            mask_txt[idx - n_img] = 1
+                    masked_images.append(
+                        apply_image_mask(
+                            image,
+                            regions,
+                            mask_img,
+                            mode=self.config.mask_mode,
+                        )
+                    )
+                    masked_text = " ".join(
+                        text_features[i] for i, flag in enumerate(mask_txt) if flag
+                    ).strip()
+                    batch_prompts.append(masked_text)
+
+                try:
+                    batch_scores = _score_multimodal_vllm_batch(
+                        self.model,
+                        batch_prompts,
+                        answer,
+                        masked_images,
+                    )
+                except Exception:
+                    batch_scores = []
+                    for batch_prompt, masked_image in zip(batch_prompts, masked_images):
+                        score = _score_multimodal_vllm(
+                            self.model,
+                            batch_prompt,
+                            answer,
+                            masked_image,
+                        )
+                        avg_logprob = score.get("avg_logprob")
+                        if avg_logprob is None or not math.isfinite(avg_logprob):
+                            raise ValueError("Multimodal scoring returned a non-finite logprob.")
+                        batch_scores.append(float(avg_logprob))
+
+                for row_idx, active_tuple, score in zip(
+                    batch_indices,
+                    batch_actives,
+                    batch_scores,
+                ):
+                    if score is None or not math.isfinite(score):
+                        raise ValueError("Multimodal scoring returned a non-finite logprob.")
+                    cache[active_tuple] = float(score)
+                    results[row_idx] = float(score)
+
+            return results
+
+        set_function.batch = batch_set_function  # type: ignore[attr-defined]
+
+        return set_function
+
+    def compute_attributions(self, image_b64: str, text_context: str, answer: str) -> Dict[str, Any]:
+        image = _decode_image_b64(image_b64)
+        regions = build_regions(image, self.config.mask_level_image, self.config.grid_size)
+        text_features, text_info = self._mask_text(text_context)
+        n_img = len(regions)
+
+        labels = [f"img:{region.label}" for region in regions] + [
+            f"txt:{feature}" for feature in text_features
+        ]
+        set_function = self._build_set_function(image, text_features, answer, regions)
+        mobius = run_proxyspex(
+            set_function,
+            labels,
+            max_order=self.config.interaction_order,
+            seed=self.config.random_seed,
+        )
+
+        if self.config.method == "shapley":
+            values = mobius_to_shapley(mobius)
+            interactions = shapley_interactions(
+                mobius,
+                order=self.config.interaction_order,
+                top_k=self.config.top_k_interactions,
+            )
+        else:
+            values = mobius_to_banzhaf(mobius)
+            interactions = banzhaf_interactions(
+                mobius,
+                order=self.config.interaction_order,
+                top_k=self.config.top_k_interactions,
+            )
+
+        features_payload = []
+        image_values = []
+        for idx, label in enumerate(labels):
+            value = float(values.get((idx,), 0.0))
+            modality = "image" if idx < n_img else "text"
+            ref_index = idx if modality == "image" else idx - n_img
+            features_payload.append(
+                {
+                    "index": idx,
+                    "feature": label,
+                    "value": value,
+                    "method": self.config.method,
+                    "modality": modality,
+                    "ref_index": ref_index,
+                }
+            )
+            if modality == "image":
+                image_values.append(value)
+
+        overlay = render_overlay(image, regions, image_values)
+        image_payload = {
+            "image_b64": _encode_image_b64(image),
+            "overlay_b64": _encode_image_b64(overlay),
+            "width": image.size[0],
+            "height": image.size[1],
+            "mask_level": self.config.mask_level_image,
+            "grid_size": self.config.grid_size,
+        }
+
+        region_payload = [
+            {"index": region.index, "label": region.label, "bbox": list(region.bbox)}
+            for region in regions
+        ]
+
+        return {
+            "method": self.config.method,
+            "features": features_payload,
+            "interactions": _format_interactions(interactions, labels, self.config.method),
+            "regions": region_payload,
+            "image": image_payload,
+            "text": {
+                "features": text_features,
+                "masking": asdict(text_info),
+            },
+        }

attribution/proxyspex.py

@@ -0,0 +1,311 @@
+from __future__ import annotations
+
+import math
+from itertools import combinations
+from typing import Callable, Dict, Optional, Sequence, Tuple
+from sklearn.model_selection import GridSearchCV, KFold
+import pandas as pd
+import lightgbm as lgb
+import numpy as np
+from collections import defaultdict
+from typing import Any
+
+from .utils import fourier_to_mobius
+
+
+def check_lgbm() -> Dict[str, Any]:
+    """
+    Return LightGBM status information.
+
+    Since LightGBM is a hard dependency (imported at module level), this
+    should always succeed after a successful import of this module.
+    """
+    return {
+        "installed": True,
+        "version": lgb.__version__,
+        "interaction_mode": "proxyspex",
+    }
+
+
+def _pairwise_mobius_small_n(
+    set_function: Callable[[Tuple[int, ...]], float],
+    n_features: int,
+) -> Dict[Tuple[int, ...], float]:
+    """
+    Exact Mobius up to order 2 (empty, singletons, pairs) for small n_features.
+
+    m({i})   = f(1_i) - f(0)
+    m({i,j}) = f(1_i,1_j) - f(1_i,0_j) - f(0_i,1_j) + f(0,0)
+    """
+    mobius: Dict[Tuple[int, ...], float] = {}
+
+    # --- empty set
+    empty_selection = tuple(0 for _ in range(n_features))
+    f0 = float(set_function(empty_selection))
+    mobius[()] = f0
+
+    # --- singletons
+    f_single: Dict[int, float] = {}
+    for i in range(n_features):
+        mask = [0] * n_features
+        mask[i] = 1
+        sel = tuple(mask)
+        fi = float(set_function(sel))
+        f_single[i] = fi
+        mobius[(i,)] = fi - f0
+
+    # --- pairs
+    for i, j in combinations(range(n_features), 2):
+        mask_ij = [0] * n_features
+        mask_ij[i] = 1
+        mask_ij[j] = 1
+        sel_ij = tuple(mask_ij)
+        f_ij = float(set_function(sel_ij))
+
+        # Mobius coefficient for pair {i,j}
+        m_ij = f_ij - f_single[i] - f_single[j] + f0
+        mobius[(i, j)] = m_ij
+
+    return mobius
+
+
+def run_proxyspex(
+    set_function: Callable[[Tuple[int, ...]], float],
+    features: Sequence[str],
+    max_order: int = 3,
+    seed: Optional[int] = None,
+) -> Dict[Tuple[int, ...], float]:
+    """
+    Execute ProxySPEX once to obtain Mobius interaction coefficients.
+
+    For small n_features (<= SMALL_N_THRESHOLD), we bypass ProxySPEX and
+    compute Mobius terms exactly up to order 2. This gives reliable singleton
+    and pairwise interactions for the web UI.
+    """
+    # global ProxySPEX, _shapiq_import_error
+
+    n_features = len(features)
+    if n_features == 0:
+        return {}
+
+    # ---- SMALL-N EXACT PATH (no shapiq, no lightgbm, always safe) ----
+    # if n_features <= SMALL_N_THRESHOLD:
+    #     print(f"[run_proxyspex] Using small-n pairwise backend (n={n_features}).")
+    #     return _pairwise_mobius_small_n(set_function, n_features)
+
+    # ---- LARGE-N PATH: try ProxySPEX lazily, but never crash ----
+
+    # Lazy import of ProxySPEX the first time we need it
+    # if ProxySPEX is None and _shapiq_import_error is None:  # pragma: no cover
+    #     try:
+    #         from shapiq import ProxySPEX as _ProxySPEX  # type: ignore
+    #         ProxySPEX = _ProxySPEX
+    #         logger.info("[run_proxyspex] Loaded ProxySPEX successfully.")
+    #     except Exception as e:
+    #         _shapiq_import_error = e
+    #         ProxySPEX = None
+    #         logger.warning(
+    #             "[run_proxyspex] ProxySPEX unavailable (%s); "
+    #             "falling back to singleton-only Mobius.",
+    #             e,
+    #         )
+
+    # If ProxySPEX is still not available, fall back instead of raising
+    # if ProxySPEX is None:  # pragma: no cover
+    #     print("[run_proxyspex] ProxySPEX unavailable; ")
+    #     print(
+    #         "[run_proxyspex] Using small-n pairwise backend (n=32). "
+    #         f"(n={n_features})."
+    #     )
+    #     return _singleton_mobius(set_function, n_features)
+
+    # We have ProxySPEX; run it for real
+    order = max(1, min(max_order, n_features))
+
+    def game(coalitions):
+        arr = np.asarray(coalitions)
+        if arr.ndim == 1:
+            outputs = set_function(tuple(int(x) for x in arr))
+            return outputs - np.mean(outputs)
+        if arr.ndim == 2:
+            if hasattr(set_function, "batch"):
+                outputs = np.asarray(set_function.batch(arr), dtype=float)
+            else:
+                outputs = np.array(
+                    [set_function(tuple(int(x) for x in row)) for row in arr],
+                    dtype=float,
+                )
+            return outputs - np.mean(outputs)
+        raise ValueError("Coalitions input must be 1D or 2D.")
+
+    # try:
+    #     approximator = ProxySPEX(
+    #         n=n_features,
+    #         index="FSII",
+    #         max_order=n_features,
+    #         random_state=seed,
+    #     )
+    # except Exception as e:
+    #     print(e)
+    #     print(
+    #         "[run_proxyspex] Could not construct ProxySPEX (%s); "
+    #         "falling back to singleton-only Mobius." % e
+    #     )
+    #     return _singleton_mobius(set_function, n_features)
+
+    budget = int(10 * n_features * math.log(n_features+1, 2))
+
+    # try:
+    #     print(budget)
+    #     raw_result = approximator.approximate(budget=budget, game=game)
+    # except Exception as e:
+    #     print(
+    #         f"[run_proxyspex] ProxySPEX failed with {type(e).__name__}: {e}. "
+    #         "Falling back to singleton-only Mobius."
+    #     )
+    #     return _singleton_mobius(set_function, n_features)
+    dict_values = proxyspex(n_features, budget, game, max_order=order, random_seed=seed)
+
+    mobius_dict: Dict[Tuple[int, ...], float] = {}
+    for loc, value in dict_values.items():
+        normalized = tuple(sorted(loc))
+        mobius_dict[normalized] = float(np.real(value))
+
+    return mobius_dict
+
+
+def proxyspex(n, budget, game, max_order=None, random_seed=0):
+    # Sample budget uniform coalitions (boolean lists) from game
+    # create a np array of shape (buget, n) of 0s and 1s
+    print("the random seed is: ", random_seed)
+    print(f"[proxyspex] n_features={n} budget={budget} max_order={max_order}")
+    rng = np.random.default_rng(random_seed)
+    X_samples = rng.integers(0, 2, size=(budget, n), dtype=np.int8)
+
+    train_X = pd.DataFrame(
+        X_samples,
+        columns=np.array([f"f{i}" for i in range(n)]),
+    ).astype(np.int8)
+
+    y_samples = game(X_samples)
+
+
+    if len(train_X) < 5:
+        # fit just one XGBoost with depth 1 
+        base_model = lgb.LGBMRegressor(verbose=-1, n_jobs=1, 
+                                       random_state=random_seed,
+                                       max_depth=1,
+                                       )
+        base_model.fit(train_X, y_samples)
+        best_model = base_model
+    else:
+        base_model = lgb.LGBMRegressor(verbose=-1, n_jobs=1, 
+                                random_state=random_seed,
+                                )
+        param_grid = {
+            "max_depth": [3, 5],
+            "n_estimators": [500, 1000],
+            "learning_rate": [0.01, 0.1],
+        }
+        cv_strategy = KFold(n_splits=5, shuffle=True, random_state=random_seed)
+
+        grid_search = GridSearchCV(
+            estimator=base_model,
+            param_grid=param_grid,
+            scoring="r2",
+            cv=cv_strategy,
+            verbose=0,
+            n_jobs=1,
+        )
+
+        grid_search.fit(train_X, y_samples)
+
+        best_model = grid_search.best_estimator_
+    
+    initial_transform = lgboost_to_fourier(best_model.booster_.dump_model())
+
+    mobius_transform = fourier_to_mobius(initial_transform)
+
+    max_order = max_order if max_order is not None else n
+
+    filtered_transform = {k: v for k, v in mobius_transform.items() if len(k) <= max_order}
+    return filtered_transform
+
+
+def lgboost_to_fourier(model_dict: dict[str, Any]) -> dict[tuple[int, ...], float]:
+    """Extracts the aggregated Fourier coefficients from an LGBoost model dictionary.
+
+    This method iterates over all trees in the LightGBM ensemble, computes the
+    Fourier coefficients for each individual tree using the `_lgboost_tree_to_fourier`
+    helper method, and then sums these coefficients to get the final Fourier
+    representation of the complete model.
+
+    Args:
+    model_dict: A dictionary representing the trained LGBoost model, as
+        produced by `model.booster_.dump_model()`.
+
+    Returns:
+        A dictionary that maps interaction tuples (representing Fourier frequencies)
+        to their aggregated Fourier coefficients.
+    """
+    aggregated_coeffs = defaultdict(float)
+
+    for tree_info in model_dict["tree_info"]:
+        tree_coeffs = lgboost_tree_to_fourier(tree_info)
+        for interaction, value in tree_coeffs.items():
+            aggregated_coeffs[interaction] += value
+
+    # Convert defaultdict to a standard dict, removing zero-valued coefficients
+    return {k: v for k, v in aggregated_coeffs.items() if v != 0.0}
+
+def lgboost_tree_to_fourier(tree_info: dict[str, Any]) -> dict[tuple[int, ...], float]:
+    """Recursively strips the Fourier coefficients from a single LGBoost tree.
+
+    This method traverses a tree's structure, as provided by LightGBM's `dump_model`
+    method, and computes the Fourier representation of the piecewise-constant
+    function that the tree defines. The logic is adapted from the work by Gorji et al. (2024).
+
+    Args:
+        tree_info: A dictionary representing a single decision tree from an LGBM model.
+
+    Returns:
+        A dictionary mapping interaction tuples to their corresponding coefficients for
+        the single tree.
+
+    References:
+        Gorji, Ali, Andisheh Amrollahi, and Andreas Krause.
+        "SHAP values via sparse Fourier representation"
+        arXiv preprint arXiv:2410.06300 (2024).
+    """
+
+    def combine_coeffs(
+        left_coeffs: dict[tuple[int, ...], float],
+        right_coeffs: dict[tuple[int, ...], float],
+        feature_idx: int,
+    ) -> dict[tuple[int, ...], float]:
+        """Combines Fourier coefficients from the left and right children of a split node."""
+        combined_coeffs = {}
+        all_interactions = set(left_coeffs.keys()) | set(right_coeffs.keys())
+
+        for interaction in all_interactions:
+            left_val = left_coeffs.get(interaction, 0.0)
+            right_val = right_coeffs.get(interaction, 0.0)
+            combined_coeffs[interaction] = (left_val + right_val) / 2
+
+            new_interaction = tuple(sorted(set(interaction) | {feature_idx}))
+            combined_coeffs[new_interaction] = (left_val - right_val) / 2
+        return combined_coeffs
+
+    def dfs_traverse(node: dict[str, Any]) -> dict[tuple[int, ...], float]:
+        """Performs a depth-first traversal of the tree to compute coefficients."""
+        # Base case: if the node is a leaf, its function is a constant.
+        if "leaf_value" in node:
+            # The only non-zero coefficient is for the empty interaction (the bias term).
+            return {(): node["leaf_value"]}
+        # Recursive step: if the node is a split node.
+        left_coeffs = dfs_traverse(node["left_child"])
+        right_coeffs = dfs_traverse(node["right_child"])
+        feature_idx = node["split_feature"]
+        return combine_coeffs(left_coeffs, right_coeffs, feature_idx)
+
+    return dfs_traverse(tree_info["tree_structure"])

attribution/set_mm.py

@@ -0,0 +1,2397 @@
+"""
+set_mm.py  –  Cross-modal image×text attribution via mask-and-score.
+
+Extension of set.py that treats BOTH image regions AND caption tokens
+as players in the cooperative game, enabling cross-modal interactions.
+
+Players:
+  - Image regions:  seg_0 … seg_{n-1}   (from SAM or patch grid)
+  - Caption tokens:  tok_0 … tok_{m-1}   (from CLIP tokenizer)
+
+f(S) = CLIP_similarity(masked_image_keeping_S∩image, masked_caption_keeping_S∩text)
+
+This enables computing:
+  - Image-only Shapley values:     seg_i contribution
+  - Text-only Shapley values:      tok_j contribution
+  - Cross-modal interactions:      seg_i × tok_j  (image region × text token)
+
+Usage:
+    python set_mm.py                                  # default: UnSAM + all levels + caption tokens
+    python set_mm.py --mode patch --grid 4            # 4×4 grid
+    python set_mm.py --mode sam --max_regions 16      # SlimSAM segmentation
+    python set_mm.py --mode unsam --segment_level fine --max_regions 24
+    python set_mm.py --num_images 3 --method banzhaf  # 3 images, Banzhaf values
+    python set_mm.py --mask_style gray                # gray infill for image
+    python set_mm.py --mask_style lama                # LaMa inpainting
+"""
+
+from __future__ import annotations
+
+import argparse
+import csv
+import io
+import math
+import os
+import re
+import sys
+import time
+import warnings
+from collections import defaultdict
+from dataclasses import dataclass, field
+from itertools import combinations
+from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
+
+import numpy as np
+from PIL import Image, ImageFilter
+
+# ---------------------------------------------------------------------------
+# Lazy / optional imports
+# ---------------------------------------------------------------------------
+try:
+    import torch
+except ImportError:
+    torch = None
+
+try:
+    from transformers import pipeline as hf_pipeline
+    from transformers import SamModel, SamProcessor, CLIPModel, AutoProcessor, CLIPTokenizerFast
+    from transformers.image_utils import load_image as hf_load_image
+except ImportError:
+    hf_pipeline = None
+    SamModel = SamProcessor = CLIPModel = AutoProcessor = CLIPTokenizerFast = None
+    hf_load_image = None
+
+try:
+    from datasets import load_dataset
+except ImportError:
+    load_dataset = None
+
+try:
+    import lightgbm as lgb
+    from sklearn.model_selection import GridSearchCV, KFold
+except ImportError:
+    lgb = None
+    GridSearchCV = KFold = None
+
+try:
+    import pandas as pd
+except ImportError:
+    pd = None
+
+try:
+    from diffusers import (
+        AutoPipelineForInpainting,
+        LCMScheduler,
+        StableDiffusionInpaintPipeline,
+    )
+except ImportError:
+    AutoPipelineForInpainting = None
+    LCMScheduler = None
+    StableDiffusionInpaintPipeline = None
+
+try:
+    from diffusers import FluxImg2ImgPipeline
+except ImportError:
+    FluxImg2ImgPipeline = None
+
+try:
+    from simple_lama_inpainting import SimpleLama
+except ImportError:
+    SimpleLama = None
+
+try:
+    import cv2 as cv2_mod
+except ImportError:
+    cv2_mod = None
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# 0.  Configuration
+# ═══════════════════════════════════════════════════════════════════════════
+
+@dataclass
+class PipelineConfig:
+    """Central knob-box for the pipeline."""
+
+    # ---- dataset ----
+    coco_split: str = "train"
+    num_images: int = 1
+
+    # ---- featurisation (step 1) ----
+    mode: str = "unsam"                # "unsam" | "sam" | "patch"
+    sam_model_name: str = "yujunwei04/unsam-whole-image-segmentation"
+    sam_points_per_batch: int = 64
+    segment_level: str = "all"         # "all" | "coarse" | "medium" | "fine"
+    grid_size: int = 8
+    max_regions: int = 32
+
+    # ---- masking (steps 2-3) ----
+    mask_style: str = "blur"
+    blur_radius: int = 12
+    diffusion_model_name: str = "runwayml/stable-diffusion-inpainting"
+    diffusion_prompt: str = ""
+    diffusion_steps: int = 20
+    diffusion_local_crop: bool = True
+    diffusion_crop_pad: int = 48
+    diffusion_crop_min_side: int = 256
+    diffusion_local_max_cover: float = 0.7
+
+    # ---- LCM-LoRA + SDXL inpainting (fast) ----
+    lcm_base_model_name: str = "diffusers/stable-diffusion-xl-1.0-inpainting-0.1"
+    lcm_lora_id: str = "latent-consistency/lcm-lora-sdxl"
+    lcm_steps: int = 4
+    lcm_guidance_scale: float = 1.0
+
+    # ---- FLUX.1-Schnell img2img (fast) ----
+    flux_model_name: str = "black-forest-labs/FLUX.1-schnell"
+    flux_steps: int = 4
+    flux_guidance_scale: float = 0.0
+    flux_strength: float = 0.85
+    flux_max_sequence_length: int = 256
+
+    # ---- scoring / f(s) (step 4) ----
+    clip_model_name: str = "openai/clip-vit-base-patch32"
+
+    # ---- cross-modal text masking ----
+    text_mask_token: str = ""          # replace masked tokens with this (empty = UNK token)
+    max_tokens: int = 20               # cap caption tokens for tractability
+
+    # ---- ProxySPEX attribution ----
+    method: str = "shapley"
+    max_order: int = 2
+    top_k_interactions: int = 15
+    random_seed: int = 24
+    # random_seed: int = 42
+
+    # ---- runtime ----
+    device: str = "cuda" if (torch is not None and torch.cuda.is_available()) else "cpu"
+    record_mask_timing: bool = False
+    print_fourier_mobius: bool = False
+    print_coeff_top_k: int = 20
+
+
+@dataclass
+class MaskTimingRecorder:
+    """Collect per-mask timing and phase-level aggregates."""
+    cocoid: int
+    mask_style: str
+    out_dir: str
+    detail_rows: List[Dict[str, Any]] = field(default_factory=list)
+    phase_stats: Dict[str, Dict[str, float]] = field(default_factory=dict)
+
+    def add(
+        self,
+        *,
+        phase: str,
+        event: str,
+        total_s: float,
+        infill_s: float,
+        n_regions: int,
+        n_masked_regions: int,
+        masked_pixels: int,
+        image_hw: Tuple[int, int],
+        store_detail: bool = False,
+    ) -> None:
+        phase = phase or "unknown"
+        event = event or ""
+        height, width = image_hw
+        total_pixels = max(1, int(height) * int(width))
+        masked_ratio = float(masked_pixels) / float(total_pixels)
+
+        st = self.phase_stats.setdefault(
+            phase,
+            {
+                "count": 0.0,
+                "total_s_sum": 0.0,
+                "infill_s_sum": 0.0,
+                "total_s_min": float("inf"),
+                "total_s_max": 0.0,
+                "infill_s_min": float("inf"),
+                "infill_s_max": 0.0,
+                "masked_ratio_sum": 0.0,
+            },
+        )
+        st["count"] += 1.0
+        st["total_s_sum"] += float(total_s)
+        st["infill_s_sum"] += float(infill_s)
+        st["total_s_min"] = min(st["total_s_min"], float(total_s))
+        st["total_s_max"] = max(st["total_s_max"], float(total_s))
+        st["infill_s_min"] = min(st["infill_s_min"], float(infill_s))
+        st["infill_s_max"] = max(st["infill_s_max"], float(infill_s))
+        st["masked_ratio_sum"] += masked_ratio
+
+        if store_detail:
+            self.detail_rows.append(
+                {
+                    "cocoid": self.cocoid,
+                    "mask_style": self.mask_style,
+                    "phase": phase,
+                    "event": event,
+                    "n_regions": int(n_regions),
+                    "n_masked_regions": int(n_masked_regions),
+                    "masked_pixels": int(masked_pixels),
+                    "masked_ratio": masked_ratio,
+                    "total_ms": float(total_s) * 1000.0,
+                    "infill_ms": float(infill_s) * 1000.0,
+                    "non_infill_ms": max(float(total_s - infill_s), 0.0) * 1000.0,
+                }
+            )
+
+    def write_csvs(self) -> Tuple[str, str]:
+        os.makedirs(self.out_dir, exist_ok=True)
+        detail_path = os.path.join(
+            self.out_dir, f"coco_{self.cocoid}_mask_timing_details.csv"
+        )
+        summary_path = os.path.join(
+            self.out_dir, f"coco_{self.cocoid}_mask_timing_summary.csv"
+        )
+
+        detail_cols = [
+            "cocoid",
+            "mask_style",
+            "phase",
+            "event",
+            "n_regions",
+            "n_masked_regions",
+            "masked_pixels",
+            "masked_ratio",
+            "total_ms",
+            "infill_ms",
+            "non_infill_ms",
+        ]
+        with open(detail_path, "w", newline="") as f:
+            w = csv.DictWriter(f, fieldnames=detail_cols)
+            w.writeheader()
+            for row in self.detail_rows:
+                w.writerow(row)
+
+        summary_cols = [
+            "cocoid",
+            "mask_style",
+            "phase",
+            "calls",
+            "avg_total_ms",
+            "min_total_ms",
+            "max_total_ms",
+            "avg_infill_ms",
+            "min_infill_ms",
+            "max_infill_ms",
+            "avg_non_infill_ms",
+            "avg_masked_ratio",
+        ]
+        with open(summary_path, "w", newline="") as f:
+            w = csv.DictWriter(f, fieldnames=summary_cols)
+            w.writeheader()
+            for phase in sorted(self.phase_stats.keys()):
+                st = self.phase_stats[phase]
+                cnt = max(1.0, st["count"])
+                avg_total = st["total_s_sum"] / cnt
+                avg_infill = st["infill_s_sum"] / cnt
+                w.writerow(
+                    {
+                        "cocoid": self.cocoid,
+                        "mask_style": self.mask_style,
+                        "phase": phase,
+                        "calls": int(st["count"]),
+                        "avg_total_ms": avg_total * 1000.0,
+                        "min_total_ms": st["total_s_min"] * 1000.0,
+                        "max_total_ms": st["total_s_max"] * 1000.0,
+                        "avg_infill_ms": avg_infill * 1000.0,
+                        "min_infill_ms": st["infill_s_min"] * 1000.0,
+                        "max_infill_ms": st["infill_s_max"] * 1000.0,
+                        "avg_non_infill_ms": max(avg_total - avg_infill, 0.0) * 1000.0,
+                        "avg_masked_ratio": st["masked_ratio_sum"] / cnt,
+                    }
+                )
+        return detail_path, summary_path
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# 1.  Image featurisation  (reused from set.py)
+# ═══════════════════════════════════════════════════════════════════════════
+
+@dataclass(frozen=True)
+class ImageRegion:
+    """One region of an image."""
+    index: int
+    bbox: Tuple[int, int, int, int]
+    label: str
+    mask: Optional[np.ndarray] = field(default=None, repr=False)
+
+
+def _load_sam_pipeline(cfg: PipelineConfig):
+    if hf_pipeline is None:
+        raise ImportError("transformers is required")
+    model_name = (cfg.sam_model_name or "").lower()
+    use_image_seg_task = cfg.mode == "unsam" or "unsam" in model_name
+    task = "image-segmentation" if use_image_seg_task else "mask-generation"
+    kwargs: Dict[str, Any] = {
+        "model": cfg.sam_model_name,
+        "device": cfg.device,
+    }
+    if task == "mask-generation":
+        kwargs["points_per_batch"] = cfg.sam_points_per_batch
+    return hf_pipeline(task, **kwargs)
+
+
+def _normalize_bool_mask(raw_mask: Any, image_size: Tuple[int, int]) -> np.ndarray:
+    width, height = image_size
+    if isinstance(raw_mask, Image.Image):
+        arr = np.asarray(raw_mask)
+    elif torch is not None and isinstance(raw_mask, torch.Tensor):
+        arr = raw_mask.detach().cpu().numpy()
+    else:
+        arr = np.asarray(raw_mask)
+
+    if arr.ndim >= 3:
+        arr = arr[..., 0]
+    if arr.ndim != 2:
+        raise ValueError(f"Unsupported mask shape: {arr.shape}")
+
+    if np.issubdtype(arr.dtype, np.floating):
+        mask = arr > 0.5
+    else:
+        mask = arr > 0
+
+    if mask.shape != (height, width):
+        resized = Image.fromarray(mask.astype(np.uint8) * 255, mode="L").resize((width, height), Image.NEAREST)
+        mask = np.asarray(resized) > 0
+    return mask
+
+
+def _extract_masks_and_scores(outputs: Any, image_size: Tuple[int, int]) -> List[Tuple[np.ndarray, float, int]]:
+    records: List[Tuple[np.ndarray, float, int]] = []
+
+    if isinstance(outputs, dict) and "masks" in outputs:
+        masks = outputs.get("masks") or []
+        scores = outputs.get("scores") or [1.0] * len(masks)
+        for i, raw_mask in enumerate(masks):
+            try:
+                mask = _normalize_bool_mask(raw_mask, image_size)
+            except Exception:
+                continue
+            if not mask.any():
+                continue
+            score = float(scores[i]) if i < len(scores) else 1.0
+            area = int(mask.sum())
+            records.append((mask, score, area))
+        return records
+
+    if isinstance(outputs, list):
+        for item in outputs:
+            if not isinstance(item, dict):
+                continue
+            raw_mask = item.get("mask")
+            if raw_mask is None:
+                raw_mask = item.get("segmentation")
+            if raw_mask is None:
+                continue
+            try:
+                mask = _normalize_bool_mask(raw_mask, image_size)
+            except Exception:
+                continue
+            if not mask.any():
+                continue
+            raw_score = (
+                item.get("score")
+                or item.get("confidence")
+                or item.get("stability_score")
+                or item.get("predicted_iou")
+                or 1.0
+            )
+            score = float(raw_score)
+            area = int(mask.sum())
+            records.append((mask, score, area))
+        return records
+
+    raise ValueError(f"Unsupported segmentation output format: {type(outputs)}")
+
+
+def _select_masks_by_level(
+    records: List[Tuple[np.ndarray, float, int]],
+    level: str,
+    max_regions: int,
+) -> List[Tuple[np.ndarray, float, int]]:
+    if not records:
+        return []
+
+    level = (level or "all").lower()
+    by_score = sorted(records, key=lambda row: row[1], reverse=True)
+    if level == "all":
+        return by_score[:max_regions]
+
+    by_area_desc = sorted(records, key=lambda row: row[2], reverse=True)
+    n = len(by_area_desc)
+    if n < 3:
+        return by_score[:max_regions]
+
+    cut1 = max(1, n // 3)
+    cut2 = max(cut1 + 1, (2 * n) // 3)
+    buckets: Dict[str, List[Tuple[np.ndarray, float, int]]] = {
+        "coarse": by_area_desc[:cut1],
+        "medium": by_area_desc[cut1:cut2],
+        "fine": by_area_desc[cut2:],
+    }
+    chosen = buckets.get(level) or by_area_desc
+    chosen_by_score = sorted(chosen, key=lambda row: row[1], reverse=True)
+    return chosen_by_score[:max_regions]
+
+
+def featurise_sam(image: Image.Image, cfg: PipelineConfig, sam_pipe=None) -> List[ImageRegion]:
+    if sam_pipe is None:
+        sam_pipe = _load_sam_pipeline(cfg)
+
+    if cfg.mode == "sam":
+        outputs = sam_pipe(image, points_per_batch=cfg.sam_points_per_batch)
+    else:
+        outputs = sam_pipe(image)
+
+    raw_records = _extract_masks_and_scores(outputs, image.size)
+    picked = _select_masks_by_level(raw_records, cfg.segment_level, cfg.max_regions)
+    regions = []
+    for rank, (m, _, _) in enumerate(picked):
+        regions.append(ImageRegion(rank, _mask_to_bbox(m), f"seg_{rank}", m))
+    seg_tag = "UnSAM" if cfg.mode == "unsam" else "SAM"
+    print(
+        f"  [{seg_tag}] {len(regions)} regions selected (level={cfg.segment_level}, raw={len(raw_records)})"
+    )
+    return regions
+
+
+def featurise_patch(image: Image.Image, cfg: PipelineConfig) -> List[ImageRegion]:
+    w, h = image.size
+    g = cfg.grid_size
+    xe = _linspace_edges(w, g)
+    ye = _linspace_edges(h, g)
+    regions = []
+    idx = 0
+    for r in range(len(ye) - 1):
+        for c in range(len(xe) - 1):
+            regions.append(ImageRegion(idx, (xe[c], ye[r], xe[c+1], ye[r+1]), f"patch_{r}_{c}"))
+            idx += 1
+    print(f"  [Patch] {len(regions)} regions ({g}×{g} grid)")
+    return regions
+
+
+def featurise(image, cfg, sam_pipe=None):
+    if cfg.mode in ("sam", "unsam"):
+        return featurise_sam(image, cfg, sam_pipe=sam_pipe)
+    elif cfg.mode == "patch":
+        return featurise_patch(image, cfg)
+    else:
+        raise ValueError(f"Unknown mode: {cfg.mode}")
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# 2.  Caption tokenisation  –  caption → list of TokenPlayer
+# ═══════════════════════════════════════════════════════════════════════════
+
+@dataclass(frozen=True)
+class TokenPlayer:
+    """One caption token treated as a player in the game."""
+    index: int         # global player index (offset by n_image_regions)
+    token_pos: int     # position in the token sequence (1-based, skipping BOS)
+    token_id: int      # vocabulary ID
+    token_str: str     # decoded string
+    label: str         # display label, e.g. "tok:dog"
+
+
+def tokenise_caption(
+    caption: str,
+    processor: "AutoProcessor",
+    cfg: PipelineConfig,
+    offset: int = 0,
+) -> Tuple[List[TokenPlayer], List[int]]:
+    """
+    Tokenise caption and return a list of TokenPlayers (one per content token).
+
+    Args:
+        caption:   The caption string.
+        processor: CLIP processor (has .tokenizer).
+        cfg:       Config with max_tokens.
+        offset:    Global player index offset (= number of image regions).
+
+    Returns:
+        (players, full_token_ids)
+        - players: List[TokenPlayer] for content tokens (excluding BOS/EOS)
+        - full_token_ids: The complete token ID list including BOS/EOS
+    """
+    tokenizer = processor.tokenizer
+    encoding = tokenizer(caption, add_special_tokens=True, return_tensors="pt")
+    token_ids = encoding["input_ids"][0].tolist()  # includes BOS and EOS
+
+    # Decode each token individually
+    players: List[TokenPlayer] = []
+    # token_ids[0] = BOS, token_ids[-1] = EOS — skip both
+    content_ids = token_ids[1:-1]  # just content tokens
+    if len(content_ids) > cfg.max_tokens:
+        content_ids = content_ids[:cfg.max_tokens]
+
+    for i, tid in enumerate(content_ids):
+        decoded = tokenizer.decode([tid]).strip()
+        # Clean up subword artifacts
+        if not decoded:
+            decoded = f"[{tid}]"
+        lbl = f"tok:{decoded}"
+        players.append(TokenPlayer(
+            index=offset + i,
+            token_pos=i + 1,      # 1-based position (after BOS)
+            token_id=tid,
+            token_str=decoded,
+            label=lbl,
+        ))
+
+    # Build full_token_ids = BOS + content (possibly truncated) + EOS
+    full_ids = [token_ids[0]] + content_ids + [token_ids[-1]]
+
+    print(f"  [Tokens] {len(players)} caption tokens: {[p.token_str for p in players]}")
+    return players, full_ids
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# 3.  Image masking  (same as set.py)
+# ════════════════════════════════��══════════════════════════════════════════
+
+_diffusion_pipe_cache: Dict[str, Any] = {}
+_PRINT_FOURIER_MOBIUS = False
+_PRINT_COEF_TOP_K = 20
+_lama_pipe_cache: Dict[str, Any] = {}
+_lcm_pipe_cache: Dict[str, Any] = {}
+_flux_schnell_pipe_cache: Dict[str, Any] = {}
+
+
+def _get_diffusion_pipe(model_name: str, device: str):
+    if model_name in _diffusion_pipe_cache:
+        return _diffusion_pipe_cache[model_name]
+    if StableDiffusionInpaintPipeline is None:
+        raise ImportError("diffusers is required for diffusion inpainting")
+    if torch is None:
+        raise ImportError("torch is required for diffusion inpainting")
+    pipe = StableDiffusionInpaintPipeline.from_pretrained(
+        model_name, torch_dtype=torch.float16 if device != "cpu" else torch.float32,
+    ).to(device)
+    pipe.safety_checker = None
+    _diffusion_pipe_cache[model_name] = pipe
+    return pipe
+
+
+def _get_lama_pipe(device: str):
+    cache_key = device or "cpu"
+    if cache_key in _lama_pipe_cache:
+        return _lama_pipe_cache[cache_key]
+    if SimpleLama is None:
+        raise ImportError(
+            "simple-lama-inpainting is required for LaMa inpainting. "
+            "Install with: pip install simple-lama-inpainting"
+        )
+    try:
+        pipe = SimpleLama(device=device)
+    except TypeError:
+        pipe = SimpleLama()
+    _lama_pipe_cache[cache_key] = pipe
+    return pipe
+
+
+def _select_dtype(device: str):
+    if torch is None:
+        return None
+    if device == "cpu":
+        return torch.float32
+    if hasattr(torch, "bfloat16"):
+        return torch.bfloat16
+    return torch.float16
+
+
+def _get_lcm_pipe(base_model_name: str, lcm_lora_id: str, device: str):
+    cache_key = f"{base_model_name}::{lcm_lora_id}::{device}"
+    if cache_key in _lcm_pipe_cache:
+        return _lcm_pipe_cache[cache_key]
+    if AutoPipelineForInpainting is None:
+        raise ImportError("diffusers AutoPipelineForInpainting is required for LCM-LoRA SDXL")
+    if LCMScheduler is None:
+        raise ImportError("LCMScheduler is required for LCM-LoRA SDXL")
+    if torch is None:
+        raise ImportError("torch is required for LCM-LoRA SDXL")
+
+    pipe = AutoPipelineForInpainting.from_pretrained(
+        base_model_name,
+        torch_dtype=_select_dtype(device),
+    ).to(device)
+
+    if hasattr(pipe, "safety_checker"):
+        pipe.safety_checker = None
+
+    # LCM-LoRA can be stacked on top of SDXL inpaint without re-distillation.
+    pipe.load_lora_weights(lcm_lora_id)
+    pipe.fuse_lora()
+    pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
+
+    _lcm_pipe_cache[cache_key] = pipe
+    return pipe
+
+
+def _lcm_inpaint(image, binary_mask, cfg):
+    pipe = _get_lcm_pipe(cfg.lcm_base_model_name, cfg.lcm_lora_id, cfg.device)
+    prompt = cfg.diffusion_prompt or "high quality photo, realistic"
+
+    if not binary_mask.any():
+        return image
+
+    use_local_crop = bool(getattr(cfg, "diffusion_local_crop", True))
+    local_max_cover = float(getattr(cfg, "diffusion_local_max_cover", 0.7))
+    local_max_cover = min(1.0, max(0.0, local_max_cover))
+    mask_cover = float(binary_mask.mean())
+
+    lcm_steps = int(getattr(cfg, "lcm_steps", 4))
+    lcm_guidance = float(getattr(cfg, "lcm_guidance_scale", 1.0))
+
+    if (not use_local_crop) or mask_cover >= local_max_cover:
+        orig_size = image.size
+        img_r = image.resize((1024, 1024), Image.LANCZOS)
+        mask_pil = Image.fromarray((binary_mask.astype(np.uint8) * 255), mode="L").resize((1024, 1024), Image.NEAREST)
+        with torch.inference_mode():
+            result = pipe(
+                prompt=prompt,
+                image=img_r,
+                mask_image=mask_pil,
+                num_inference_steps=lcm_steps,
+                guidance_scale=lcm_guidance,
+            ).images[0]
+        return result.resize(orig_size, Image.LANCZOS)
+
+    h, w = binary_mask.shape
+    x0, y0, x1, y1 = _mask_to_bbox(binary_mask)
+    crop_pad = max(0, int(getattr(cfg, "diffusion_crop_pad", 48)))
+    x0, y0 = max(0, x0 - crop_pad), max(0, y0 - crop_pad)
+    x1, y1 = min(w, x1 + crop_pad), min(h, y1 + crop_pad)
+    if x1 <= x0 or y1 <= y0:
+        return image
+
+    crop_img = image.crop((x0, y0, x1, y1))
+    crop_mask = binary_mask[y0:y1, x0:x1]
+    if not crop_mask.any():
+        return image
+
+    rw, rh = crop_img.size
+    min_side = max(8, int(getattr(cfg, "diffusion_crop_min_side", 256)))
+    scale = 1.0
+    short_side = max(1, min(rw, rh))
+    if short_side < min_side:
+        scale = float(min_side) / float(short_side)
+
+    tw = max(8, int(round((rw * scale) / 8.0) * 8))
+    th = max(8, int(round((rh * scale) / 8.0) * 8))
+    img_r = crop_img.resize((tw, th), Image.LANCZOS)
+    mask_pil = Image.fromarray((crop_mask.astype(np.uint8) * 255), mode="L").resize((tw, th), Image.NEAREST)
+
+    with torch.inference_mode():
+        result = pipe(
+            prompt=prompt,
+            image=img_r,
+            mask_image=mask_pil,
+            num_inference_steps=lcm_steps,
+            guidance_scale=lcm_guidance,
+        ).images[0]
+
+    result_crop = result.resize((rw, rh), Image.LANCZOS)
+    out = np.asarray(image).copy()
+    roi = out[y0:y1, x0:x1]
+    result_arr = np.asarray(result_crop)
+
+    mh, mw = crop_mask.shape
+    rh2, rw2 = result_arr.shape[:2]
+    if (rh2, rw2) != (mh, mw):
+        result_arr = np.asarray(result_crop.resize((mw, mh), Image.LANCZOS))
+        rh2, rw2 = result_arr.shape[:2]
+    if roi.shape[:2] != (mh, mw) or (rh2, rw2) != (mh, mw):
+        hh = min(roi.shape[0], mh, rh2)
+        ww = min(roi.shape[1], mw, rw2)
+        roi = roi[:hh, :ww]
+        result_arr = result_arr[:hh, :ww]
+        crop_mask = crop_mask[:hh, :ww]
+
+    roi[crop_mask] = result_arr[crop_mask]
+    out[y0:y1, x0:x1] = roi
+    return Image.fromarray(out)
+
+
+def _get_flux_schnell_pipe(model_name: str, device: str):
+    cache_key = f"{model_name}::{device}"
+    if cache_key in _flux_schnell_pipe_cache:
+        return _flux_schnell_pipe_cache[cache_key]
+    if FluxImg2ImgPipeline is None:
+        raise ImportError("diffusers FluxImg2ImgPipeline is required for FLUX.1-Schnell")
+    if torch is None:
+        raise ImportError("torch is required for FLUX.1-Schnell")
+
+    pipe = FluxImg2ImgPipeline.from_pretrained(
+        model_name,
+        torch_dtype=_select_dtype(device),
+    ).to(device)
+    _flux_schnell_pipe_cache[cache_key] = pipe
+    return pipe
+
+
+def _run_flux_schnell(pipe, image: Image.Image, cfg):
+    prompt = cfg.diffusion_prompt or "high quality photo, realistic"
+    kwargs = {
+        "prompt": prompt,
+        "image": image,
+        "strength": float(getattr(cfg, "flux_strength", 0.85)),
+        "num_inference_steps": int(getattr(cfg, "flux_steps", 4)),
+        "guidance_scale": float(getattr(cfg, "flux_guidance_scale", 0.0)),
+        "max_sequence_length": int(getattr(cfg, "flux_max_sequence_length", 256)),
+    }
+    with torch.inference_mode():
+        try:
+            return pipe(**kwargs).images[0]
+        except TypeError:
+            kwargs.pop("max_sequence_length", None)
+            return pipe(**kwargs).images[0]
+
+
+def _flux_schnell_inpaint(image, binary_mask, cfg):
+    pipe = _get_flux_schnell_pipe(cfg.flux_model_name, cfg.device)
+
+    if not binary_mask.any():
+        return image
+
+    use_local_crop = bool(getattr(cfg, "diffusion_local_crop", True))
+    local_max_cover = float(getattr(cfg, "diffusion_local_max_cover", 0.7))
+    local_max_cover = min(1.0, max(0.0, local_max_cover))
+    mask_cover = float(binary_mask.mean())
+
+    if (not use_local_crop) or mask_cover >= local_max_cover:
+        gen = _run_flux_schnell(pipe, image, cfg)
+        out = np.asarray(image).copy()
+        gen_arr = np.asarray(gen.resize(image.size, Image.LANCZOS))
+        out[binary_mask] = gen_arr[binary_mask]
+        return Image.fromarray(out)
+
+    h, w = binary_mask.shape
+    x0, y0, x1, y1 = _mask_to_bbox(binary_mask)
+    crop_pad = max(0, int(getattr(cfg, "diffusion_crop_pad", 48)))
+    x0, y0 = max(0, x0 - crop_pad), max(0, y0 - crop_pad)
+    x1, y1 = min(w, x1 + crop_pad), min(h, y1 + crop_pad)
+    if x1 <= x0 or y1 <= y0:
+        return image
+
+    crop_img = image.crop((x0, y0, x1, y1))
+    crop_mask = binary_mask[y0:y1, x0:x1]
+    if not crop_mask.any():
+        return image
+
+    result_crop = _run_flux_schnell(pipe, crop_img, cfg).resize(crop_img.size, Image.LANCZOS)
+
+    out = np.asarray(image).copy()
+    roi = out[y0:y1, x0:x1]
+    result_arr = np.asarray(result_crop)
+
+    mh, mw = crop_mask.shape
+    rh2, rw2 = result_arr.shape[:2]
+    if (rh2, rw2) != (mh, mw):
+        result_arr = np.asarray(result_crop.resize((mw, mh), Image.LANCZOS))
+        rh2, rw2 = result_arr.shape[:2]
+    if roi.shape[:2] != (mh, mw) or (rh2, rw2) != (mh, mw):
+        hh = min(roi.shape[0], mh, rh2)
+        ww = min(roi.shape[1], mw, rw2)
+        roi = roi[:hh, :ww]
+        result_arr = result_arr[:hh, :ww]
+        crop_mask = crop_mask[:hh, :ww]
+
+    roi[crop_mask] = result_arr[crop_mask]
+    out[y0:y1, x0:x1] = roi
+    return Image.fromarray(out)
+
+
+def _diffusion_inpaint(image, binary_mask, cfg):
+    pipe = _get_diffusion_pipe(cfg.diffusion_model_name, cfg.device)
+    prompt = cfg.diffusion_prompt or "high quality photo, realistic"
+
+    if not binary_mask.any():
+        return image
+
+    use_local_crop = bool(getattr(cfg, "diffusion_local_crop", True))
+    local_max_cover = float(getattr(cfg, "diffusion_local_max_cover", 0.7))
+    local_max_cover = min(1.0, max(0.0, local_max_cover))
+    mask_cover = float(binary_mask.mean())
+
+    if (not use_local_crop) or mask_cover >= local_max_cover:
+        orig_size = image.size
+        img_r = image.resize((512, 512), Image.LANCZOS)
+        mask_pil = Image.fromarray((binary_mask.astype(np.uint8) * 255), mode="L").resize((512, 512), Image.NEAREST)
+        with torch.inference_mode():
+            result = pipe(
+                prompt=prompt,
+                image=img_r,
+                mask_image=mask_pil,
+                num_inference_steps=cfg.diffusion_steps,
+                guidance_scale=7.5,
+            ).images[0]
+        return result.resize(orig_size, Image.LANCZOS)
+
+    h, w = binary_mask.shape
+    x0, y0, x1, y1 = _mask_to_bbox(binary_mask)
+    crop_pad = max(0, int(getattr(cfg, "diffusion_crop_pad", 48)))
+    x0, y0 = max(0, x0 - crop_pad), max(0, y0 - crop_pad)
+    x1, y1 = min(w, x1 + crop_pad), min(h, y1 + crop_pad)
+    if x1 <= x0 or y1 <= y0:
+        return image
+
+    crop_img = image.crop((x0, y0, x1, y1))
+    crop_mask = binary_mask[y0:y1, x0:x1]
+    if not crop_mask.any():
+        return image
+
+    rw, rh = crop_img.size
+    min_side = max(8, int(getattr(cfg, "diffusion_crop_min_side", 256)))
+    scale = 1.0
+    short_side = max(1, min(rw, rh))
+    if short_side < min_side:
+        scale = float(min_side) / float(short_side)
+
+    tw = max(8, int(round((rw * scale) / 8.0) * 8))
+    th = max(8, int(round((rh * scale) / 8.0) * 8))
+    img_r = crop_img.resize((tw, th), Image.LANCZOS)
+    mask_pil = Image.fromarray((crop_mask.astype(np.uint8) * 255), mode="L").resize((tw, th), Image.NEAREST)
+
+    with torch.inference_mode():
+        result = pipe(
+            prompt=prompt,
+            image=img_r,
+            mask_image=mask_pil,
+            num_inference_steps=cfg.diffusion_steps,
+            guidance_scale=7.5,
+        ).images[0]
+
+    result_crop = result.resize((rw, rh), Image.LANCZOS)
+    out = np.asarray(image).copy()
+    roi = out[y0:y1, x0:x1]
+    result_arr = np.asarray(result_crop)
+
+    # Guard against occasional off-by-one patch sizes from inpainting backends.
+    mh, mw = crop_mask.shape
+    rh2, rw2 = result_arr.shape[:2]
+    if (rh2, rw2) != (mh, mw):
+        result_arr = np.asarray(result_crop.resize((mw, mh), Image.LANCZOS))
+        rh2, rw2 = result_arr.shape[:2]
+    if roi.shape[:2] != (mh, mw) or (rh2, rw2) != (mh, mw):
+        hh = min(roi.shape[0], mh, rh2)
+        ww = min(roi.shape[1], mw, rw2)
+        roi = roi[:hh, :ww]
+        result_arr = result_arr[:hh, :ww]
+        crop_mask = crop_mask[:hh, :ww]
+
+    roi[crop_mask] = result_arr[crop_mask]
+    out[y0:y1, x0:x1] = roi
+    return Image.fromarray(out)
+
+
+def _lama_inpaint(image, binary_mask, cfg):
+    pipe = _get_lama_pipe(cfg.device)
+
+    if not binary_mask.any():
+        return image
+
+    use_local_crop = bool(getattr(cfg, "diffusion_local_crop", True))
+    local_max_cover = float(getattr(cfg, "diffusion_local_max_cover", 0.7))
+    local_max_cover = min(1.0, max(0.0, local_max_cover))
+    mask_cover = float(binary_mask.mean())
+
+    if (not use_local_crop) or mask_cover >= local_max_cover:
+        mask_pil = Image.fromarray((binary_mask.astype(np.uint8) * 255), mode="L")
+        return pipe(image, mask_pil)
+
+    h, w = binary_mask.shape
+    x0, y0, x1, y1 = _mask_to_bbox(binary_mask)
+    crop_pad = max(0, int(getattr(cfg, "diffusion_crop_pad", 48)))
+    x0, y0 = max(0, x0 - crop_pad), max(0, y0 - crop_pad)
+    x1, y1 = min(w, x1 + crop_pad), min(h, y1 + crop_pad)
+    if x1 <= x0 or y1 <= y0:
+        return image
+
+    crop_img = image.crop((x0, y0, x1, y1))
+    crop_mask = binary_mask[y0:y1, x0:x1]
+    if not crop_mask.any():
+        return image
+
+    mask_pil = Image.fromarray((crop_mask.astype(np.uint8) * 255), mode="L")
+    result_crop = pipe(crop_img, mask_pil)
+
+    out = np.asarray(image).copy()
+    roi = out[y0:y1, x0:x1]
+    result_arr = np.asarray(result_crop)
+
+    # Guard against occasional off-by-one patch sizes from inpainting backends.
+    mh, mw = crop_mask.shape
+    rh2, rw2 = result_arr.shape[:2]
+    if (rh2, rw2) != (mh, mw):
+        result_arr = np.asarray(result_crop.resize((mw, mh), Image.LANCZOS))
+        rh2, rw2 = result_arr.shape[:2]
+    if roi.shape[:2] != (mh, mw) or (rh2, rw2) != (mh, mw):
+        hh = min(roi.shape[0], mh, rh2)
+        ww = min(roi.shape[1], mw, rw2)
+        roi = roi[:hh, :ww]
+        result_arr = result_arr[:hh, :ww]
+        crop_mask = crop_mask[:hh, :ww]
+
+    roi[crop_mask] = result_arr[crop_mask]
+    out[y0:y1, x0:x1] = roi
+    return Image.fromarray(out)
+
+
+def _cv2_inpaint(image, binary_mask, method="telea"):
+    if cv2_mod is None:
+        raise ImportError("opencv-python required for cv2 inpainting")
+    arr = cv2_mod.cvtColor(np.asarray(image), cv2_mod.COLOR_RGB2BGR)
+    mask_u8 = binary_mask.astype(np.uint8) * 255
+    flag = cv2_mod.INPAINT_TELEA if method == "telea" else cv2_mod.INPAINT_NS
+    res = cv2_mod.inpaint(arr, mask_u8, inpaintRadius=5, flags=flag)
+    return Image.fromarray(cv2_mod.cvtColor(res, cv2_mod.COLOR_BGR2RGB))
+
+
+def apply_image_mask(
+    image,
+    regions,
+    coalition_mask,
+    style="blur",
+    blur_radius=12,
+    cfg=None,
+    timing_meta: Optional[Dict[str, Any]] = None,
+):
+    """Mask image regions based on coalition vector (1=keep, 0=mask)."""
+    t0 = time.perf_counter()
+    image = _ensure_rgb(image)
+    arr = np.asarray(image).copy()
+    h, w = arr.shape[:2]
+    style = (style or "blur").lower()
+    infill_s = 0.0
+
+    if len(coalition_mask) != len(regions):
+        raise ValueError(f"coalition length {len(coalition_mask)} != regions {len(regions)}")
+
+    n_regions = len(regions)
+    n_masked_regions = sum(1 for v in coalition_mask if not int(v))
+
+    if style in ("diffusion", "lama", "lcm", "flux", "flux_schnell", "diffusion_cv2", "cv2_telea", "cv2_ns"):
+        combined = np.zeros((h, w), dtype=bool)
+        for i, reg in enumerate(regions):
+            if not int(coalition_mask[i]):
+                combined |= _region_mask(reg, (h, w))
+        masked_pixels = int(combined.sum())
+        if not combined.any():
+            out_img = image
+        elif style == "diffusion":
+            t_infill = time.perf_counter()
+            out_img = _diffusion_inpaint(image, combined, cfg)
+            infill_s = time.perf_counter() - t_infill
+        elif style == "lama":
+            t_infill = time.perf_counter()
+            out_img = _lama_inpaint(image, combined, cfg)
+            infill_s = time.perf_counter() - t_infill
+        elif style == "lcm":
+            t_infill = time.perf_counter()
+            out_img = _lcm_inpaint(image, combined, cfg)
+            infill_s = time.perf_counter() - t_infill
+        elif style in ("flux", "flux_schnell"):
+            t_infill = time.perf_counter()
+            out_img = _flux_schnell_inpaint(image, combined, cfg)
+            infill_s = time.perf_counter() - t_infill
+        elif style in ("diffusion_cv2", "cv2_telea"):
+            t_infill = time.perf_counter()
+            out_img = _cv2_inpaint(image, combined, method="telea")
+            infill_s = time.perf_counter() - t_infill
+        else:  # cv2_ns
+            t_infill = time.perf_counter()
+            out_img = _cv2_inpaint(image, combined, method="ns")
+            infill_s = time.perf_counter() - t_infill
+        total_s = time.perf_counter() - t0
+        if timing_meta and timing_meta.get("recorder") is not None:
+            timing_meta["recorder"].add(
+                phase=str(timing_meta.get("phase") or "apply_image_mask"),
+                event=str(timing_meta.get("event") or ""),
+                total_s=total_s,
+                infill_s=infill_s,
+                n_regions=n_regions,
+                n_masked_regions=n_masked_regions,
+                masked_pixels=masked_pixels,
+                image_hw=(h, w),
+                store_detail=bool(timing_meta.get("store_detail", False)),
+            )
+        return out_img
+
+    if style == "blur":
+        fill_mode = "blur"
+    elif style == "gray":
+        fill_mode = "gray"
+    elif style == "mean":
+        fill_mode = "mean"
+    else:
+        raise ValueError(f"Unknown mask_style: {style}")
+
+    combined = np.zeros((h, w), dtype=bool)
+    for i, reg in enumerate(regions):
+        if int(coalition_mask[i]):
+            continue
+        rmask = _region_mask(reg, (h, w))
+        if rmask.any():
+            combined |= rmask
+
+    masked_pixels = int(combined.sum())
+    if combined.any():
+        t_infill = time.perf_counter()
+        if fill_mode == "blur":
+            fill = np.asarray(image.filter(ImageFilter.GaussianBlur(radius=blur_radius)))
+        elif fill_mode == "gray":
+            fill = np.full_like(arr, 128)
+        else:
+            fill = np.broadcast_to(arr.mean(axis=(0, 1)).astype(np.uint8), arr.shape).copy()
+        arr[combined] = fill[combined]
+        infill_s = time.perf_counter() - t_infill
+
+    out_img = Image.fromarray(arr)
+    total_s = time.perf_counter() - t0
+    if timing_meta and timing_meta.get("recorder") is not None:
+        timing_meta["recorder"].add(
+            phase=str(timing_meta.get("phase") or "apply_image_mask"),
+            event=str(timing_meta.get("event") or ""),
+            total_s=total_s,
+            infill_s=infill_s,
+            n_regions=n_regions,
+            n_masked_regions=n_masked_regions,
+            masked_pixels=masked_pixels,
+            image_hw=(h, w),
+            store_detail=bool(timing_meta.get("store_detail", False)),
+        )
+    return out_img
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# 4.  Text masking  –  mask caption tokens by replacing with PAD
+# ═══════════════════════════════════════════════════════════════════════════
+
+def mask_token_ids(
+    full_token_ids: List[int],
+    token_players: List[TokenPlayer],
+    token_coalition: Sequence[int],
+    unk_token_id: int,
+) -> List[int]:
+    """
+    Given the full token ID list (BOS + content + EOS) and a binary coalition
+    over content tokens, replace masked tokens with unk_token_id.
+
+    Returns a new token ID list of the same length.
+    """
+    # full_token_ids = [BOS, tok1, tok2, ..., tokN, EOS]
+    # token_players have token_pos = 1..N (1-based positions in full_token_ids)
+    result = list(full_token_ids)
+    for player, keep in zip(token_players, token_coalition):
+        if not int(keep):
+            result[player.token_pos] = unk_token_id
+    return result
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# 5.  Cross-modal CLIP scorer
+# ═══════════════════════════════════════════════════════════════════════════
+
+def _make_open_clip_scorer(model_name: str, device: str):
+    """
+    Load an open_clip model (e.g. BiomedCLIP) and return (model_adapter, processor_adapter)
+    whose interface matches what CrossModalCLIPScorer expects from a transformers CLIPModel.
+    """
+    try:
+        import open_clip
+    except ImportError:
+        raise ImportError(
+            "open_clip is required for BiomedCLIP. "
+            "Install with: pip install open-clip-torch"
+        )
+
+    hf_hub_id = model_name if model_name.startswith("hf-hub:") else f"hf-hub:{model_name}"
+    oc_model, _, preprocess_val = open_clip.create_model_and_transforms(hf_hub_id)
+    oc_tokenizer = open_clip.get_tokenizer(hf_hub_id)
+    oc_model = oc_model.to(device)
+    oc_model.eval()
+
+    # Dict subclass that supports .to(device) like HF BatchFeature
+    class _TensorDict(dict):
+        def to(self, d):
+            return _TensorDict(
+                {k: v.to(d) if isinstance(v, torch.Tensor) else v for k, v in self.items()}
+            )
+
+    # Tokenizer adapter — delegates to the underlying HF tokenizer when available
+    class _TokAdapter:
+        def __init__(self, oc_tok):
+            self._hf = getattr(oc_tok, "tokenizer", None)
+            self.unk_token_id = getattr(self._hf, "unk_token_id", None) or 100
+            self.pad_token_id = getattr(self._hf, "pad_token_id", None) or 0
+
+        def __call__(self, text, add_special_tokens=True, return_tensors=None):
+            if self._hf is not None:
+                return self._hf(
+                    text, add_special_tokens=add_special_tokens, return_tensors=return_tensors
+                )
+            # Fallback: use open_clip tokenizer directly
+            ids = oc_tokenizer([text])  # (1, context_length)
+            return {"input_ids": ids}
+
+        def decode(self, ids, skip_special_tokens=False):
+            if self._hf is not None:
+                return self._hf.decode(ids, skip_special_tokens=skip_special_tokens)
+            return ""
+
+        def encode(self, text, add_special_tokens=False):
+            if self._hf is not None:
+                return self._hf.encode(text, add_special_tokens=add_special_tokens)
+            return []
+
+    # Model adapter — wraps open_clip encode_image / encode_text
+    class _ModelAdapter:
+        def __init__(self, m):
+            self._m = m
+            self.logit_scale = m.logit_scale
+
+        def get_image_features(self, pixel_values):
+            return self._m.encode_image(pixel_values)
+
+        def get_text_features(self, input_ids, attention_mask=None):
+            try:
+                return self._m.encode_text(input_ids, attention_mask=attention_mask)
+            except TypeError:
+                return self._m.encode_text(input_ids)
+
+        def eval(self):
+            self._m.eval()
+            return self
+
+        def to(self, d):
+            self._m = self._m.to(d)
+            return self
+
+    # Processor adapter — wraps torchvision preprocess_val
+    class _ProcAdapter:
+        def __init__(self, pre, tok):
+            self._pre = pre
+            self.tokenizer = tok
+
+        def __call__(self, images=None, return_tensors="pt"):
+            if isinstance(images, list):
+                tensors = torch.stack([self._pre(img) for img in images])
+            else:
+                tensors = self._pre(images).unsqueeze(0)
+            return _TensorDict({"pixel_values": tensors})
+
+    tok_adapter = _TokAdapter(oc_tokenizer)
+    return _ModelAdapter(oc_model), _ProcAdapter(preprocess_val, tok_adapter)
+
+
+class CrossModalCLIPScorer:
+    """
+    CLIP scorer that supports masking BOTH image regions AND text tokens.
+
+    Players are indexed as:
+      [0 .. n_img-1]  = image regions
+      [n_img .. n_img+n_tok-1]  = caption tokens
+    """
+
+    def __init__(self, cfg: PipelineConfig):
+        self.device = cfg.device
+        # Detect open_clip models (BiomedCLIP etc.) by name pattern
+        _use_open_clip = (
+            "BiomedCLIP" in cfg.clip_model_name
+            or cfg.clip_model_name.startswith("hf-hub:")
+        )
+        if _use_open_clip:
+            self.model, self.processor = _make_open_clip_scorer(
+                cfg.clip_model_name, self.device
+            )
+        else:
+            if CLIPModel is None:
+                raise ImportError("transformers is required")
+            self.model = CLIPModel.from_pretrained(cfg.clip_model_name).to(self.device)
+            self.processor = AutoProcessor.from_pretrained(cfg.clip_model_name)
+            self.model.eval()
+        # Get UNK token id for text masking
+        self.unk_token_id = self.processor.tokenizer.unk_token_id
+        if self.unk_token_id is None:
+            self.unk_token_id = 0  # fallback
+        # Get PAD token id for true padding (aligning sequence lengths)
+        self.pad_token_id = self.processor.tokenizer.pad_token_id
+        if self.pad_token_id is None:
+            self.pad_token_id = 0  # fallback
+
+    def score_with_masked_tokens(
+        self,
+        image: Image.Image,
+        token_ids: List[int],
+        attention_mask: Optional[List[int]] = None,
+    ) -> float:
+        """
+        Score image against a specific token ID sequence (possibly masked).
+
+        This bypasses normal text processing and injects token IDs directly.
+        """
+        # Process image only
+        image_inputs = self.processor(images=image, return_tensors="pt").to(self.device)
+
+        # Build text inputs manually
+        input_ids = torch.tensor([token_ids], dtype=torch.long, device=self.device)
+        if attention_mask is None:
+            attention_mask = [1] * len(token_ids)
+        attn = torch.tensor([attention_mask], dtype=torch.long, device=self.device)
+
+        with torch.inference_mode():
+            image_embeds = self.model.get_image_features(pixel_values=image_inputs["pixel_values"])
+            text_embeds = self.model.get_text_features(input_ids=input_ids, attention_mask=attn)
+
+            # Normalize
+            image_embeds = image_embeds / image_embeds.norm(p=2, dim=-1, keepdim=True)
+            text_embeds = text_embeds / text_embeds.norm(p=2, dim=-1, keepdim=True)
+
+            # Cosine similarity scaled by CLIP's logit_scale
+            logit_scale = self.model.logit_scale.exp()
+            similarity = (image_embeds @ text_embeds.T) * logit_scale
+
+        return similarity[0, 0].item()
+
+    def score_batch_with_masked_tokens(
+        self,
+        images: List[Image.Image],
+        token_ids_list: List[List[int]],
+    ) -> List[float]:
+        """Batch score: multiple (image, token_ids) pairs."""
+        if not images:
+            return []
+
+        # Process all images
+        image_inputs = self.processor(images=images, return_tensors="pt").to(self.device)
+
+        # Pad token_ids to same length
+        max_len = max(len(t) for t in token_ids_list)
+        padded_ids = []
+        attn_masks = []
+        for tids in token_ids_list:
+            pad_len = max_len - len(tids)
+            padded_ids.append(tids + [self.pad_token_id] * pad_len)
+            attn_masks.append([1] * len(tids) + [0] * pad_len)
+
+        input_ids = torch.tensor(padded_ids, dtype=torch.long, device=self.device)
+        attn = torch.tensor(attn_masks, dtype=torch.long, device=self.device)
+
+        with torch.inference_mode():
+            image_embeds = self.model.get_image_features(pixel_values=image_inputs["pixel_values"])
+            text_embeds = self.model.get_text_features(input_ids=input_ids, attention_mask=attn)
+
+            image_embeds = image_embeds / image_embeds.norm(p=2, dim=-1, keepdim=True)
+            text_embeds = text_embeds / text_embeds.norm(p=2, dim=-1, keepdim=True)
+
+            logit_scale = self.model.logit_scale.exp()
+            # Per-pair similarity (diagonal)
+            sims = (image_embeds * text_embeds).sum(dim=-1) * logit_scale
+
+        return sims.cpu().tolist()
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# 6.  Cross-modal set function  f(S)
+# ═══════════════════════════════════════════════════════════════════════════
+
+def build_cross_modal_set_function(
+    image: Image.Image,
+    regions: List[ImageRegion],
+    token_players: List[TokenPlayer],
+    full_token_ids: List[int],
+    scorer: CrossModalCLIPScorer,
+    cfg: PipelineConfig,
+    timing_recorder: Optional[MaskTimingRecorder] = None,
+):
+    """
+    Build f : {0,1}^{n_img + n_tok} → ℝ
+
+    The coalition vector has:
+      - first n_img bits: image regions (1=keep, 0=mask/infill)
+      - next n_tok bits:  caption tokens (1=keep, 0=replace with PAD)
+
+    f(S) = CLIP_similarity(masked_image, masked_text)
+    """
+    n_img = len(regions)
+    n_tok = len(token_players)
+    n_total = n_img + n_tok
+    cache: Dict[Tuple[int, ...], float] = {}
+
+    def _key(sel):
+        return tuple(int(x) for x in sel)
+
+    def set_function(selection: Sequence[int]) -> float:
+        if len(selection) != n_total:
+            raise ValueError(f"Expected {n_total} bits, got {len(selection)}")
+        key = _key(selection)
+        if key in cache:
+            return cache[key]
+
+        img_coalition = list(selection[:n_img])
+        tok_coalition = list(selection[n_img:])
+
+        # Mask image
+        masked_img = apply_image_mask(
+            image, regions, img_coalition,
+            style=cfg.mask_style, blur_radius=cfg.blur_radius, cfg=cfg,
+            timing_meta={
+                "recorder": timing_recorder,
+                "phase": "set_function_single",
+                "event": "",
+                "store_detail": False,
+            },
+        )
+
+        # Mask text tokens
+        masked_tids = mask_token_ids(full_token_ids, token_players, tok_coalition, scorer.unk_token_id)
+
+        val = scorer.score_with_masked_tokens(masked_img, masked_tids)
+        cache[key] = val
+        return val
+
+    def batch_set_function(selections: Sequence[Sequence[int]]) -> List[float]:
+        num = len(selections)
+        if num == 0:
+            return []
+        results = [0.0] * num
+        pending = []  # (idx, key, sel)
+
+        for i, sel in enumerate(selections):
+            key = _key(sel)
+            if key in cache:
+                results[i] = cache[key]
+            else:
+                pending.append((i, key, sel))
+
+        batch_sz = 16
+        for start in range(0, len(pending), batch_sz):
+            batch = pending[start:start + batch_sz]
+            imgs = []
+            tids_list = []
+            for _, _, sel in batch:
+                img_coal = list(sel[:n_img])
+                tok_coal = list(sel[n_img:])
+                imgs.append(apply_image_mask(
+                    image, regions, img_coal,
+                    style=cfg.mask_style, blur_radius=cfg.blur_radius, cfg=cfg,
+                    timing_meta={
+                        "recorder": timing_recorder,
+                        "phase": "set_function_batch",
+                        "event": "",
+                        "store_detail": False,
+                    },
+                ))
+                tids_list.append(mask_token_ids(full_token_ids, token_players, tok_coal, scorer.unk_token_id))
+
+            scores = scorer.score_batch_with_masked_tokens(imgs, tids_list)
+            for (row_i, key, _), sc in zip(batch, scores):
+                cache[key] = sc
+                results[row_i] = sc
+        return results
+
+    set_function.batch = batch_set_function
+    return set_function
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# 7.  ProxySPEX  (same as set.py)
+# ═══════════════════════════════════════════════════════════════════════════
+
+def powerset(loc, max_order=None):
+    if max_order is None:
+        max_order = len(loc)
+    out = []
+    for r in range(max_order + 1):
+        for sub in combinations(loc, r):
+            out.append(sub)
+    return out
+
+
+def _print_coeff_scores(name: str, coeffs: Dict[Tuple[int, ...], float], top_k: int = 20) -> None:
+    """Pretty-print Fourier/Mobius coefficient scores in terminal."""
+    items = sorted(coeffs.items(), key=lambda kv: abs(float(kv[1])), reverse=True)
+    total = len(items)
+    if top_k > 0:
+        items = items[:top_k]
+    print(f"  [ProxySPEX] {name} scores: showing {len(items)}/{total} terms")
+    for loc, val in items:
+        label = "{}" if len(loc) == 0 else str(tuple(loc))
+        print(f"    {label:>16s}  {float(val):+.8f}")
+
+
+def fourier_to_mobius(fourier):
+    if not fourier:
+        return {}
+    unscaled = defaultdict(float)
+    for loc, coef in fourier.items():
+        real_coef = float(np.real(coef))
+        for sub in powerset(loc):
+            unscaled[sub] += real_coef
+    mobius = {loc: val * np.power(-2.0, len(loc)) for loc, val in unscaled.items() if abs(val) > 1e-12}
+    if _PRINT_FOURIER_MOBIUS:
+        _print_coeff_scores("Fourier", fourier, top_k=_PRINT_COEF_TOP_K)
+        _print_coeff_scores("Mobius", mobius, top_k=_PRINT_COEF_TOP_K)
+    return mobius
+
+
+def mobius_to_fourier(mobius):
+    """
+    Convert sparse Mobius coefficients to sparse Fourier coefficients.
+    """
+    if not mobius:
+        return {}
+    unscaled = defaultdict(float)
+    for loc, coef in mobius.items():
+        loc = tuple(sorted(loc))
+        scale = float(np.real(coef)) / (2 ** len(loc))
+        for sub in powerset(loc):
+            unscaled[sub] += scale
+    return {
+        loc: val * np.power(-1.0, len(loc))
+        for loc, val in unscaled.items()
+        if abs(val) > 1e-12
+    }
+
+
+def _lgboost_tree_to_fourier(tree_info):
+    def combine(left, right, feat_idx):
+        combined = {}
+        for inter in set(left) | set(right):
+            lv, rv = left.get(inter, 0.0), right.get(inter, 0.0)
+            combined[inter] = (lv + rv) / 2
+            new = tuple(sorted(set(inter) | {feat_idx}))
+            combined[new] = (lv - rv) / 2
+        return combined
+
+    def dfs(node):
+        if "leaf_value" in node:
+            return {(): node["leaf_value"]}
+        return combine(dfs(node["left_child"]), dfs(node["right_child"]), node["split_feature"])
+
+    return dfs(tree_info["tree_structure"])
+
+
+def lgboost_to_fourier(model_dict):
+    agg = defaultdict(float)
+    for ti in model_dict["tree_info"]:
+        for inter, val in _lgboost_tree_to_fourier(ti).items():
+            agg[inter] += val
+    return {k: v for k, v in agg.items() if v != 0.0}
+
+
+def proxyspex(n, budget, game, max_order=None, random_seed=0):
+    if lgb is None:
+        raise ImportError("lightgbm required")
+    if pd is None:
+        raise ImportError("pandas required")
+
+    rng = np.random.default_rng(random_seed)
+    X = rng.integers(0, 2, size=(budget, n), dtype=np.int8)
+    train_X = pd.DataFrame(X, columns=[f"f{i}" for i in range(n)]).astype(np.int8)
+    y = game(X)
+
+    if len(train_X) < 5:
+        model = lgb.LGBMRegressor(verbose=-1, n_jobs=1, random_state=random_seed, max_depth=1)
+        model.fit(train_X, y)
+        best = model
+    else:
+        base = lgb.LGBMRegressor(verbose=-1, n_jobs=1, random_state=random_seed)
+        param_grid = {"max_depth": [3, 5], "n_estimators": [500, 1000], "learning_rate": [0.01, 0.1]}
+        cv = KFold(n_splits=min(5, len(train_X)), shuffle=True, random_state=random_seed)
+        gs = GridSearchCV(base, param_grid, scoring="r2", cv=cv, verbose=0, n_jobs=1)
+        gs.fit(train_X, y)
+        best = gs.best_estimator_
+
+    fourier = lgboost_to_fourier(best.booster_.dump_model())
+    mobius = fourier_to_mobius(fourier)
+    max_order = max_order if max_order is not None else n
+    return {k: v for k, v in mobius.items() if len(k) <= max_order}
+
+
+def run_proxyspex(set_function, features, max_order=3, seed=None):
+    n = len(features)
+    if n == 0:
+        return {}
+
+    def game(coalitions):
+        arr = np.asarray(coalitions)
+        if arr.ndim == 1:
+            out = set_function(tuple(int(x) for x in arr))
+            return out - np.mean(out)
+        if arr.ndim == 2:
+            if hasattr(set_function, "batch"):
+                out = np.asarray(set_function.batch(arr), dtype=float)
+            else:
+                out = np.array([set_function(tuple(int(x) for x in row)) for row in arr], dtype=float)
+            return out - np.mean(out)
+        raise ValueError("coalitions must be 1-D or 2-D")
+
+    order = max(1, min(max_order, n))
+    budget = int(10 * n * math.log(n + 1, 2))
+    print(f"  [ProxySPEX] n={n}, budget={budget}, max_order={order}")
+    return proxyspex(n, budget, game, max_order=order, random_seed=seed or 0)
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# 8.  Möbius → Shapley / Banzhaf  (same as set.py)
+# ═══════════════════════════════════════════════════════════════════════════
+
+def mobius_to_shapley(mobius):
+    sii = {}
+    for loc, coef in mobius.items():
+        loc = tuple(sorted(loc))
+        t = len(loc)
+        for sub in powerset(loc, max_order=1):
+            sii[sub] = sii.get(sub, 0.0) + float(np.real(coef)) / (1 + t - len(sub))
+    return {k: v for k, v in sii.items() if len(k) == 1 and not math.isclose(v, 0.0)}
+
+
+def mobius_to_banzhaf(mobius):
+    bii = {}
+    for loc, coef in mobius.items():
+        loc = tuple(sorted(loc))
+        t = len(loc)
+        for sub in powerset(loc, max_order=1):
+            bii[sub] = bii.get(sub, 0.0) + float(np.real(coef)) / (2 ** (t - len(sub)))
+    return {k: v for k, v in bii.items() if len(k) == 1 and not math.isclose(v, 0.0)}
+
+
+def mobius_to_influence_ii(mobius, max_order=None):
+    """
+    Convert Mobius coefficients to Influence interaction indices.
+
+    This mirrors AttrLLM_dev/attribution/utils.py logic:
+      1) mobius -> fourier
+      2) for each target location A, influence(A) = sum_{B overlaps A} |hat{f}(B)|^2
+    """
+    fourier = mobius_to_fourier(mobius)
+    if not fourier:
+        return {}
+
+    squared = {loc: abs(coef) ** 2 for loc, coef in fourier.items()}
+    out = {}
+    for target_loc in fourier.keys():
+        if max_order is not None and len(target_loc) > max_order:
+            continue
+        target = set(target_loc)
+        score = 0.0
+        for loc, sq in squared.items():
+            if not target.isdisjoint(loc):
+                score += float(sq)
+        out[tuple(sorted(target_loc))] = score
+    return out
+
+
+def mobius_to_influence(mobius):
+    iii = mobius_to_influence_ii(mobius, max_order=1)
+    return {k: v for k, v in iii.items() if len(k) == 1 and not math.isclose(v, 0.0)}
+
+
+def influence_interactions(mobius, order=2, top_k=15, player_filter=None):
+    """
+    Extract top-k Influence interactions by order from mobius coefficients.
+    """
+    iii = mobius_to_influence_ii(mobius)
+    filtered = {}
+    for loc, val in iii.items():
+        if len(loc) != order or math.isclose(float(val), 0.0):
+            continue
+        if player_filter is not None and not player_filter(loc):
+            continue
+        filtered[tuple(sorted(loc))] = float(val)
+    top = sorted(filtered.items(), key=lambda kv: abs(kv[1]), reverse=True)
+    if top_k is None or top_k <= 0:
+        return top
+    return top[:top_k]
+
+
+def extract_interactions(mobius, order=2, top_k=15, player_filter=None):
+    """
+    Extract top-k interaction indices of the given order.
+    Works for any pair including cross-modal (image × text).
+
+    Args:
+        player_filter: optional callable(loc) -> bool.  When provided, only
+                       interactions for which player_filter(loc) is True are
+                       included.  This lets callers extract e.g. only
+                       cross-modal pairs with their own dedicated top_k.
+    """
+    scores = {}
+    for T, mval in mobius.items():
+        T = tuple(sorted(T))
+        k = len(T)
+        if k < order:
+            continue
+        weight = 1.0 / (k - order + 1)
+        for sub in combinations(T, order):
+            if player_filter is not None and not player_filter(sub):
+                continue
+            scores[sub] = scores.get(sub, 0.0) + weight * float(np.real(mval))
+    top = sorted(scores.items(), key=lambda kv: abs(kv[1]), reverse=True)
+    return top[:top_k] if top_k else top
+
+
+def extract_cross_per_token(mobius, n_img, n_tok, order=2, per_token_k=5):
+    """
+    Extract cross-modal interactions with **per-token top-k** guarantee.
+
+    Returns:
+        per_token_dict : dict[int, list]   – token local index j → [(loc, val), …] (up to per_token_k)
+        global_top5    : list              – overall top-5 cross-modal pairs by |value|
+    """
+    cross_filter = lambda loc: any(i < n_img for i in loc) and any(i >= n_img for i in loc)
+    # Get ALL cross-modal interactions (no top_k cap)
+    all_cross = extract_interactions(mobius, order=order, top_k=0, player_filter=cross_filter)
+
+    # Per-token buckets
+    per_token_dict: Dict[int, list] = {j: [] for j in range(n_tok)}
+    for loc, val in all_cross:
+        for idx in loc:
+            if idx >= n_img:  # this is a token player
+                tok_j = idx - n_img
+                if tok_j < n_tok:
+                    per_token_dict[tok_j].append((loc, val))
+
+    # Keep only top-k per token
+    for j in per_token_dict:
+        per_token_dict[j].sort(key=lambda x: abs(x[1]), reverse=True)
+        per_token_dict[j] = per_token_dict[j][:per_token_k]
+
+    # Global top-5
+    global_top5 = all_cross[:5]
+
+    return per_token_dict, global_top5
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# 9.  Dataset loader  (same as set.py)
+# ═══════════════════════════════════════════════════════════════════════════
+
+def load_coco_samples(cfg):
+    samples = _try_load_flickr30k(cfg)
+    if samples is not None:
+        return samples
+    samples = _try_load_coco_urls(cfg)
+    if samples is not None:
+        return samples
+    raise RuntimeError("Could not load any image dataset.")
+
+
+def _try_load_flickr30k(cfg):
+    if load_dataset is None:
+        return None
+    try:
+        print("[Dataset] Trying nlphuji/flickr30k …")
+        ds = load_dataset("nlphuji/flickr30k", split="test")
+        total = len(ds)
+        rng = np.random.default_rng(cfg.random_seed)
+        indices = rng.choice(total, size=min(cfg.num_images, total), replace=False)
+        samples = []
+        for idx in indices:
+            row = ds[int(idx)]
+            pil_img = row["image"].convert("RGB")
+            cap_field = row.get("caption", row.get("captions", []))
+            caption = cap_field[0] if isinstance(cap_field, list) and cap_field else (cap_field if isinstance(cap_field, str) else "a photo")
+            cocoid = row.get("img_id", row.get("id", int(idx)))
+            if isinstance(cocoid, str):
+                cocoid = hash(cocoid) % 10**8
+            samples.append({"image": pil_img, "caption": caption, "cocoid": int(cocoid)})
+            print(f"  loaded id={cocoid}  caption='{caption[:60]}…'")
+        return samples
+    except Exception as e:
+        print(f"[Dataset] flickr30k failed: {e}")
+        return None
+
+
+def _try_load_coco_urls(cfg):
+    import urllib.request, io, json, zipfile
+    print("[Dataset] Trying direct COCO val2017 download (with real annotations) …")
+    base_img_url = "http://images.cocodataset.org/val2017/"
+    ann_zip_url = "http://images.cocodataset.org/annotations/annotations_trainval2017.zip"
+    cache_dir = os.environ.get("HF_HOME", os.path.join(os.path.expanduser("~"), ".cache", "huggingface"))
+    captions_cache = os.path.join(cache_dir, "coco_captions_val2017.json")
+    id_to_captions = {}
+    if os.path.isfile(captions_cache):
+        with open(captions_cache) as f:
+            ann_data = json.load(f)
+        for ann in ann_data["annotations"]:
+            id_to_captions.setdefault(ann["image_id"], []).append(ann["caption"])
+    else:
+        print("  Downloading COCO annotations zip (~252 MB, one-time) …")
+        try:
+            with urllib.request.urlopen(ann_zip_url, timeout=300) as resp:
+                zip_bytes = resp.read()
+            with zipfile.ZipFile(io.BytesIO(zip_bytes)) as zf:
+                with zf.open("annotations/captions_val2017.json") as cf:
+                    ann_data = json.load(cf)
+            os.makedirs(os.path.dirname(captions_cache), exist_ok=True)
+            with open(captions_cache, "w") as f:
+                json.dump(ann_data, f)
+            for ann in ann_data["annotations"]:
+                id_to_captions.setdefault(ann["image_id"], []).append(ann["caption"])
+        except Exception as e:
+            print(f"  WARN: Failed to download COCO annotations: {e}")
+            return None
+    if not id_to_captions:
+        return None
+    all_ids = list(id_to_captions.keys())
+    rng = np.random.default_rng(cfg.random_seed)
+    chosen = rng.choice(all_ids, size=min(cfg.num_images, len(all_ids)), replace=False)
+    samples = []
+    for cid in chosen:
+        cid = int(cid)
+        url = base_img_url + f"{cid:012d}.jpg"
+        try:
+            with urllib.request.urlopen(url, timeout=30) as resp:
+                pil_img = Image.open(io.BytesIO(resp.read())).convert("RGB")
+            samples.append({"image": pil_img, "caption": id_to_captions[cid][0], "cocoid": cid})
+        except Exception as e:
+            print(f"  WARN: cocoid={cid}: {e}")
+    return samples or None
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# 10.  Visualisation helpers
+# ═══════════════════════════════════════════════════════════════════════════
+
+def render_segmentation_map(image, regions):
+    """Colour-coded segmentation map with labels."""
+    from PIL import ImageDraw, ImageFont
+    import colorsys
+
+    image = _ensure_rgb(image)
+    w, h = image.size
+    base = image.copy().convert("RGBA")
+    n = len(regions)
+    colours = []
+    for i in range(n):
+        hue = (i * 0.618033988749895) % 1.0
+        r, g, b = colorsys.hsv_to_rgb(hue, 0.65, 0.90)
+        colours.append((int(r*255), int(g*255), int(b*255)))
+
+    overlay = Image.new("RGBA", (w, h), (0, 0, 0, 0))
+    oa = np.array(overlay)
+    for i, reg in enumerate(regions):
+        mask = _region_mask(reg, (h, w))
+        if not mask.any(): continue
+        oa[mask, :3] = colours[i]
+        oa[mask, 3] = 100
+    combined = Image.alpha_composite(base, Image.fromarray(oa, "RGBA"))
+
+    # Edges
+    ea = np.zeros((h, w, 4), dtype=np.uint8)
+    for i, reg in enumerate(regions):
+        mask = _region_mask(reg, (h, w))
+        if not mask.any(): continue
+        try:
+            from scipy import ndimage
+            dilated = ndimage.binary_dilation(mask, iterations=2)
+        except ImportError:
+            dilated = mask.copy()
+            for sa, sd in [(0,1),(0,-1),(1,1),(1,-1)]:
+                dilated |= np.roll(mask, sd, axis=sa)
+                dilated |= np.roll(mask, sd*2, axis=sa)
+        edge = dilated & ~mask
+        ea[edge, :3] = [max(0,c-60) for c in colours[i]]
+        ea[edge, 3] = 200
+    combined = Image.alpha_composite(combined, Image.fromarray(ea, "RGBA"))
+
+    # Labels
+    draw = ImageDraw.Draw(combined)
+    try:
+        font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", size=max(12, min(h,w)//40))
+    except (IOError, OSError):
+        try:
+            font = ImageFont.truetype("/usr/share/fonts/dejavu-sans-fonts/DejaVuSans-Bold.ttf", size=max(12, min(h,w)//40))
+        except (IOError, OSError):
+            font = ImageFont.load_default()
+    for i, reg in enumerate(regions):
+        mask = _region_mask(reg, (h, w))
+        if not mask.any(): continue
+        ys, xs = np.where(mask)
+        cx, cy = int(xs.mean()), int(ys.mean())
+        for dx in (-1,0,1):
+            for dy in (-1,0,1):
+                if dx==0 and dy==0: continue
+                draw.text((cx+dx,cy+dy), reg.label, fill=(0,0,0,220), font=font, anchor="mm")
+        draw.text((cx,cy), reg.label, fill=(255,255,255,255), font=font, anchor="mm")
+    return combined
+
+
+def render_overlay(image, regions, values, alpha=0.6):
+    """Red/green overlay based on attribution values."""
+    image = _ensure_rgb(image)
+    w, h = image.size
+    ov = np.zeros((h, w, 4), dtype=np.uint8)
+    max_abs = max((abs(v) for v in values), default=0.0) or 1.0
+    for reg, val in zip(regions, values):
+        mask = _region_mask(reg, (h, w))
+        if not mask.any(): continue
+        norm = min(1.0, abs(val) / max_abs)
+        if norm <= 0: continue
+        rgb = np.array([1,109,1], dtype=np.uint8) if val >= 0 else np.array([221,19,19], dtype=np.uint8)
+        ov[mask, :3] = rgb
+        ov[mask, 3] = int(round(255 * alpha * norm))
+    return Image.fromarray(ov, "RGBA")
+
+
+def save_results_mm(
+    image, regions, token_players, all_labels,
+    values, interactions_img, interactions_tok,
+    cross_per_token, cross_global_top5,
+    cross_pair_scores, mobius,
+    method, cocoid, caption,
+    n_img, n_tok, out_dir="results_mm",
+):
+    """Save cross-modal attribution results."""
+    os.makedirs(out_dir, exist_ok=True)
+
+    # Image overlay (only image-region values)
+    img_vals = [float(values.get((i,), 0.0)) for i in range(n_img)]
+    overlay = render_overlay(image, regions, img_vals)
+    base = image.convert("RGBA")
+    Image.alpha_composite(base, overlay).save(os.path.join(out_dir, f"coco_{cocoid}_overlay.png"))
+    image.save(os.path.join(out_dir, f"coco_{cocoid}_original.png"))
+
+    # Segmentation map
+    try:
+        render_segmentation_map(image, regions).save(os.path.join(out_dir, f"coco_{cocoid}_segmap.png"))
+    except Exception as e:
+        print(f"  WARN: segmap: {e}")
+
+    # Text summary
+    lines = [f"=== COCO {cocoid} – {method} CROSS-MODAL attributions ===\n"]
+    if caption:
+        lines.append(f"Caption: {caption}\n")
+
+    # Image region values
+    lines.append(f"Image region values ({n_img} regions):")
+    sorted_img = sorted([(i, img_vals[i]) for i in range(n_img)], key=lambda x: abs(x[1]), reverse=True)
+    for idx, val in sorted_img:
+        lines.append(f"  {all_labels[idx]:>25s}  {val:+.6f}")
+
+    # Token values
+    tok_vals = [float(values.get((n_img + j,), 0.0)) for j in range(n_tok)]
+    lines.append(f"\nCaption token values ({n_tok} tokens):")
+    sorted_tok = sorted([(j, tok_vals[j]) for j in range(n_tok)], key=lambda x: abs(x[1]), reverse=True)
+    for j, val in sorted_tok:
+        lines.append(f"  {all_labels[n_img + j]:>25s}  {val:+.6f}")
+
+    # Each category already extracted with its own top_k
+    lines.append(f"\nImage×Image interactions ({len(interactions_img)}):")
+    for loc, val in interactions_img:
+        names = " × ".join(all_labels[i] for i in loc)
+        lines.append(f"  {names:>45s}  {val:+.6f}")
+
+    lines.append(f"\nToken×Token interactions ({len(interactions_tok)}):")
+    for loc, val in interactions_tok:
+        names = " × ".join(all_labels[i] for i in loc)
+        lines.append(f"  {names:>45s}  {val:+.6f}")
+
+    # ── Global top-5 cross-modal interactions ──
+    lines.append(f"\nTop-5 Cross-Modal interactions (seg × tok):")
+    for loc, val in cross_global_top5:
+        names = " × ".join(all_labels[i] for i in loc)
+        lines.append(f"  {names:>45s}  {val:+.6f}")
+
+    # ── Focused cross-modal: top-2 positive tokens + bottom-2 negative tokens × top-3 segs ──
+    pos_toks = [(j, tok_vals[j]) for j in range(n_tok) if tok_vals[j] > 0]
+    neg_toks = [(j, tok_vals[j]) for j in range(n_tok) if tok_vals[j] < 0]
+    pos_toks.sort(key=lambda x: x[1], reverse=True)
+    neg_toks.sort(key=lambda x: x[1])  # most negative first
+
+    for rank, (pj, pv) in enumerate(pos_toks[:2], 1):
+        plabel = all_labels[n_img + pj]
+        lines.append(f"\n★ Positive token #{rank}: {plabel}  ({pv:+.6f})  × top-3 segments:")
+        for loc, val in cross_per_token.get(pj, [])[:3]:
+            names = " × ".join(all_labels[i] for i in loc)
+            lines.append(f"  {names:>45s}  {val:+.6f}")
+    if not pos_toks:
+        lines.append(f"\n★ Positive tokens: (none — no positive token values)")
+
+    for rank, (nj, nv) in enumerate(neg_toks[:2], 1):
+        nlabel = all_labels[n_img + nj]
+        lines.append(f"\n★ Negative token #{rank}: {nlabel}  ({nv:+.6f})  × top-3 segments:")
+        for loc, val in cross_per_token.get(nj, [])[:3]:
+            names = " × ".join(all_labels[i] for i in loc)
+            lines.append(f"  {names:>45s}  {val:+.6f}")
+
+    # ── Per-token top-5 cross-modal (full dump) ──
+    lines.append(f"\n{'='*60}")
+    lines.append(f"Per-token cross-modal interactions (top-5 each):")
+    for j in range(n_tok):
+        tlabel = all_labels[n_img + j]
+        entries = cross_per_token.get(j, [])
+        if not entries:
+            lines.append(f"\n  {tlabel}: (none)")
+            continue
+        lines.append(f"\n  {tlabel}  (Shapley={tok_vals[j]:+.6f}):")
+        for loc, val in entries:
+            names = " × ".join(all_labels[i] for i in loc)
+            lines.append(f"    {names:>45s}  {val:+.6f}")
+
+    # ── Full seg×tok influence table (all pairs) ──
+    lines.append(f"\n{'='*60}")
+    lines.append("All seg×tok influence scores (complete matrix):")
+    for i in range(n_img):
+        for j in range(n_tok):
+            loc = (i, n_img + j)
+            names = f"{all_labels[i]} × {all_labels[n_img + j]}"
+            val = float(cross_pair_scores.get(loc, 0.0))
+            lines.append(f"  {names:>45s}  {val:+.6f}")
+
+    # ── Mobius coefficients (all terms) ──
+    lines.append(f"\n{'='*60}")
+    lines.append("Mobius scores (all terms, sorted by |score|):")
+    sorted_mobius = sorted(
+        ((tuple(sorted(loc)), float(val)) for loc, val in mobius.items()),
+        key=lambda kv: abs(kv[1]),
+        reverse=True,
+    )
+    for loc, val in sorted_mobius:
+        if len(loc) == 0:
+            names = "{}"
+        else:
+            names = " × ".join(all_labels[i] for i in loc)
+        lines.append(f"  {str(loc):>20s}  {names:>45s}  {val:+.6f}")
+
+    summary_path = os.path.join(out_dir, f"coco_{cocoid}_summary.txt")
+    with open(summary_path, "w") as f:
+        f.write("\n".join(lines))
+
+    # Also dump machine-readable per-pair influence scores.
+    csv_path = os.path.join(out_dir, f"coco_{cocoid}_cross_pair_influence.csv")
+    with open(csv_path, "w") as f:
+        f.write("seg_index,seg_label,token_index,token_label,player_pair,influence_score\n")
+        for i in range(n_img):
+            for j in range(n_tok):
+                loc = (i, n_img + j)
+                pair = f"({loc[0]},{loc[1]})"
+                f.write(
+                    f"{i},{all_labels[i]},{j},{all_labels[n_img + j]},\"{pair}\",{float(cross_pair_scores.get(loc, 0.0)):.10f}\n"
+                )
+
+    print(f"  Saved overlay → {out_dir}/coco_{cocoid}_overlay.png")
+    print(f"  Saved summary → {summary_path}")
+    print(f"  Saved pair influence CSV → {csv_path}")
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# Geometry helpers
+# ═══════════════════════════════════════════════════════════════════════════
+
+def _ensure_rgb(img):
+    return img if img.mode == "RGB" else img.convert("RGB")
+
+def _mask_to_bbox(mask):
+    ys, xs = np.where(mask)
+    if ys.size == 0:
+        return (0, 0, 0, 0)
+    return (int(xs.min()), int(ys.min()), int(xs.max())+1, int(ys.max())+1)
+
+def _region_mask(region, shape):
+    if region.mask is not None:
+        return region.mask
+    h, w = shape
+    x0, y0, x1, y1 = region.bbox
+    m = np.zeros((h, w), dtype=bool)
+    x0, x1 = max(0, min(w, x0)), max(0, min(w, x1))
+    y0, y1 = max(0, min(h, y0)), max(0, min(h, y1))
+    if x1 > x0 and y1 > y0:
+        m[y0:y1, x0:x1] = True
+    return m
+
+def _linspace_edges(size, bins):
+    bins = max(1, int(bins))
+    return [int(round(v)) for v in np.linspace(0, size, bins + 1)]
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# 11.  Main  –  cross-modal evaluation pipeline
+# ═══════════════════════════════════════════════════════════════════════════
+
+def evaluate_one(
+    image: Image.Image,
+    caption: str,
+    cocoid: int,
+    cfg: PipelineConfig,
+    scorer: CrossModalCLIPScorer,
+    sam_pipe=None,
+    out_dir: str = "results_mm",
+) -> Dict[str, Any]:
+    """Run the cross-modal pipeline on a single image."""
+    print(f"\n══ Image cocoid={cocoid} (cross-modal) ══")
+    timing_recorder = (
+        MaskTimingRecorder(cocoid=cocoid, mask_style=cfg.mask_style, out_dir=out_dir)
+        if cfg.record_mask_timing
+        else None
+    )
+
+    # Step 1: image regions
+    regions = featurise(image, cfg, sam_pipe=sam_pipe)
+    n_img = len(regions)
+
+    # Step 2: caption tokens
+    token_players, full_token_ids = tokenise_caption(caption, scorer.processor, cfg, offset=n_img)
+    n_tok = len(token_players)
+    n_total = n_img + n_tok
+
+    # All player labels
+    all_labels = [r.label for r in regions] + [tp.label for tp in token_players]
+    print(f"  Total players: {n_total} = {n_img} image + {n_tok} text")
+
+    # Build cross-modal f(S)
+    sf = build_cross_modal_set_function(
+        image, regions, token_players, full_token_ids, scorer, cfg,
+        timing_recorder=timing_recorder,
+    )
+
+    # Sanity check
+    f_full = sf(tuple([1] * n_total))
+    f_empty = sf(tuple([0] * n_total))
+    f_img_only = sf(tuple([1]*n_img + [0]*n_tok))
+    f_txt_only = sf(tuple([0]*n_img + [1]*n_tok))
+    print(f"  f(all)       = {f_full:.4f}")
+    print(f"  f(∅)         = {f_empty:.4f}")
+    print(f"  f(img only)  = {f_img_only:.4f}")
+    print(f"  f(txt only)  = {f_txt_only:.4f}")
+
+    # ── Save intermediate masked images for visual comparison ──
+    masked_dir = os.path.join(out_dir, f"coco_{cocoid}_masked_{cfg.mask_style}")
+    os.makedirs(masked_dir, exist_ok=True)
+
+    # 1. Fully masked image (all regions removed)
+    all_masked = apply_image_mask(image, regions, [0]*n_img,
+                                  style=cfg.mask_style, blur_radius=cfg.blur_radius, cfg=cfg,
+                                  timing_meta={
+                                      "recorder": timing_recorder,
+                                      "phase": "viz_mask",
+                                      "event": "all_masked",
+                                      "store_detail": True,
+                                  })
+    all_masked.save(os.path.join(masked_dir, "all_masked.png"))
+
+    # 2. Each region kept alone (solo), others masked
+    for i, reg in enumerate(regions):
+        solo = [0]*n_img
+        solo[i] = 1
+        img_solo = apply_image_mask(image, regions, solo,
+                                    style=cfg.mask_style, blur_radius=cfg.blur_radius, cfg=cfg,
+                                    timing_meta={
+                                        "recorder": timing_recorder,
+                                        "phase": "viz_mask",
+                                        "event": f"{reg.label}_solo",
+                                        "store_detail": True,
+                                    })
+        img_solo.save(os.path.join(masked_dir, f"{reg.label}_solo.png"))
+
+    # 3. Each region removed (others kept)
+    for i, reg in enumerate(regions):
+        removed = [1]*n_img
+        removed[i] = 0
+        img_removed = apply_image_mask(image, regions, removed,
+                                       style=cfg.mask_style, blur_radius=cfg.blur_radius, cfg=cfg,
+                                       timing_meta={
+                                           "recorder": timing_recorder,
+                                           "phase": "viz_mask",
+                                           "event": f"{reg.label}_removed",
+                                           "store_detail": True,
+                                       })
+        img_removed.save(os.path.join(masked_dir, f"{reg.label}_removed.png"))
+
+    print(f"  Saved {1 + 2*n_img} masked images → {masked_dir}/")
+
+    # ProxySPEX
+    mobius = run_proxyspex(sf, all_labels, max_order=cfg.max_order, seed=cfg.random_seed)
+
+    # Shapley / Banzhaf / Influence singleton values
+    if cfg.method == "shapley":
+        values = mobius_to_shapley(mobius)
+    elif cfg.method == "banzhaf":
+        values = mobius_to_banzhaf(mobius)
+    else:
+        values = mobius_to_influence(mobius)
+
+    # Extract interactions by category, each with its own top_k ──────────
+    # This avoids the old bug where a single global top_k=15 was shared
+    # across img×img, tok×tok, and cross-modal, causing most cross-modal
+    # pairs to be silently dropped.
+    img_img_filter = lambda loc: all(i < n_img for i in loc)
+    tok_tok_filter = lambda loc: all(i >= n_img for i in loc)
+    cross_filter   = lambda loc: any(i < n_img for i in loc) and any(i >= n_img for i in loc)
+
+    if cfg.method == "influence":
+        interactions_img = influence_interactions(
+            mobius, order=cfg.max_order, top_k=cfg.top_k_interactions, player_filter=img_img_filter
+        )
+        interactions_tok = influence_interactions(
+            mobius, order=cfg.max_order, top_k=cfg.top_k_interactions, player_filter=tok_tok_filter
+        )
+    else:
+        interactions_img = extract_interactions(
+            mobius, order=cfg.max_order, top_k=cfg.top_k_interactions, player_filter=img_img_filter
+        )
+        interactions_tok = extract_interactions(
+            mobius, order=cfg.max_order, top_k=cfg.top_k_interactions, player_filter=tok_tok_filter
+        )
+    # Cross-modal: per-token top-5 (so no token is silently dropped)
+    # Always compute complete pairwise cross-modal influence scores for export.
+    all_cross_influence = influence_interactions(mobius, order=2, top_k=0, player_filter=cross_filter)
+    all_cross_influence_dict = {tuple(sorted(loc)): float(val) for loc, val in all_cross_influence}
+
+    if cfg.method == "influence":
+        all_cross = all_cross_influence
+
+        cross_per_token = {j: [] for j in range(n_tok)}
+        for loc, val in all_cross:
+            for idx in loc:
+                if idx >= n_img:
+                    tok_j = idx - n_img
+                    if tok_j < n_tok:
+                        cross_per_token[tok_j].append((loc, val))
+        for j in cross_per_token:
+            cross_per_token[j].sort(key=lambda x: abs(x[1]), reverse=True)
+            cross_per_token[j] = cross_per_token[j][:5]
+        cross_global_top5 = all_cross[:5]
+    else:
+        cross_per_token, cross_global_top5 = extract_cross_per_token(
+            mobius, n_img, n_tok, order=cfg.max_order, per_token_k=5,
+        )
+
+    # Build a complete seg×tok influence matrix (every pair has a score; missing -> 0).
+    cross_pair_scores = {}
+    for i in range(n_img):
+        for j in range(n_tok):
+            pair = tuple(sorted((i, n_img + j)))
+            cross_pair_scores[(i, n_img + j)] = float(all_cross_influence_dict.get(pair, 0.0))
+
+    # Combined list (for backward compat in return value)
+    # Deduplicate the union of all per-token lists
+    _seen = set()
+    cross_interactions_all = []
+    for j in range(n_tok):
+        for item in cross_per_token[j]:
+            key = (item[0], item[1])
+            if item[0] not in _seen:
+                _seen.add(item[0])
+                cross_interactions_all.append(item)
+    interactions = interactions_img + interactions_tok + cross_interactions_all
+
+    # Print summary
+    print(f"\n  {cfg.method.upper()} values (image regions):")
+    for i in range(n_img):
+        v = values.get((i,), 0.0)
+        if not math.isclose(v, 0.0):
+            print(f"    {all_labels[i]:>20s}  {v:+.6f}")
+
+    print(f"\n  {cfg.method.upper()} values (caption tokens):")
+    for j in range(n_tok):
+        v = values.get((n_img + j,), 0.0)
+        if not math.isclose(v, 0.0):
+            print(f"    {all_labels[n_img + j]:>20s}  {v:+.6f}")
+
+    # ── Global top-5 cross-modal interactions ──
+    print(f"\n  Top-5 cross-modal interactions (seg × tok):")
+    for loc, v in cross_global_top5:
+        names = " × ".join(all_labels[i] for i in loc)
+        print(f"    {names:>35s}  {v:+.6f}")
+
+    # ── Focused cross-modal: top-2 pos / bottom-2 neg tokens × top-3 segments ──
+    tok_val_list = [(j, float(values.get((n_img + j,), 0.0))) for j in range(n_tok)]
+    pos_toks = sorted([(j, v) for j, v in tok_val_list if v > 0], key=lambda x: -x[1])
+    neg_toks = sorted([(j, v) for j, v in tok_val_list if v < 0], key=lambda x: x[1])
+
+    for rank, (pj, pv) in enumerate(pos_toks[:2], 1):
+        print(f"\n  ★ Positive token #{rank}: {all_labels[n_img+pj]} ({pv:+.6f}) × top-3 segs:")
+        for loc, v in cross_per_token[pj][:3]:
+            names = " × ".join(all_labels[i] for i in loc)
+            print(f"    {names:>35s}  {v:+.6f}")
+
+    for rank, (nj, nv) in enumerate(neg_toks[:2], 1):
+        print(f"\n  ★ Negative token #{rank}: {all_labels[n_img+nj]} ({nv:+.6f}) × top-3 segs:")
+        for loc, v in cross_per_token[nj][:3]:
+            names = " × ".join(all_labels[i] for i in loc)
+            print(f"    {names:>35s}  {v:+.6f}")
+
+    # Save
+    save_results_mm(
+        image, regions, token_players, all_labels,
+        values, interactions_img, interactions_tok,
+        cross_per_token, cross_global_top5,
+        cross_pair_scores, mobius,
+        cfg.method, cocoid, caption,
+        n_img, n_tok, out_dir=out_dir,
+    )
+
+    timing_detail_path = None
+    timing_summary_path = None
+    if timing_recorder is not None:
+        timing_detail_path, timing_summary_path = timing_recorder.write_csvs()
+        print(f"  Saved mask timing details → {timing_detail_path}")
+        print(f"  Saved mask timing summary → {timing_summary_path}")
+
+    return {
+        "cocoid": cocoid,
+        "n_img": n_img,
+        "n_tok": n_tok,
+        "all_labels": all_labels,
+        "mobius": mobius,
+        "values": values,
+        "interactions": interactions,
+        "mask_timing_detail_path": timing_detail_path,
+        "mask_timing_summary_path": timing_summary_path,
+    }
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Cross-modal image×text attribution via mask-and-CLIP with ProxySPEX"
+    )
+    parser.add_argument("--mode", choices=["unsam", "sam", "patch"], default="unsam")
+    parser.add_argument("--grid", type=int, default=8)
+    parser.add_argument("--segment_level", choices=["all", "coarse", "medium", "fine"], default="all")
+    parser.add_argument("--mask_style", choices=["blur", "gray", "mean", "diffusion", "lama", "lcm", "flux", "flux_schnell", "diffusion_cv2", "cv2_telea", "cv2_ns"], default="blur")
+    parser.add_argument("--diffusion_model", default="runwayml/stable-diffusion-inpainting")
+    parser.add_argument("--diffusion_prompt", default="")
+    parser.add_argument("--diffusion_steps", type=int, default=20)
+    parser.add_argument("--diffusion_crop_pad", type=int, default=48)
+    parser.add_argument("--diffusion_crop_min_side", type=int, default=256)
+    parser.add_argument("--diffusion_local_max_cover", type=float, default=0.7)
+    parser.add_argument("--diffusion_local_crop", dest="diffusion_local_crop", action="store_true")
+    parser.add_argument("--no_diffusion_local_crop", dest="diffusion_local_crop", action="store_false")
+    parser.set_defaults(diffusion_local_crop=True)
+    parser.add_argument("--lcm_base_model", default="diffusers/stable-diffusion-xl-1.0-inpainting-0.1",
+                        help="Base SDXL inpainting model for LCM-LoRA")
+    parser.add_argument("--lcm_lora_id", default="latent-consistency/lcm-lora-sdxl",
+                        help="LCM-LoRA adapter id for SDXL")
+    parser.add_argument("--lcm_steps", type=int, default=4,
+                        help="Inference steps for LCM-LoRA (recommended 4)")
+    parser.add_argument("--lcm_guidance_scale", type=float, default=1.0,
+                        help="Guidance scale for LCM-LoRA")
+    parser.add_argument("--flux_model", default="black-forest-labs/FLUX.1-schnell",
+                        help="FLUX.1-Schnell model id")
+    parser.add_argument("--flux_steps", type=int, default=4,
+                        help="Inference steps for FLUX.1-Schnell")
+    parser.add_argument("--flux_guidance_scale", type=float, default=0.0,
+                        help="Guidance scale for FLUX.1-Schnell")
+    parser.add_argument("--flux_strength", type=float, default=0.85,
+                        help="Img2img strength for FLUX.1-Schnell")
+    parser.add_argument("--flux_max_sequence_length", type=int, default=256,
+                        help="Max sequence length for FLUX.1-Schnell")
+    parser.add_argument("--method", choices=["shapley", "banzhaf", "influence"], default="shapley")
+    parser.add_argument("--num_images", type=int, default=1)
+    parser.add_argument("--max_regions", type=int, default=16, help="Max image regions (default lower for cross-modal)")
+    parser.add_argument("--max_tokens", type=int, default=15, help="Max caption tokens to include as players")
+    parser.add_argument("--max_order", type=int, default=2)
+    parser.add_argument("--top_k", type=int, default=15)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--split", default="train")
+    parser.add_argument("--clip_model", default="openai/clip-vit-base-patch32")
+    parser.add_argument("--sam_model", default="yujunwei04/unsam-whole-image-segmentation")
+    parser.add_argument("--out_dir", default="results_mm")
+    parser.add_argument("--record_mask_timing", action="store_true",
+                        help="Write per-mask timing CSVs (total vs infill time)")
+    parser.add_argument("--print_fourier_mobius", action="store_true",
+                        help="Print Fourier/Mobius coefficient scores from ProxySPEX")
+    parser.add_argument("--print_coeff_top_k", type=int, default=20,
+                        help="How many Fourier/Mobius terms to print (<=0 prints all)")
+    args = parser.parse_args()
+
+    cfg = PipelineConfig(
+        coco_split=args.split,
+        num_images=args.num_images,
+        mode=args.mode,
+        sam_model_name=args.sam_model,
+        segment_level=args.segment_level,
+        grid_size=args.grid,
+        max_regions=args.max_regions,
+        mask_style=args.mask_style,
+        clip_model_name=args.clip_model,
+        method=args.method,
+        max_order=args.max_order,
+        top_k_interactions=args.top_k,
+        random_seed=args.seed,
+        diffusion_model_name=args.diffusion_model,
+        diffusion_prompt=args.diffusion_prompt,
+        diffusion_steps=args.diffusion_steps,
+        diffusion_local_crop=args.diffusion_local_crop,
+        diffusion_crop_pad=args.diffusion_crop_pad,
+        diffusion_crop_min_side=args.diffusion_crop_min_side,
+        diffusion_local_max_cover=args.diffusion_local_max_cover,
+        lcm_base_model_name=args.lcm_base_model,
+        lcm_lora_id=args.lcm_lora_id,
+        lcm_steps=args.lcm_steps,
+        lcm_guidance_scale=args.lcm_guidance_scale,
+        flux_model_name=args.flux_model,
+        flux_steps=args.flux_steps,
+        flux_guidance_scale=args.flux_guidance_scale,
+        flux_strength=args.flux_strength,
+        flux_max_sequence_length=args.flux_max_sequence_length,
+        max_tokens=args.max_tokens,
+        record_mask_timing=args.record_mask_timing,
+        print_fourier_mobius=args.print_fourier_mobius,
+        print_coeff_top_k=args.print_coeff_top_k,
+    )
+
+    global _PRINT_FOURIER_MOBIUS, _PRINT_COEF_TOP_K
+    _PRINT_FOURIER_MOBIUS = bool(cfg.print_fourier_mobius)
+    _PRINT_COEF_TOP_K = int(cfg.print_coeff_top_k)
+
+    print("=" * 65)
+    print("  Cross-Modal Image×Text Attribution Pipeline")
+    print("  Players: image regions + caption tokens")
+    print("  f(S) = CLIP(masked_image, masked_caption)")
+    print("=" * 65)
+    print(f"  Mode          : {cfg.mode}")
+    print(f"  Mask style    : {cfg.mask_style}")
+    print(f"  Method        : {cfg.method}")
+    print(f"  CLIP model    : {cfg.clip_model_name}")
+    if cfg.mode in ("sam", "unsam"):
+        print(f"  Seg model     : {cfg.sam_model_name}")
+        print(f"  Seg level     : {cfg.segment_level}")
+    else:
+        print(f"  Grid size     : {cfg.grid_size}")
+    print(f"  Max regions   : {cfg.max_regions}")
+    print(f"  Max tokens    : {cfg.max_tokens}")
+    print(f"  Max order     : {cfg.max_order}")
+    print(f"  Mask timing   : {cfg.record_mask_timing}")
+    print(f"  Print F/M     : {cfg.print_fourier_mobius}")
+    print(f"  Coef top-k    : {cfg.print_coeff_top_k}")
+    print(f"  Device        : {cfg.device}")
+    print()
+
+    samples = load_coco_samples(cfg)
+
+    print("\n[Models] Loading CLIP (cross-modal scorer)…")
+    scorer = CrossModalCLIPScorer(cfg)
+
+    sam_pipe = None
+    if cfg.mode in ("sam", "unsam"):
+        print("[Models] Loading segmentation model…")
+        sam_pipe = _load_sam_pipeline(cfg)
+
+    all_results = []
+    for sample in samples:
+        result = evaluate_one(
+            image=sample["image"],
+            caption=sample["caption"],
+            cocoid=sample["cocoid"],
+            cfg=cfg,
+            scorer=scorer,
+            sam_pipe=sam_pipe,
+            out_dir=args.out_dir,
+        )
+        all_results.append(result)
+
+    print("\n" + "=" * 65)
+    print(f"  Done. Processed {len(all_results)} image(s).")
+    print(f"  Results saved to ./{args.out_dir}/")
+    print("=" * 65)
+
+
+if __name__ == "__main__":
+    main()

attribution/utils.py

@@ -0,0 +1,438 @@
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from itertools import combinations
+from typing import Dict, List, Sequence, Tuple, Union
+
+import numpy as np
+
+
+@dataclass
+class AttributionResult:
+    """Container for a single feature or interaction attribution value."""
+
+    feature: Union[str, Tuple[str, ...]]
+    value: float
+    method: str
+    interaction_order: int
+
+
+def powerset(loc_tuple: Tuple[int, ...], max_order: int | None = None) -> List[Tuple[int, ...]]:
+    """
+    Generate the powerset of a sparse location tuple up to a specified maximum order.
+
+    Parameters:
+    - loc_tuple: A sparse tuple of feature indices (e.g., (0, 2) means features 0 and 2).
+    - max_order: The maximum order of the powerset (default is None).
+
+    Returns:
+    - A list of tuples representing the powerset in sparse format.
+    """
+    # loc_tuple is already sparse (contains feature indices)
+    # e.g., (0, 2) means features 0 and 2 are active
+    
+    if max_order is None:
+        max_order = len(loc_tuple)
+    
+    # Generate all subsets of the sparse indices up to max_order
+    tuples = []
+    for r in range(max_order + 1):
+        for subset in combinations(loc_tuple, r):
+            tuples.append(subset)
+    
+    return tuples
+
+
+def fourier_to_mobius(fourier_dict: Dict[Tuple[int, ...], float]) -> Dict[Tuple[int, ...], float]:
+    """
+    Convert Fourier coefficients to Mobius coefficients.
+
+    Parameters:
+    - fourier_dict: A dictionary of Fourier coefficients.
+
+    Returns:
+    - A dictionary of Mobius coefficients.
+    """
+    if len(fourier_dict) == 0:
+        return {}
+    else:
+        unscaled_mobius_dict = {}
+        for loc, coef in fourier_dict.items():
+            real_coef = np.real(coef)
+            for subset in powerset(loc):
+                if subset in unscaled_mobius_dict:
+                    unscaled_mobius_dict[subset] += real_coef
+                else:
+                    unscaled_mobius_dict[subset] = real_coef
+
+        # multiply each entry by (-2)^(cardinality)
+        return {loc: val * np.power(-2.0, len(loc)) for loc, val in unscaled_mobius_dict.items() if
+                np.abs(val) > 1e-12}
+
+def mobius_to_fourier(mobius_dict: Dict[Tuple[int, ...], float]) -> Dict[Tuple[int, ...], float]:
+    """
+    Convert Mobius coefficients to Fourier coefficients.
+
+    Parameters:
+    - mobius_dict: A dictionary of Mobius coefficients.
+
+    Returns:
+    - A dictionary of Fourier coefficients.
+    """
+    if len(mobius_dict) == 0:
+        return {}
+    else:
+        unscaled_fourier_dict = {}
+        for loc, coef in mobius_dict.items():
+            # Sparse tuples store feature indices, so cardinality is len(loc), not sum(loc).
+            real_coef = np.real(coef) / (2 ** len(loc))
+            for subset in powerset(loc):
+                if subset in unscaled_fourier_dict:
+                    unscaled_fourier_dict[subset] += real_coef
+                else:
+                    unscaled_fourier_dict[subset] = real_coef
+
+        # multiply each entry by (-1)^(cardinality)
+        return {loc: val * np.power(-1.0, len(loc)) for loc, val in unscaled_fourier_dict.items() if
+                np.abs(val) > 1e-12}
+
+
+def mobius_to_shapley_ii(
+    mobius_dict: Dict[Tuple[int, ...], complex],
+    max_order: int | None = None,
+    **kwargs,
+) -> Dict[Tuple[int, ...], float]:
+    """
+    Convert Mobius coefficients to Shapley interaction indices.
+    
+    WARNING: This function can be expensive (exponential in feature count).
+    For large feature sets, consider using sparse pairwise extraction instead.
+    
+    Args:
+        mobius_dict: Mobius coefficients keyed by sparse tuples of indices.
+        max_order: Maximum interaction order to compute (limits powerset size)
+    """
+    sii_dict: Dict[Tuple[int, ...], float] = {}
+    
+    for loc, coef in mobius_dict.items():
+        loc = tuple(sorted(loc))
+        real_coef = float(np.real(coef))
+        t_size = len(loc)
+        for subset in powerset(loc, max_order=max_order):
+            contribution = real_coef / (1 + t_size - len(subset))
+            sii_dict[subset] = sii_dict.get(subset, 0.0) + contribution
+    return sii_dict
+
+
+def mobius_to_banzhaf_ii(
+    mobius_dict: Dict[Tuple[int, ...], complex],
+    max_order: int | None = None,
+    **kwargs,
+) -> Dict[Tuple[int, ...], float]:
+    """
+    Convert Mobius coefficients to Banzhaf interaction indices.
+    
+    WARNING: This function can be expensive (exponential in feature count).
+    For large feature sets, consider using sparse pairwise extraction instead.
+    
+    Args:
+        mobius_dict: Mobius coefficients keyed by sparse tuples of indices.
+        max_order: Maximum interaction order to compute (limits powerset size)
+    """
+    bii_dict: Dict[Tuple[int, ...], float] = {}
+    
+    for loc, coef in mobius_dict.items():
+        loc = tuple(sorted(loc))
+        real_coef = float(np.real(coef))
+        t_size = len(loc)
+        for subset in powerset(loc, max_order=max_order):
+            contribution = real_coef / math.pow(2.0, t_size - len(subset))
+            bii_dict[subset] = bii_dict.get(subset, 0.0) + contribution
+    
+    return bii_dict
+
+def mobius_to_influence_ii(
+    mobius_dict: Dict[Tuple[int, ...], complex],
+    max_order: int | None = None,
+    **kwargs,
+) -> Dict[Tuple[int, ...], float]:
+    """
+    Convert Mobius coefficients to Influence interaction indices.
+    
+    WARNING: This function can be expensive (exponential in feature count).
+    For large feature sets, consider using sparse pairwise extraction instead.
+    
+    Args:
+        mobius_dict: Mobius coefficients keyed by sparse tuples of indices.
+        max_order: Maximum interaction order to compute (limits powerset size)
+    """
+    subset_influences: Dict[Tuple[int, ...], float] = {}
+
+    fourier_dict = mobius_to_fourier(mobius_dict)
+    
+    # Step 1: Precompute the squared magnitude for every location
+    # (Doing this once saves us from recalculating abs()**2 inside the loop)
+    squared_coefs = {loc: abs(coef) ** 2 for loc, coef in fourier_dict.items()}
+    
+    # Step 2: For every nonzero location, calculate its subset influence
+    for target_loc in fourier_dict.keys():
+        target_set = set(target_loc)
+        influence = 0.0
+        
+        # Sum the squared coefficients of any location that overlaps with the target
+        for loc, sq_coef in squared_coefs.items():
+            if not target_set.isdisjoint(loc):
+                influence += sq_coef
+                
+        subset_influences[target_loc] = influence
+        
+    return subset_influences
+
+
+def _filter_order(values: Dict[Tuple[int, ...], float], order: int) -> Dict[Tuple[int, ...], float]:
+    return {
+        loc: val
+        for loc, val in values.items()
+        if len(loc) == order and not math.isclose(val, 0.0)
+    }
+
+
+def _top_k(items: Dict[Tuple[int, ...], float], top_k: int) -> List[Tuple[Tuple[int, ...], float]]:
+    sorted_items = sorted(items.items(), key=lambda kv: abs(kv[1]), reverse=True)
+    if top_k is None or top_k <= 0:
+        return sorted_items
+    return sorted_items[:top_k]
+
+
+def mobius_to_shapley(mobius_dict: Dict[Tuple[int, ...], complex]) -> Dict[Tuple[int, ...], float]:
+    """
+    Convert Mobius coefficients to singleton Shapley values.
+    
+    Optimized: Only computes up to order=1 (singletons) to avoid exponential blow-up.
+    """
+    # Only compute singletons - no need for higher-order powersets
+    sii = mobius_to_shapley_ii(mobius_dict, max_order=1)
+    return _filter_order(sii, order=1)
+
+
+def shapley_interactions(
+    mobius_dict: Dict[Tuple[int, ...], complex],
+    order: int = 2,
+    top_k: int = 10,
+) -> List[Tuple[Tuple[int, ...], float]]:
+    """
+    Extract the top-k Shapley interaction indices of a given order.
+
+    For order=2 (pairwise), uses sparse superset aggregation to avoid exponential blow-up.
+    For order>2, falls back to full conversion (may be expensive).
+    
+    Formula for pairwise (|S|=2): φ_shapley(S) = Σ_{T ⊇ S} m(T) / (|T| - |S| + 1)
+    """
+    if order == 2:
+        # SPARSE PATH: Sparse superset aggregation for pairwise interactions
+        # Formula: φ_shapley(S) = Σ_{T ⊇ S} m(T) / (|T| - |S| + 1) = Σ_{T ⊇ S} m(T) / (|T| - 1)
+        assert order == 2, "Only pairwise supported for sparse path"
+        
+        if not mobius_dict:
+            return []
+        
+        scores: Dict[Tuple[int, ...], float] = {}
+        
+        for T, mval in mobius_dict.items():
+            T = tuple(sorted(T))
+            k = len(T)
+            if k < order:
+                continue
+            # Weight: 1 / (|T| - |S| + 1) = 1 / (k - 2 + 1) = 1 / (k - 1)
+            weight = 1.0 / (k - order + 1)
+            for subset in combinations(T, order):
+                scores[subset] = scores.get(subset, 0.0) + weight * float(np.real(mval))
+        
+        return _top_k(scores, top_k)
+    else:
+        # FULL PATH: Use traditional conversion for higher orders
+        sii = mobius_to_shapley_ii(mobius_dict)
+        filtered = _filter_order(sii, order=order)
+        return _top_k(filtered, top_k)
+
+
+def mobius_to_banzhaf(mobius_dict: Dict[Tuple[int, ...], complex]) -> Dict[Tuple[int, ...], float]:
+    """
+    Convert Mobius coefficients to singleton Banzhaf values.
+    
+    Optimized: Only computes up to order=1 (singletons) to avoid exponential blow-up.
+    """
+    # Only compute singletons - no need for higher-order powersets
+    bii = mobius_to_banzhaf_ii(mobius_dict, max_order=1)
+    return _filter_order(bii, order=1)
+
+
+def banzhaf_interactions(
+    mobius_dict: Dict[Tuple[int, ...], complex],
+    order: int = 2,
+    top_k: int = 10,
+) -> List[Tuple[Tuple[int, ...], float]]:
+    """
+    Extract the top-k Banzhaf interaction indices of a given order.
+    
+    For order=2 (pairwise), uses sparse superset aggregation to avoid exponential blow-up.
+    For order>2, falls back to full conversion (may be expensive).
+    
+    Formula for pairwise (|S|=2): β_banzhaf(S) = Σ_{T ⊇ S} m(T) / 2^(|T| - |S|)
+    """
+    if order == 2:
+        # SPARSE PATH: Sparse superset aggregation for pairwise interactions
+        # Formula: β_banzhaf(S) = Σ_{T ⊇ S} m(T) / 2^(|T| - |S|) = Σ_{T ⊇ S} m(T) / 2^(|T| - 2)
+        assert order == 2, "Only pairwise supported for sparse path"
+        
+        if not mobius_dict:
+            return []
+        
+        scores: Dict[Tuple[int, ...], float] = {}
+        
+        for T, mval in mobius_dict.items():
+            T = tuple(sorted(T))
+            k = len(T)
+            if k < order:
+                continue
+            # Weight: 1 / 2^(|T| - |S|) = 1 / 2^(k - 2)
+            weight = 1.0 / (2 ** (k - order))
+            for subset in combinations(T, order):
+                scores[subset] = scores.get(subset, 0.0) + weight * float(np.real(mval))
+        
+        return _top_k(scores, top_k)
+    else:
+        # FULL PATH: Use traditional conversion for higher orders
+        bii = mobius_to_banzhaf_ii(mobius_dict)
+        filtered = _filter_order(bii, order=order)
+        return _top_k(filtered, top_k)
+
+def mobius_to_banzhaf(mobius_dict: Dict[Tuple[int, ...], complex]) -> Dict[Tuple[int, ...], float]:
+    """
+    Convert Mobius coefficients to singleton Banzhaf values.
+    
+    Optimized: Only computes up to order=1 (singletons) to avoid exponential blow-up.
+    """
+    # Only compute singletons - no need for higher-order powersets
+    bii = mobius_to_banzhaf_ii(mobius_dict, max_order=1)
+    return _filter_order(bii, order=1)
+
+def mobius_to_influence(mobius_dict: Dict[Tuple[int, ...], complex]) -> Dict[Tuple[int, ...], float]:
+    """
+    Convert Mobius coefficients to singleton Influence values.
+    
+    Optimized: Only computes up to order=1 (singletons) to avoid exponential blow-up.
+    """
+    # Only compute singletons - no need for higher-order powersets
+    iii = mobius_to_influence_ii(mobius_dict, max_order=1)
+    return _filter_order(iii, order=1)
+
+def influence_interactions(
+    mobius_dict: Dict[Tuple[int, ...], complex],
+    order: int = 2,
+    top_k: int = 10,
+) -> List[Tuple[Tuple[int, ...], float]]:
+    """
+    Extract the top-k Influence interaction indices of a given order.
+    
+    For order=2 (pairwise), uses sparse superset aggregation to avoid exponential blow-up.
+    For order>2, falls back to full conversion (may be expensive).
+    
+    Formula for pairwise (|S|=2): β_banzhaf(S) = Σ_{T ⊇ S} m(T) / 2^(|T| - |S|)
+    """
+    # FULL PATH: Use traditional conversion for higher orders
+    iii = mobius_to_influence_ii(mobius_dict)
+    filtered = _filter_order(iii, order=order)
+    return _top_k(filtered, top_k)
+
+def convert_to_heatmap_data(
+    attrs: Sequence[AttributionResult],
+    feature_order: Sequence[str] | None = None,
+) -> Dict[str, Sequence]:
+    """
+    Convert pairwise attribution results into a token × token interaction matrix.
+
+    Parameters
+    ----------
+    attrs:
+        Sequence of attribution entries where ``feature`` is either a single token
+        (diagonal contribution) or a tuple of tokens describing an interaction.
+    feature_order:
+        Optional list specifying the axis ordering. When omitted, tokens are added
+        according to first appearance in ``attrs``.
+    """
+    if not attrs:
+        # Safe default when there are no interactions
+        return {
+            "matrix": [],
+            "features": list(feature_order) if feature_order else [],
+            "methods": [],
+            "orders": [],
+            "min": 0.0,
+            "max": 0.0,
+        }
+
+    if feature_order:
+        tokens = list(feature_order)
+    else:
+        tokens = []
+        for attr in attrs:
+            feats = attr.feature if isinstance(attr.feature, (tuple, list)) else (attr.feature,)
+            for token in feats:
+                if token not in tokens:
+                    tokens.append(token)
+
+    size = len(tokens)
+    if size == 0:
+        return {"matrix": [], "features": [], "methods": []}
+
+    index = {token: idx for idx, token in enumerate(tokens)}
+    matrix = [[0.0 for _ in range(size)] for _ in range(size)]
+    methods: List[str] = []
+    orders: List[int] = []
+
+    for attr in attrs:
+        feats = attr.feature if isinstance(attr.feature, (tuple, list)) else (attr.feature,)
+        if not feats:
+            continue
+        methods.append(attr.method)
+        orders.append(attr.interaction_order)
+
+        if len(feats) == 1:
+            token = feats[0]
+            if token in index:
+                idx = index[token]
+                matrix[idx][idx] = float(attr.value)
+            continue
+
+        if len(feats) > 2:
+            # For higher-order interactions, distribute the value across all ordered pairs.
+            iterator = [
+                (index[a], index[b])
+                for i, a in enumerate(feats)
+                for j, b in enumerate(feats)
+                if i != j and a in index and b in index
+            ]
+        else:
+            a, b = feats
+            if a not in index or b not in index:
+                continue
+            iterator = [(index[a], index[b]), (index[b], index[a])]
+
+        for i, j in iterator:
+            matrix[i][j] = float(attr.value)
+
+    flat_values = [value for row in matrix for value in row]
+    if not flat_values:
+        flat_values = [0.0]
+
+    return {
+        "matrix": matrix,
+        "features": tokens,
+        "methods": methods,
+        "orders": orders,
+        "min": float(min(flat_values)),
+        "max": float(max(flat_values)),
+    }
+

audioop.py

@@ -0,0 +1,31 @@
+# audioop stub for Python 3.13+ compatibility.
+# Python 3.13 removed the built-in audioop module.
+# pydub imports audioop but only calls it for audio operations (which this app never uses).
+# This empty stub prevents the ImportError at startup.
+
+def add(*a, **k): ...
+def adpcm2lin(*a, **k): ...
+def alaw2lin(*a, **k): ...
+def avg(*a, **k): return 0
+def avgpp(*a, **k): return 0
+def bias(*a, **k): ...
+def byteswap(*a, **k): ...
+def cross(*a, **k): return 0
+def findfactor(*a, **k): return 0.0
+def findfit(*a, **k): return (0, 0.0)
+def findmax(*a, **k): return 0
+def getsample(*a, **k): return 0
+def lin2adpcm(*a, **k): return (b"", 0)
+def lin2alaw(*a, **k): return b""
+def lin2lin(*a, **k): return b""
+def lin2ulaw(*a, **k): return b""
+def max(*a, **k): return 0
+def maxpp(*a, **k): return 0
+def minmax(*a, **k): return (0, 0)
+def mul(*a, **k): return b""
+def ratecv(*a, **k): return (b"", (0, 0))
+def reverse(*a, **k): return b""
+def rms(*a, **k): return 0
+def tostereo(*a, **k): return b""
+def tomono(*a, **k): return b""
+def ulaw2lin(*a, **k): return b""

datasets/BBQ_disamb.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f0155987c1d7c2c396025e7288a7b4b299d7095f1f671fd982fa3f8af99680f1
+size 4244

datasets/BarExam_qa.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8c81699f93bdccee10a45839297d509cfea95cda6961ac2575017fa03902e66d
+size 14377

datasets/CNN_dailymail.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:211e6e255435474cbb1725327bde6e2d61f56f8aa4e4738fadf2e378e5dd8d96
+size 54840

datasets/bbh_causal_judgement.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3363c531c6eed1ed5661671f98c95e44b0876927d86228409169cbf4c9638767
+size 14737

datasets/bbh_snarks.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e06480910f969da34c642c4d581dbbbf1c3b922f9529a243b3f56293b19824e5
+size 1871

datasets/drop.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1aabf12072359977d9bb42a445d79e116e60b127a67e088346e1fef0fa89e6db
+size 18030

datasets/eSNLI.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3c57e79a35ffabf0f857be75b137d12174fee46661da71017df79028fd7be41b
+size 3109

datasets/fever.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6059f85f2d1a2ebe2cdeb9786a237f942fa159124376024b271c6bb55a7e46c9
+size 1583

datasets/hotpot_qa.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a746b3324975632a5d94c703e203fc19bd8ce6b3cc163050d29488c6e662d127
+size 66525

datasets/medical_qa.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d761128bf4e3c1641c8f6976dead0c1d78059ea46e6c3c261b75d53c2c3567cb
+size 8908

loader/__init__.py

@@ -0,0 +1,84 @@
+# loader/__init__.py
+
+from .models import (
+    BaseModel,
+    ModelInfo,
+    InferenceResult,
+    get_model,
+)
+from .data import (
+    list_datasets,
+    get_examples,
+    get_example_by_id,
+)
+from .perplexity import (
+    score_continuation,
+    score_reference_autoregressive_hf,
+    score_reference_autoregressive_vllm,
+)
+from .utils import device_info, set_seed
+
+# --- Optional vLLM hook (only works inside the vLLM container) ---
+import traceback as _traceback
+
+_VLLM_IMPORT_ERR = None
+
+try:
+    from .models_vllm import VLLMModel
+except Exception as exc:  # capture and expose real reason
+    VLLMModel = None
+    _VLLM_IMPORT_ERR = _traceback.format_exc()
+
+    def get_model_vllm(size: str = "small") -> BaseModel:
+        msg = (
+            "vLLM is not available in this environment. "
+            "Run inside the vllm_meng.sif container to use get_model_vllm()."
+        )
+        if _VLLM_IMPORT_ERR:
+            msg += f"\nUnderlying import error:\n{_VLLM_IMPORT_ERR}"
+        raise RuntimeError(msg)  # surface actual import failure
+
+else:
+
+    def get_model_vllm(size: str = "small") -> BaseModel:
+        return VLLMModel(size)
+
+# --- Optional vLLM multimodal hook ---
+try:
+    from .models_mm_vllm import VLLMMultimodalModel, list_mm_model_keys, get_mm_model_info
+
+    def get_model_mm_vllm(key: str = "vl_small") -> BaseModel:
+        return VLLMMultimodalModel(key)
+
+except Exception:
+    VLLMMultimodalModel = None
+
+    def get_model_mm_vllm(key: str = "vl_small") -> BaseModel:
+        raise RuntimeError(
+            "Multimodal vLLM is not available in this environment. "
+            "Run inside the vLLM container to use get_model_mm_vllm()."
+        )
+
+    def list_mm_model_keys() -> list[str]:
+        return []
+
+    def get_mm_model_info(_key: str):
+        return None
+
+
+def debug_vllm_status():
+    """
+    Print basic diagnostics about vLLM availability/import state.
+    """
+    import sys, os
+
+    print("cwd:", os.getcwd())
+    print("sys.path[0:3]:", sys.path[:3])
+    try:
+        import vllm  # type: ignore
+
+        print("vllm import: OK", getattr(vllm, '__version__', 'unknown'))
+    except Exception as exc:  # pragma: no cover - diagnostic helper
+        print("vllm import: FAILED", exc)
+    if _VLLM_IMPORT_ERR:
+        print("Stored VLLM import error trace:\n", _VLLM_IMPORT_ERR)

loader/data.py

@@ -0,0 +1,179 @@
+# loader/data.py
+from typing import List, Dict
+from pathlib import Path
+import csv
+
+# -------------------------------------------------------------------
+# 1. Mapping from dataset key → CSV filename
+# -------------------------------------------------------------------
+_DATASET_FILES: Dict[str, str] = {
+    "bar_exam":        "BarExam_qa.csv",
+    "causal_judgment": "bbh_causal_judgement.csv",
+    "snarks":          "bbh_snarks.csv",
+    "bbq_disamb":      "BBQ_disamb.csv",
+    "cnn_dailymail":   "CNN_dailymail.csv",
+    "drop":            "drop.csv",
+    "esnli":           "eSNLI.csv",
+    "fever":           "fever.csv",
+    "hotpot_qa":       "hotpot_qa.csv",
+    "medical_qa":      "medical_qa.csv",
+}
+
+# -------------------------------------------------------------------
+# 1b. Human-readable display names for UI
+# -------------------------------------------------------------------
+_DATASET_DISPLAY_NAMES: Dict[str, str] = {
+    "bar_exam":        "Bar Exam Questions",
+    "causal_judgment": "Causal Judgment",
+    "snarks":          "Snarks",
+    "bbq_disamb":      "BBQ Disambiguation",
+    "cnn_dailymail":   "CNN / DailyMail Summaries",
+    "drop":            "DROP Reading Comprehension",
+    "esnli":           "e-SNLI Natural Language Inference",
+    "fever":           "FEVER Fact Checking",
+    "hotpot_qa":       "HotpotQA Multi-hop Questions",
+    "medical_qa":      "Medical Questions",
+}
+
+# -------------------------------------------------------------------
+# 2. Where the CSVs live (loader/../datasets/)
+# -------------------------------------------------------------------
+def _datasets_dir() -> Path:
+    return (Path(__file__).resolve().parent.parent / "datasets").resolve()
+
+# -------------------------------------------------------------------
+# 3. Pick the first non-empty column among several candidates
+# -------------------------------------------------------------------
+def _pick_first_nonempty(raw: Dict[str, str], candidates: List[str]) -> str:
+    for c in candidates:
+        val = raw.get(c)
+        if val is not None and str(val).strip() != "":
+            return str(val)
+    return ""
+
+# -------------------------------------------------------------------
+# 4. Load a single CSV file and normalize it to our schema
+# -------------------------------------------------------------------
+def _load_one_dataset(name: str, filename: str) -> List[Dict[str, str]]:
+    """
+    Reads a CSV file and converts each row to our standard format:
+
+        {
+            "id": "example_1",
+            "context": "...",
+            "prompt":  "...",
+            "answer":  "..."   # optional
+        }
+
+    Only the first 10 rows are kept.
+    """
+    path = _datasets_dir() / filename
+    rows: List[Dict[str, str]] = []
+
+    # errors="replace" avoids Unicode crashes for imperfect CSVs
+    try:
+        with path.open("r", encoding="utf-8", errors="replace", newline="") as f:
+            reader = csv.DictReader(f)
+
+            for i, raw in enumerate(reader, start=1):
+
+                ex_id = raw.get("id") or raw.get("example_id") \
+                     or raw.get("uid") or f"example_{i}"
+
+                context = _pick_first_nonempty(raw, [
+                    "Context", "context",
+                    "passage", "article", "story", "premise",
+                    "paragraph", "document", "sentence1", "sent1", "background",
+                ])
+
+                prompt = _pick_first_nonempty(raw, [
+                    "Prompt", "prompt",
+                    "question", "input", "query",
+                    "sentence2", "sent2", "hypothesis",
+                    "qa_question", "title",
+                ])
+
+                answer = _pick_first_nonempty(raw, [
+                    "Answer", "answer",
+                    "target", "gold", "label", "output", "reference",
+                    "highlights",
+                ])
+
+                ex = {
+                    "id": str(ex_id),
+                    "context": context,
+                    "prompt": prompt,
+                }
+                if answer:
+                    ex["answer"] = answer
+
+                rows.append(ex)
+    except FileNotFoundError:
+        return []
+    except Exception:
+        # Keep import resilient in constrained environments (e.g., Spaces).
+        return []
+
+    return rows[:10]  # keep exactly 10 examples
+
+# -------------------------------------------------------------------
+# 5. Load all datasets ONCE when the module is imported
+# -------------------------------------------------------------------
+def _load_all_datasets() -> Dict[str, List[Dict[str, str]]]:
+    return {
+        name: _load_one_dataset(name, filename)
+        for name, filename in _DATASET_FILES.items()
+    }
+
+_DATA: Dict[str, List[Dict[str, str]]] = _load_all_datasets()  # ← cached
+
+# -------------------------------------------------------------------
+# 6. Public Functions — these are used by the app
+# -------------------------------------------------------------------
+def list_datasets() -> List[str]:
+    """Return all dataset names, sorted alphabetically."""
+    return sorted(_DATA.keys())
+
+
+def get_dataset_display_name(dataset: str) -> str:
+    """Return human-readable display name for a dataset."""
+    return _DATASET_DISPLAY_NAMES.get(dataset, dataset)
+
+
+def get_dataset_key_from_display_name(display_name: str) -> str:
+    """Convert display name back to internal key."""
+    # Create reverse mapping
+    for key, name in _DATASET_DISPLAY_NAMES.items():
+        if name == display_name:
+            return key
+    # If not found, assume it's already a key
+    return display_name
+
+
+def list_datasets_with_display_names() -> List[tuple[str, str]]:
+    """Return list of (key, display_name) tuples, sorted by display name."""
+    pairs = [(key, _DATASET_DISPLAY_NAMES.get(key, key)) for key in _DATA.keys()]
+    return sorted(pairs, key=lambda x: x[1])
+
+
+def list_dataset_display_names() -> List[str]:
+    """Return list of display names only, sorted alphabetically."""
+    names = [_DATASET_DISPLAY_NAMES.get(key, key) for key in _DATA.keys()]
+    return sorted(names)
+
+
+def get_examples(dataset: str, n: int = 10) -> List[Dict[str, str]]:
+    """Return up to n examples for a dataset."""
+    if dataset not in _DATA:
+        raise KeyError(f"Unknown dataset: {dataset}")
+    return _DATA[dataset][:n]
+
+
+def get_example_by_id(dataset: str, ex_id: str) -> Dict[str, str]:
+    """Return a single example whose ID matches ex_id."""
+    if dataset not in _DATA:
+        raise KeyError(f"Unknown dataset: {dataset}")
+    for ex in _DATA[dataset]:
+        if ex["id"] == ex_id:
+            return ex
+    raise KeyError(f"Example id '{ex_id}' not found in dataset '{dataset}'")

loader/models.py

@@ -0,0 +1,65 @@
+from dataclasses import dataclass, field
+from typing import List, Optional
+
+@dataclass
+class ModelInfo:
+    size: str
+    name: str
+    max_model_len: int
+    is_chat: bool
+    is_multimodal: bool = False
+    image_placeholder: Optional[str] = None
+    mm_data_key: Optional[str] = None
+
+@dataclass
+class InferenceResult:
+    prompt: str
+    text: str
+    token_ids: List[int] = field(default_factory=list)
+    # one scalar logprob per generated token (or None if not requested)
+    logprobs: Optional[List[float]] = None
+
+class BaseModel:
+    def info(self) -> ModelInfo:
+        raise NotImplementedError
+
+    def infer(
+        self,
+        prompts: List[str],
+        *,
+        max_tokens: int = 256,
+        temperature: float = 0.7,
+        top_p: float = 0.95,
+        logprobs: Optional[int] = None,
+        stop: Optional[List[str]] = None,
+    ) -> List[InferenceResult]:
+        raise NotImplementedError
+
+# --- Minimal working dummy model so imports/tests succeed ---
+class _DummyModel(BaseModel):
+    def __init__(self, size: str):
+        self._info = ModelInfo(
+            size=size,
+            name=f"dummy-{size}",
+            max_model_len=4096,
+            is_chat=False,
+        )
+
+    def info(self) -> ModelInfo:
+        return self._info
+
+    def infer(self, prompts: List[str], **kwargs) -> List[InferenceResult]:
+        outs: List[InferenceResult] = []
+        for p in prompts:
+            outs.append(InferenceResult(
+                prompt=p,
+                text=f"[dummy completion for]: {p}",
+                token_ids=[0, 1, 2],
+                logprobs=None
+            ))
+        return outs
+
+def get_model(size: str) -> BaseModel:
+    if size not in {"small", "medium", "large"}:
+        raise ValueError("size must be one of {'small','medium','large'}")
+    return _DummyModel(size)

loader/models_mm_hf.py

@@ -0,0 +1,252 @@
+# loader/models_mm_hf.py
+"""
+HuggingFace Transformers wrapper for multimodal VLMs (e.g. LLaVA).
+
+Provides the same interface as VLLMMultimodalModel so the attribution
+pipeline in mm_game.py can use it without changes.  Slower than vLLM
+(no continuous batching, no PagedAttention) but works on any GPU without
+specific vLLM version requirements.
+"""
+
+import math
+import os
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Sequence
+
+import torch
+from PIL import Image
+
+from .models import BaseModel, InferenceResult, ModelInfo
+
+
+# ---------------------------------------------------------------------------
+# Fake vLLM-compatible output objects (just enough for _score_from_prompt_logprobs)
+# ---------------------------------------------------------------------------
+
+@dataclass
+class _FakeCompletionOutput:
+    text: str
+    token_ids: List[int]
+
+
+@dataclass
+class _FakeRequestOutput:
+    outputs: List[_FakeCompletionOutput]
+    prompt_logprobs: Optional[List[Optional[Dict[int, float]]]]
+
+
+class _FakeSamplingParams:
+    """Mirrors vllm.SamplingParams fields used by mm_game.py."""
+    def __init__(self, **kwargs):
+        self.max_tokens = kwargs.get("max_tokens", 256)
+        self.temperature = kwargs.get("temperature", 0.0)
+        self.top_p = kwargs.get("top_p", 1.0)
+        self.logprobs = kwargs.get("logprobs", 0)
+        self.prompt_logprobs = kwargs.get("prompt_logprobs", None)
+
+
+# ---------------------------------------------------------------------------
+# HF wrapper that mimics vLLM LLM interface
+# ---------------------------------------------------------------------------
+
+class _HFLLMShim:
+    """
+    Drop-in replacement for vllm.LLM that uses HuggingFace transformers.
+
+    Only implements the methods that mm_game.py actually calls:
+      - get_tokenizer()
+      - generate(requests, sampling_params)  where requests are dicts
+        with "prompt" and "multi_modal_data" keys
+    """
+
+    def __init__(self, model_name: str, device: str = "cuda"):
+        from transformers import AutoProcessor, LlavaNextForConditionalGeneration, LlavaForConditionalGeneration
+
+        self._model_name = model_name
+        self._device = device
+
+        print(f"  [HF] Loading processor for {model_name} ...")
+        self._processor = AutoProcessor.from_pretrained(model_name)
+
+        print(f"  [HF] Loading model {model_name} (torch.float16) ...")
+        # Try LlavaNext first, fall back to Llava
+        try:
+            self._model = LlavaNextForConditionalGeneration.from_pretrained(
+                model_name, torch_dtype=torch.float16, device_map=device,
+            )
+        except Exception:
+            self._model = LlavaForConditionalGeneration.from_pretrained(
+                model_name, torch_dtype=torch.float16, device_map=device,
+            )
+        self._model.eval()
+        print(f"  [HF] Model loaded on {device}.")
+
+    def get_tokenizer(self):
+        """Return the tokenizer (same interface as vllm.LLM.get_tokenizer)."""
+        return self._processor.tokenizer
+
+    def generate(
+        self,
+        requests: List[Dict[str, Any]],
+        sampling_params: Any,
+    ) -> List[_FakeRequestOutput]:
+        """
+        Mimic vllm.LLM.generate() for multimodal requests.
+
+        Each request is {"prompt": str, "multi_modal_data": {"image": PIL.Image}}.
+        Returns list of _FakeRequestOutput with .prompt_logprobs and .outputs.
+        """
+        results = []
+        for req in requests:
+            prompt = req["prompt"]
+            mm_data = req.get("multi_modal_data", {})
+
+            # Extract image
+            raw_image = mm_data.get("image")
+            if isinstance(raw_image, list):
+                raw_image = raw_image[0] if raw_image else None
+
+            want_prompt_logprobs = getattr(sampling_params, "prompt_logprobs", None)
+            max_new = getattr(sampling_params, "max_tokens", 256)
+
+            if want_prompt_logprobs:
+                result = self._score_prompt_logprobs(prompt, raw_image, max_new)
+            else:
+                result = self._generate_text(prompt, raw_image, max_new, sampling_params)
+
+            results.append(result)
+        return results
+
+    def _process_inputs(self, prompt: str, image: Optional[Image.Image]):
+        """Prepare model inputs from prompt + image."""
+        if image is not None:
+            inputs = self._processor(text=prompt, images=image, return_tensors="pt")
+        else:
+            inputs = self._processor(text=prompt, return_tensors="pt")
+        return {k: v.to(self._device) for k, v in inputs.items()}
+
+    def _score_prompt_logprobs(
+        self, prompt: str, image: Optional[Image.Image], max_new: int
+    ) -> _FakeRequestOutput:
+        """Compute per-token log-probabilities of the prompt (for log-prob attribution)."""
+        inputs = self._process_inputs(prompt, image)
+        input_ids = inputs["input_ids"]
+        seq_len = input_ids.shape[1]
+
+        with torch.inference_mode():
+            outputs = self._model(**inputs)
+            # logits shape: (1, seq_len, vocab_size)
+            logits = outputs.logits[0]  # (seq_len, vocab_size)
+
+        log_probs = torch.log_softmax(logits, dim=-1)  # (seq_len, vocab_size)
+
+        # Build prompt_logprobs in vLLM format:
+        # prompt_logprobs[i] = {token_id: logprob} for position i
+        # Position i's logprob comes from logits at position i-1 (autoregressive)
+        prompt_logprobs: List[Optional[Dict[int, float]]] = [None]  # position 0 has no predecessor
+        flat_ids = input_ids[0].tolist()
+
+        for i in range(1, seq_len):
+            tid = flat_ids[i]
+            lp = log_probs[i - 1, tid].item()
+            # Include top-5 for the position (matching vLLM's logprobs=5)
+            top_vals, top_ids = torch.topk(log_probs[i - 1], k=min(5, log_probs.shape[-1]))
+            pos_dict = {int(top_ids[j]): float(top_vals[j]) for j in range(len(top_ids))}
+            if tid not in pos_dict:
+                pos_dict[tid] = float(lp)
+            prompt_logprobs.append(pos_dict)
+
+        return _FakeRequestOutput(
+            outputs=[_FakeCompletionOutput(text="", token_ids=[])],
+            prompt_logprobs=prompt_logprobs,
+        )
+
+    def _generate_text(
+        self,
+        prompt: str,
+        image: Optional[Image.Image],
+        max_new: int,
+        sampling_params: Any,
+    ) -> _FakeRequestOutput:
+        """Generate text (for generation attribution)."""
+        inputs = self._process_inputs(prompt, image)
+        input_len = inputs["input_ids"].shape[1]
+
+        temperature = getattr(sampling_params, "temperature", 0.0)
+        do_sample = temperature > 0
+
+        with torch.inference_mode():
+            output_ids = self._model.generate(
+                **inputs,
+                max_new_tokens=max_new,
+                do_sample=do_sample,
+                temperature=temperature if do_sample else None,
+            )
+
+        new_ids = output_ids[0, input_len:].tolist()
+        text = self._processor.tokenizer.decode(new_ids, skip_special_tokens=True)
+
+        return _FakeRequestOutput(
+            outputs=[_FakeCompletionOutput(text=text, token_ids=new_ids)],
+            prompt_logprobs=None,
+        )
+
+
+# ---------------------------------------------------------------------------
+# Public model class
+# ---------------------------------------------------------------------------
+
+class HFMultimodalModel(BaseModel):
+    """
+    HuggingFace Transformers multimodal model.
+
+    Same interface as VLLMMultimodalModel but uses transformers instead of vLLM.
+    The .llm attribute is a _HFLLMShim that mimics vLLM's LLM API.
+    """
+
+    def __init__(self, model_name: str, device: str = "cuda"):
+        user = os.environ.get("USER", "user")
+        cache_dir = f"/tmp/{user}/hf_cache"
+        os.makedirs(cache_dir, exist_ok=True)
+        os.environ["HF_HOME"] = cache_dir
+        os.environ["TRANSFORMERS_CACHE"] = cache_dir
+        os.environ["HUGGINGFACE_HUB_CACHE"] = cache_dir
+
+        self._info = ModelInfo(
+            size="hf_" + model_name.split("/")[-1],
+            name=model_name,
+            max_model_len=4096,
+            is_chat=True,
+            is_multimodal=True,
+            image_placeholder="<image>",
+            mm_data_key="image",
+        )
+
+        self.llm = _HFLLMShim(model_name, device)
+
+    @property
+    def info(self) -> ModelInfo:
+        return self._info
+
+    def infer(
+        self,
+        prompts: List[str],
+        max_new_tokens: int = 256,
+        temperature: float = 0.0,
+        top_p: float = 1.0,
+        logprobs: int = 0,
+    ) -> List[InferenceResult]:
+        """Generate text for a list of prompts (no image — text-only generation)."""
+        sp = _FakeSamplingParams(
+            max_tokens=max_new_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            logprobs=logprobs,
+        )
+        results = []
+        for prompt in prompts:
+            req = {"prompt": prompt, "multi_modal_data": {}}
+            out = self.llm.generate([req], sp)[0]
+            seq = out.outputs[0]
+            results.append(InferenceResult(prompt=prompt, text=seq.text, token_ids=seq.token_ids))
+        return results

loader/models_mm_vllm.py

@@ -0,0 +1,165 @@
+# loader/models_mm_vllm.py
+
+import os
+from typing import List, Optional
+
+from vllm import LLM, SamplingParams
+
+from .models import BaseModel, InferenceResult, ModelInfo
+
+# ---------------------------------------------------------------------
+# Multimodal model registry (vLLM-compatible VLMs only)
+# ---------------------------------------------------------------------
+_MM_MODEL_CONFIGS = {
+    "vl_small": ModelInfo(
+        size="vl_small",
+        name="Qwen/Qwen2-VL-2B-Instruct",
+        max_model_len=4096,
+        is_chat=True,
+        is_multimodal=True,
+        image_placeholder="<|image_pad|>",
+        mm_data_key="image",
+    ),
+    "vl_large": ModelInfo(
+        size="vl_large",
+        name="Qwen/Qwen2-VL-7B-Instruct",
+        max_model_len=4096,
+        is_chat=True,
+        is_multimodal=True,
+        image_placeholder="<|image_pad|>",
+        mm_data_key="image",
+    ),
+    "llavamed": ModelInfo(
+        size="llavamed",
+        name="llava-hf/llava-v1.6-mistral-7b-hf",
+        max_model_len=4096,
+        is_chat=True,
+        is_multimodal=True,
+        image_placeholder="<image>",
+        mm_data_key="image",
+    ),
+}
+
+
+def list_mm_model_keys() -> List[str]:
+    return list(_MM_MODEL_CONFIGS.keys())
+
+
+def get_mm_model_info(key: str) -> Optional[ModelInfo]:
+    return _MM_MODEL_CONFIGS.get(key)
+
+
+class VLLMMultimodalModel(BaseModel):
+    def __init__(self, key: str = "vl_small") -> None:
+        if key not in _MM_MODEL_CONFIGS:
+            raise ValueError(f"Unknown multimodal model key: {key}")
+
+        info = _MM_MODEL_CONFIGS[key]
+        gpu_utilization = _resolve_gpu_memory_utilization(key)
+        max_model_len = _resolve_max_model_len(key, info.max_model_len)
+        self._info = ModelInfo(
+            size=info.size,
+            name=info.name,
+            max_model_len=max_model_len,
+            is_chat=info.is_chat,
+            is_multimodal=True,
+            image_placeholder=info.image_placeholder,
+            mm_data_key=info.mm_data_key,
+        )
+
+        user = os.environ.get("USER", "user")
+        cache_dir = f"/tmp/{user}/hf_cache"
+        os.makedirs(cache_dir, exist_ok=True)
+        os.environ["HF_HOME"] = cache_dir
+        os.environ["TRANSFORMERS_CACHE"] = cache_dir
+        os.environ["HUGGINGFACE_HUB_CACHE"] = cache_dir
+        os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+        download_dir = cache_dir
+
+        self.llm = LLM(
+            model=info.name,
+            tensor_parallel_size=1,
+            download_dir=download_dir,
+            trust_remote_code=True,
+            disable_log_stats=True,
+            max_model_len=max_model_len,
+            gpu_memory_utilization=gpu_utilization,
+        )
+
+    @property
+    def info(self) -> ModelInfo:
+        return self._info
+
+    def infer(
+        self,
+        prompts: List[str],
+        max_new_tokens: int = 256,
+        temperature: float = 0.0,
+        top_p: float = 1.0,
+        logprobs: int = 0,
+    ) -> List[InferenceResult]:
+        sp = SamplingParams(
+            max_tokens=max_new_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            logprobs=logprobs,
+        )
+
+        outs = self.llm.generate(prompts, sp)
+
+        results: List[InferenceResult] = []
+        for prompt, out in zip(prompts, outs):
+            seq = out.outputs[0]
+            text = seq.text
+            token_ids = list(seq.token_ids)
+            results.append(
+                InferenceResult(
+                    prompt=prompt,
+                    text=text,
+                    token_ids=token_ids,
+                )
+            )
+
+        return results
+
+
+def _parse_env_float(name: str) -> Optional[float]:
+    value = os.getenv(name)
+    if value is None or value.strip() == "":
+        return None
+    try:
+        return float(value)
+    except ValueError as exc:
+        raise ValueError(f"{name} must be a float, got {value!r}") from exc
+
+
+def _parse_env_int(name: str) -> Optional[int]:
+    value = os.getenv(name)
+    if value is None or value.strip() == "":
+        return None
+    try:
+        return int(value)
+    except ValueError as exc:
+        raise ValueError(f"{name} must be an int, got {value!r}") from exc
+
+
+def _resolve_gpu_memory_utilization(key: str) -> float:
+    key_upper = key.upper()
+    override = _parse_env_float(f"VLLM_GPU_MEMORY_UTILIZATION_{key_upper}")
+    if override is not None:
+        return override
+    override = _parse_env_float("VLLM_GPU_MEMORY_UTILIZATION")
+    if override is not None:
+        return override
+    return 0.6
+
+
+def _resolve_max_model_len(key: str, default: int) -> int:
+    key_upper = key.upper()
+    override = _parse_env_int(f"VLLM_MAX_MODEL_LEN_{key_upper}")
+    if override is not None:
+        return override
+    override = _parse_env_int("VLLM_MAX_MODEL_LEN")
+    if override is not None:
+        return override
+    return default

loader/models_vllm.py

@@ -0,0 +1,218 @@
+# loader/models_vllm.py
+
+import logging
+import os
+from typing import List, Optional
+
+from vllm import LLM, SamplingParams
+
+from .models import BaseModel, InferenceResult, ModelInfo
+
+logger = logging.getLogger(__name__)
+
+# ---------------------------------------------------------------------
+# Model registry
+# ---------------------------------------------------------------------
+_MODEL_CONFIGS = {
+    "small": ModelInfo(
+        size="small",
+        name="Qwen/Qwen3-4B-Instruct-2507",
+        max_model_len=8192,   # we cap to 32k, no need for 262k
+        is_chat=True,
+    ),
+    "medium": ModelInfo(
+        size="medium",
+        name="mistralai/Mistral-7B-Instruct-v0.2",
+        max_model_len=8192,
+        is_chat=True,
+    ),
+    "large": ModelInfo(
+        size="large",
+        name="Qwen/Qwen3-30B-A3B-Instruct-2507",
+        max_model_len=2048,   # reduced to fit KV cache within GPU memory
+        is_chat=True,
+    ),
+}
+
+
+class VLLMModel(BaseModel):
+    def __init__(self, size: str = "small") -> None:
+        if size not in _MODEL_CONFIGS:
+            raise ValueError(f"Unknown VLLM model size: {size}")
+
+        info = _MODEL_CONFIGS[size]
+        gpu_utilization = _resolve_gpu_memory_utilization(size)
+        max_model_len = _resolve_max_model_len(size, info.max_model_len)
+        max_num_seqs = _resolve_max_num_seqs(size)
+        max_batched_tokens = _resolve_max_batched_tokens(size)
+        self._info = ModelInfo(
+            size=info.size,
+            name=info.name,
+            max_model_len=max_model_len,
+            is_chat=info.is_chat,
+        )
+
+        # ---------------------------------------------------
+        # Use node-local /tmp for all HuggingFace / vLLM cache
+        # ---------------------------------------------------
+        user = os.environ.get("USER", "user")
+        cache_dir = f"/tmp/{user}/hf_cache"
+
+        # Make sure the directory exists
+        os.makedirs(cache_dir, exist_ok=True)
+
+        # Tell HF + transformers + vLLM to use this cache
+        os.environ["HF_HOME"] = cache_dir
+        os.environ["TRANSFORMERS_CACHE"] = cache_dir
+        os.environ["HUGGINGFACE_HUB_CACHE"] = cache_dir
+        os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+
+        # vLLM will also use this as download_dir
+        download_dir = cache_dir
+
+        # ----------------------------------
+        # Log GPU state & resolved config
+        # ----------------------------------
+        try:
+            import torch
+            if torch.cuda.is_available():
+                free, total = torch.cuda.mem_get_info(0)
+                logger.info(
+                    "GPU memory: free=%.2f GiB  total=%.2f GiB  (before vLLM init)",
+                    free / (1 << 30), total / (1 << 30),
+                )
+        except Exception:
+            pass
+
+        enforce_eager = os.environ.get("VLLM_ENFORCE_EAGER", "1") == "1"
+
+        logger.info(
+            "VLLMModel(%s): model=%s  gpu_util=%.2f  max_model_len=%d  "
+            "enforce_eager=%s  max_num_seqs=%s  max_batched_tokens=%s",
+            size, info.name, gpu_utilization, max_model_len,
+            enforce_eager, max_num_seqs, max_batched_tokens,
+        )
+
+        # ----------------------------------
+        # Create the vLLM LLM instance
+        # ----------------------------------
+        self.llm = LLM(
+            model=info.name,
+            tensor_parallel_size=1,
+            download_dir=download_dir,
+            trust_remote_code=True,
+            disable_log_stats=True,
+            max_model_len=max_model_len,
+            gpu_memory_utilization=gpu_utilization,
+            max_num_seqs=max_num_seqs,
+            max_num_batched_tokens=max_batched_tokens,
+            enforce_eager=enforce_eager,
+        )
+
+    # For convenience, expose model metadata
+    @property
+    def info(self) -> ModelInfo:
+        return self._info
+
+    def infer(
+    self,
+    prompts: List[str],
+    max_new_tokens: int = 256,
+    temperature: float = 0.0,
+    top_p: float = 1.0,
+    logprobs: int = 0,
+    ) -> List[InferenceResult]:
+        sp = SamplingParams(
+            max_tokens=max_new_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            logprobs=logprobs,
+        )
+
+        outs = self.llm.generate(prompts, sp)
+
+        results: List[InferenceResult] = []
+        for prompt, out in zip(prompts, outs):
+            seq = out.outputs[0]
+            text = seq.text
+
+            # vLLM gives a list of int token IDs
+            token_ids = list(seq.token_ids)
+
+            # You *can* still compute token_logprobs here if you want,
+            # but don't pass them into InferenceResult yet.
+            # (You can add this field to the dataclass later if needed.)
+
+            results.append(
+                InferenceResult(
+                    prompt=prompt,
+                    text=text,
+                    token_ids=token_ids,
+                )
+            )
+
+        return results
+
+
+def _parse_env_float(name: str) -> Optional[float]:
+    value = os.getenv(name)
+    if value is None or value.strip() == "":
+        return None
+    try:
+        return float(value)
+    except ValueError as exc:
+        raise ValueError(f"{name} must be a float, got {value!r}") from exc
+
+
+def _parse_env_int(name: str) -> Optional[int]:
+    value = os.getenv(name)
+    if value is None or value.strip() == "":
+        return None
+    try:
+        return int(value)
+    except ValueError as exc:
+        raise ValueError(f"{name} must be an int, got {value!r}") from exc
+
+
+def _resolve_max_num_seqs(size: str) -> Optional[int]:
+    size_key = size.upper()
+    override = _parse_env_int(f"VLLM_MAX_NUM_SEQS_{size_key}")
+    if override is not None:
+        return override
+    return _parse_env_int("VLLM_MAX_NUM_SEQS")
+
+
+def _resolve_max_batched_tokens(size: str) -> Optional[int]:
+    size_key = size.upper()
+    override = _parse_env_int(f"VLLM_MAX_BATCHED_TOKENS_{size_key}")
+    if override is not None:
+        return override
+    return _parse_env_int("VLLM_MAX_BATCHED_TOKENS")
+
+
+def _resolve_gpu_memory_utilization(size: str) -> float:
+    size_key = size.upper()
+    # Per-size override: VLLM_GPU_MEMORY_UTILIZATION_LARGE, etc.
+    override = _parse_env_float(f"VLLM_GPU_MEMORY_UTILIZATION_{size_key}")
+    if override is not None:
+        return override
+    # Global override: VLLM_GPU_MEMORY_UTILIZATION or short alias VLLM_GPU_UTIL
+    override = _parse_env_float("VLLM_GPU_MEMORY_UTILIZATION")
+    if override is not None:
+        return override
+    override = _parse_env_float("VLLM_GPU_UTIL")
+    if override is not None:
+        return override
+    # Defaults
+    return 0.90
+
+
+def _resolve_max_model_len(size: str, default: int) -> int:
+    size_key = size.upper()
+    override = _parse_env_int(f"VLLM_MAX_MODEL_LEN_{size_key}")
+    if override is not None:
+        return override
+    override = _parse_env_int("VLLM_MAX_MODEL_LEN")
+    if override is not None:
+        return override
+    return default

loader/perplexity.py

@@ -0,0 +1,538 @@
+import math
+import os
+import logging
+from typing import Callable, List, Dict, Any, Optional
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from .models import BaseModel
+
+logger = logging.getLogger(__name__)
+
+# Signature: (question, context, answer_tag, reference_text) -> full_prompt_str
+PromptFormatter = Callable[[str, str, str, str], str]
+
+
+def default_prompt_formatter(
+    question: str,
+    context: str,
+    answer_tag: str,
+    reference_text: str,
+) -> str:
+    """
+    Basic prompt layout used by default.
+
+    You can override this via the `prompt_formatter` argument if a
+    specific model needs a different template (e.g., chat template).
+    """
+    return f"{question}\nContext: {context}\n{answer_tag} {reference_text}"
+
+
+def _join_prefix_continuation(prefix: str, continuation: str) -> str:
+    """Join prefix and continuation with a single space when needed."""
+    if not prefix:
+        return continuation
+    if not continuation:
+        return prefix
+    if prefix[-1].isspace() or continuation[0].isspace():
+        return prefix + continuation
+    return prefix + " " + continuation
+
+
+def _continuation_prompt_formatter(
+    _question: str,
+    context: str,
+    _answer_tag: str,
+    reference_text: str,
+) -> str:
+    """Prompt formatter that concatenates context and continuation text."""
+    return _join_prefix_continuation(context, reference_text)
+
+
+def score_continuation(
+    model: BaseModel,
+    prefix: str,
+    continuation: str,
+    *,
+    max_new_tokens: Optional[int] = None,
+) -> Dict[str, Any]:
+    """
+    Compute teacher-forced logprob and perplexity of `continuation` given `prefix`.
+    Uses vLLM prompt logprobs when available.
+    """
+    if not continuation:
+        return {
+            "avg_logprob": float("-inf"),
+            "perplexity": float("inf"),
+            "per_token": [],
+            "target_len": 0,
+            "sequence_logprobs": [],
+        }
+
+    if hasattr(model, "llm"):
+        result = score_reference_autoregressive_vllm(
+            model,
+            question="",
+            masked_inputs=[prefix],
+            reference_text=continuation,
+            answer_tag="",
+            prompt_formatter=_continuation_prompt_formatter,
+            max_new_tokens=max_new_tokens,
+        )
+        token_logprobs = result.get("token_logprobs", [])
+        per_token = token_logprobs[0] if token_logprobs else []
+        target_len = int(result.get("target_len", 0) or 0)
+        seq_logprobs = result.get("sequence_logprobs", [])
+        if seq_logprobs:
+            total_logprob = float(seq_logprobs[0])
+        else:
+            avg_lp = float(result.get("avg_logprob", float("-inf")))
+            total_logprob = avg_lp * target_len if target_len else float("-inf")
+        if target_len > 0 and math.isfinite(total_logprob):
+            avg_nll = -total_logprob / target_len
+            ppl = math.exp(avg_nll)
+        else:
+            avg_nll = float("inf")
+            ppl = float("inf")
+        return {
+            "avg_logprob": result.get("avg_logprob", float("-inf")),
+            "perplexity": result.get("perplexity", float("inf")),
+            "total_logprob": total_logprob,
+            "avg_nll": avg_nll,
+            "per_token": per_token,
+            "target_len": target_len,
+            "sequence_logprobs": seq_logprobs,
+        }
+
+    raise RuntimeError(
+        "score_continuation requires a VLLMModel with .llm "
+        "(use loader.get_model_vllm inside the vLLM container)."
+    )
+
+# --- Real (teacher-forced) scorer using Hugging Face generate() ---
+def score_reference_autoregressive_hf(
+    model,                 # HF AutoModelForCausalLM (on CUDA or CPU)
+    tokenizer,             # matching HF tokenizer
+    question: str,         # e.g., "Count the number of r's in strawberry."
+    masked_inputs: List[str],  # batch of masked contexts (strings)
+    reference_token_ids: List[int],  # tokenized reference answer ids
+    *,
+    answer_tag: str = "Answer:",
+    prompt_formatter: PromptFormatter = default_prompt_formatter,
+) -> Dict[str, Any]:
+    """
+    Teacher-forced next-token log-probs of the reference answer across a batch of masked contexts.
+    Returns:
+      {
+        "avg_logprob": float,                # mean over tokens, then mean over batch
+        "perplexity": float,                 # exp(-avg_logprob)
+        "target_len": int,                   # number of next-token steps (len(ref_ids)-1)
+        "sequence_logprobs": List[float],    # total logprob per masked input (len= batch)
+        "token_logprobs": List[List[float]]  # per-token logprobs per masked input
+      }
+    """
+    # --- ensure tokenizer/model have a pad token ---
+    if tokenizer.pad_token is None:
+        if tokenizer.eos_token is not None:
+            tokenizer.pad_token = tokenizer.eos_token
+        else:
+            tokenizer.add_special_tokens({"pad_token": "[PAD]"})
+            # only call if model is an HF CausalLM with embeddings
+            try:
+                model.resize_token_embeddings(len(tokenizer))
+            except Exception:
+                pass
+    tokenizer.padding_side = "left"
+    try:
+        model.config.pad_token_id = tokenizer.pad_token_id
+    except Exception:
+        pass
+    B = len(masked_inputs)
+    # We will accumulate total log-prob of the reference per masked input
+    batch_seq_logprobs = np.zeros(B, dtype=np.float64)
+
+    T = max(0, len(reference_token_ids) - 1)  # number of next-token predictions
+    if T == 0:
+        return {
+            "avg_logprob": float("-inf"),
+            "perplexity": float("inf"),
+            "target_len": 0,
+            "sequence_logprobs": batch_seq_logprobs.tolist(),
+        }
+
+    for j in range(T):
+        # Reference prefix up to and including token j
+        prefix_ids = reference_token_ids[: j + 1]
+        # Decode the entire prefix to keep spacing/special-tokens consistent
+        prefix_str = tokenizer.decode(prefix_ids, skip_special_tokens=False)
+
+        # Build prompts for this step
+        prompts = [
+            prompt_formatter(question, ctx, answer_tag, prefix_str)
+            for ctx in masked_inputs
+        ]
+
+        inputs = tokenizer(
+            prompts, return_tensors="pt", padding=True, truncation=True
+        ).to(model.device)
+
+        with torch.inference_mode():
+            out = model.generate(
+                **inputs,
+                max_new_tokens=1,                 # predict exactly one next token
+                output_scores=True,
+                return_dict_in_generate=True,
+                pad_token_id=tokenizer.eos_token_id,
+            )
+
+        # out["scores"] is a list (len = #new tokens = 1) of logits [batch, vocab]
+        logits = torch.stack(out["scores"]).swapaxes(0, 1)[:, 0, :]   # (B, V)
+        logprobs = F.log_softmax(logits, dim=-1)
+
+        target_next_id = reference_token_ids[j + 1]
+        step_lp = logprobs[:, target_next_id].detach().cpu().numpy()  # (B,)
+
+        batch_seq_logprobs += step_lp
+
+        # clean up to keep memory low
+        del inputs, out, logits, logprobs
+
+    # Average over tokens, then average across the batch
+    avg_lp = float((batch_seq_logprobs / T).mean())
+    ppl = math.exp(-avg_lp)
+
+    return {
+        "avg_logprob": avg_lp,
+        "perplexity": ppl,
+        "target_len": T,
+        "sequence_logprobs": batch_seq_logprobs.tolist(),
+    }
+
+def score_reference_autoregressive_vllm(
+    vllm_model,              # your loader.VLLMModel instance
+    question: str,
+    masked_inputs: List[str],
+    reference_text: str,     # e.g., " The answer is 3."
+    *,
+    answer_tag: str = "Answer:",
+    prompt_formatter: PromptFormatter = default_prompt_formatter,
+    top_k: int = 5,
+    debug: bool = False,
+    debug_index: int = 0,
+    max_new_tokens: Optional[int] = None,
+) -> Dict[str, Any]:
+    """
+    vLLM-native teacher-forced perplexity over a *given* reference answer.
+
+    Batching behavior:
+      - Build one full prompt per masked context.
+      - Send ALL prompts to vLLM in a single `llm.generate` call.
+      - vLLM returns prompt-level logprobs for each token position.
+      - For each prompt, slice out the part corresponding to the answer
+        and sum the logprobs of those answer tokens.
+
+    If debug=True, we print detailed info for the example at `debug_index`:
+      - Full prompt string
+      - Base vs full token lengths
+      - Each answer token + its logprob and running perplexity.
+    """
+    # 1) Get underlying vLLM engine & tokenizer
+    assert hasattr(vllm_model, "llm"), f"Expected VLLMModel wrapper, got {type(vllm_model)}"
+
+    llm = vllm_model.llm
+    try:
+        tokenizer = llm.get_tokenizer()
+    except Exception as e:
+        raise RuntimeError("Could not access vLLM tokenizer; check vLLM version / API") from e
+
+    B = len(masked_inputs)
+    if B == 0:
+        return {
+            "avg_logprob": float("-inf"),
+            "perplexity": float("inf"),
+            "target_len": 0,
+            "sequence_logprobs": [],
+            "token_logprobs": [],
+        }
+
+    # 2) Build all prompts and tokenizations in a batch
+    full_prompts: List[str] = []
+    base_ids_list: List[List[int]] = []
+    full_ids_list: List[List[int]] = []
+    ref_ids_list: List[List[int]] = []
+
+    ref_token_count: Optional[int] = None
+
+    for ctx in masked_inputs:
+        # base prompt: no answer text
+        base_prompt = prompt_formatter(question, ctx, answer_tag, "")
+        # full prompt: includes gold answer text
+        full_prompt = prompt_formatter(question, ctx, answer_tag, reference_text)
+
+        base_ids = tokenizer(base_prompt, add_special_tokens=False)["input_ids"]
+        full_ids = tokenizer(full_prompt, add_special_tokens=False)["input_ids"]
+
+        # Safety check
+        if not full_ids or len(full_ids) <= len(base_ids):
+            base_ids_list.append(base_ids)
+            full_ids_list.append(full_ids)
+            ref_ids_list.append([])
+            full_prompts.append(full_prompt)
+            continue
+
+        ref_ids = full_ids[len(base_ids):]  # tokens of the answer region
+
+        if ref_token_count is None:
+            ref_token_count = len(ref_ids)
+        else:
+            # make all answers share a common length (minimum across examples)
+            ref_token_count = min(ref_token_count, len(ref_ids))
+
+        base_ids_list.append(base_ids)
+        full_ids_list.append(full_ids)
+        ref_ids_list.append(ref_ids)
+        full_prompts.append(full_prompt)
+
+    # If no valid ref tokens at all, bail out
+    if ref_token_count is None or ref_token_count == 0:
+        return {
+            "avg_logprob": float("-inf"),
+            "perplexity": float("inf"),
+            "target_len": 0,
+            "sequence_logprobs": [float("-inf")] * B,
+            "token_logprobs": [[] for _ in range(B)],
+        }
+
+    # Truncate all ref_ids to the common length
+    ref_len = ref_token_count
+    ref_ids_list = [ref_ids[:ref_len] for ref_ids in ref_ids_list]
+
+    # 3) Call vLLM *once* on all full prompts (batched)
+    from vllm import SamplingParams
+
+    sp = SamplingParams(
+        max_tokens=max_new_tokens or 1,  # keep tiny generation; default minimal
+        temperature=0.0,
+        top_p=1.0,
+        logprobs=top_k,         # store top-k logprobs per position
+        prompt_logprobs=1,      # request logprobs for prompt tokens
+    )
+
+    results = llm.generate(full_prompts, sp)
+
+    # 4) For each prompt, sum the logprobs over the answer segment
+    per_seq_logprobs: List[float] = []
+    per_seq_token_logprobs: List[List[float]] = []
+
+    def to_float(v):
+        # v might be a Logprob object or a raw float
+        return float(getattr(v, "logprob", v))
+
+    for idx_ex, (base_ids, full_ids, ref_ids) in enumerate(
+        zip(base_ids_list, full_ids_list, ref_ids_list)
+    ):
+        if not ref_ids or not full_ids:
+            per_seq_logprobs.append(float("-inf"))
+            per_seq_token_logprobs.append([])
+            continue
+
+        req_out = results[idx_ex]
+        prompt_logprobs = getattr(req_out, "prompt_logprobs", None)
+        if prompt_logprobs is None:
+            raise RuntimeError(
+                "req_out.prompt_logprobs is None. Adjust to match your vLLM version."
+            )
+
+        # Align lengths: vLLM may add BOS token
+        if len(prompt_logprobs) == len(full_ids) + 1:
+            prompt_logprobs = prompt_logprobs[1:]
+        elif len(prompt_logprobs) != len(full_ids):
+            L = min(len(prompt_logprobs), len(full_ids))
+            prompt_logprobs = prompt_logprobs[:L]
+            full_ids = full_ids[:L]
+
+        seq_lp = 0.0
+        token_logprobs: List[float] = []
+
+        # Optional debug header
+        if debug and idx_ex == debug_index:
+            print("\n=== DEBUG vLLM PPL example", idx_ex, "===")
+            print("Full prompt:\n", full_prompts[idx_ex])
+            print("\nBase token length:", len(base_ids))
+            print("Full token length:", len(full_ids))
+            print("Answer token length (ref_len):", ref_len)
+            print("\nAnswer tokens and logprobs:")
+
+        for offset, token_id in enumerate(ref_ids[:ref_len]):
+            pos = len(base_ids) + offset  # index in full sequence
+
+            if pos >= len(prompt_logprobs):
+                token_lp = -20.0
+                seq_lp += token_lp
+                token_logprobs.append(token_lp)
+                if debug and idx_ex == debug_index:
+                    tok_str = tokenizer.convert_ids_to_tokens([token_id])[0]
+                    print(f"  pos={pos:<3} token={tok_str!r:>10}  logprob={token_lp: .4f}  (OUT OF RANGE)")
+                continue
+
+            cand_dict = prompt_logprobs[pos] or {}
+
+            if token_id in cand_dict:
+                token_lp = to_float(cand_dict[token_id])
+            else:
+                if cand_dict:
+                    floor = min(to_float(v) for v in cand_dict.values())
+                    token_lp = floor - 5.0
+                else:
+                    token_lp = -20.0
+
+            seq_lp += token_lp
+            token_logprobs.append(token_lp)
+
+            if debug and idx_ex == debug_index:
+                tok_str = tokenizer.convert_ids_to_tokens([token_id])[0]
+                avg_lp_so_far = seq_lp / (offset + 1)
+                ppl_so_far = math.exp(-avg_lp_so_far)
+                print(
+                    f"  pos={pos:<3} token={tok_str!r:>10}  "
+                    f"logprob={token_lp: .4f}  "
+                    f"cum_avg_lp={avg_lp_so_far: .4f}  "
+                    f"cum_ppl={ppl_so_far: .4f}"
+                )
+
+        per_seq_logprobs.append(seq_lp)
+        per_seq_token_logprobs.append(token_logprobs)
+
+    # 5) Aggregate across examples
+    arr = np.array(per_seq_logprobs, dtype=np.float64)
+    T = int(ref_len)
+
+    if T == 0:
+        return {
+            "avg_logprob": float("-inf"),
+            "perplexity": float("inf"),
+            "target_len": 0,
+            "sequence_logprobs": per_seq_logprobs,
+            "token_logprobs": per_seq_token_logprobs,
+        }
+
+    avg_lp = float((arr / T).mean())
+    ppl = math.exp(-avg_lp)
+
+    return {
+        "avg_logprob": avg_lp,
+        "perplexity": ppl,
+        "target_len": T,
+        "sequence_logprobs": per_seq_logprobs,
+        "token_logprobs": per_seq_token_logprobs,
+    }
+
+
+def score_continuation_batch(
+    model: BaseModel,
+    prefixes: List[str],
+    continuation: str,
+    *,
+    max_new_tokens: Optional[int] = None,
+    batch_size: Optional[int] = None,
+) -> List[Dict[str, Any]]:
+    """
+    Batched variant of score_continuation. Returns list aligned to prefixes.
+    """
+    if not prefixes:
+        return []
+
+    B = batch_size or int(os.getenv("VLLM_SCORE_BATCH_SIZE",
+                                     os.getenv("ATTRLLM_VLLM_BATCH_SIZE", "128")))
+    B = max(1, int(B))
+    outputs: List[Optional[Dict[str, Any]]] = [None] * len(prefixes)
+
+    def _failed_output() -> Dict[str, Any]:
+        return {
+            "avg_logprob": float("-inf"),
+            "perplexity": float("inf"),
+            "total_logprob": float("-inf"),
+            "avg_nll": float("inf"),
+            "per_token": [],
+            "target_len": 0,
+            "sequence_logprobs": [float("-inf")],
+        }
+
+    def _score_chunk(chunk_prefixes: List[str], offset: int, chunk_batch_size: int) -> None:
+        if not chunk_prefixes:
+            return
+        try:
+            res = score_reference_autoregressive_vllm(
+                model,
+                question="",
+                masked_inputs=chunk_prefixes,
+                reference_text=continuation,
+                answer_tag="",
+                prompt_formatter=_continuation_prompt_formatter,
+                max_new_tokens=max_new_tokens,
+            )
+            seq_lp = res.get("sequence_logprobs", [])
+            per_token = res.get("token_logprobs", [])
+            target_len = int(res.get("target_len", 0) or 0)
+
+            if len(seq_lp) != len(chunk_prefixes):
+                raise RuntimeError(
+                    "vLLM returned mismatched sequence_logprobs length: "
+                    f"expected={len(chunk_prefixes)} got={len(seq_lp)}"
+                )
+
+            for i, lp in enumerate(seq_lp):
+                idx = offset + i
+                total_logprob = float(lp)
+                if target_len > 0 and math.isfinite(total_logprob):
+                    avg_nll = -total_logprob / target_len
+                    ppl = math.exp(avg_nll)
+                else:
+                    avg_nll = float("inf")
+                    ppl = float("inf")
+                outputs[idx] = {
+                    "avg_logprob": res.get("avg_logprob", float("-inf")),
+                    "perplexity": res.get("perplexity", ppl),
+                    "total_logprob": total_logprob,
+                    "avg_nll": avg_nll,
+                    "per_token": per_token[i] if i < len(per_token) else [],
+                    "target_len": target_len,
+                    "sequence_logprobs": [total_logprob],
+                }
+        except Exception as exc:
+            # vLLM occasionally asserts in large-batch prompt_logprob scoring.
+            # Back off to smaller chunks instead of failing the whole request.
+            n = len(chunk_prefixes)
+            if n == 1:
+                logger.warning(
+                    "score_continuation_batch failed for single prefix; returning -inf fallback. error=%s",
+                    exc,
+                )
+                outputs[offset] = _failed_output()
+                return
+
+            next_batch_size = max(1, min(chunk_batch_size // 2, n // 2))
+            logger.warning(
+                "[BACKOFF_ACTIVE] score_continuation_batch chunk failed; retrying smaller chunks. "
+                "chunk_size=%d batch_size=%d next_batch_size=%d error=%s",
+                n,
+                chunk_batch_size,
+                next_batch_size,
+                exc,
+            )
+            for sub_start in range(0, n, next_batch_size):
+                sub_chunk = chunk_prefixes[sub_start : sub_start + next_batch_size]
+                _score_chunk(sub_chunk, offset + sub_start, next_batch_size)
+
+    for start in range(0, len(prefixes), B):
+        chunk = prefixes[start : start + B]
+        _score_chunk(chunk, start, B)
+
+    # Guarantee alignment even in worst-case failures.
+    for i, item in enumerate(outputs):
+        if item is None:
+            outputs[i] = _failed_output()
+
+    return outputs  # type: ignore[return-value]

loader/results.py

@@ -0,0 +1,468 @@
+from typing import Dict, Any, Tuple, Optional
+from pathlib import Path
+import copy
+import json
+import os
+from itertools import combinations
+try:
+    from attribution.utils import (
+        influence_interactions,
+        mobius_to_influence,
+        mobius_to_shapley,
+        mobius_to_banzhaf,
+        shapley_interactions,
+        banzhaf_interactions,
+    )
+except Exception:  # pragma: no cover - keep loader resilient if attribution stack is unavailable
+    influence_interactions = None
+    mobius_to_influence = None
+    mobius_to_shapley = None
+    mobius_to_banzhaf = None
+    shapley_interactions = None
+    banzhaf_interactions = None
+
+try:
+    from huggingface_hub import hf_hub_download
+except Exception:  # pragma: no cover
+    hf_hub_download = None
+
+
+def _resolve_results_dir() -> Path:
+    """Resolve results directory: env, repo root, or on HF Space fallback to cwd/results."""
+    env_dir = os.getenv("ATTRLLM_RESULTS_DIR")
+    if env_dir:
+        return Path(env_dir).resolve()
+    repo_root = Path(__file__).resolve().parent.parent
+    default = (repo_root / "results").resolve()
+    if default.exists():
+        return default
+    # Hugging Face Space: repo may be at /app but cwd can be workspace root with results/
+    if os.getenv("SPACE_ID") or os.getenv("HF_SPACE"):
+        cwd_results = (Path.cwd() / "results").resolve()
+        if cwd_results.exists():
+            return cwd_results
+    return default
+
+
+_RESULTS_DIR = _resolve_results_dir()
+_CACHE: Dict[Path, Dict[str, Any]] = {}
+
+
+def _maybe_download_from_space(path: Path, force_download: bool = False) -> Optional[Path]:
+    """Best-effort download from Hugging Face Space repo if running in Spaces.
+    When force_download=True, re-fetches from Hub (can resolve LFS/Xet pointers in container)."""
+    if hf_hub_download is None:
+        return None
+    repo_id = os.getenv("SPACE_ID")
+    if not repo_id:
+        return None
+    rel = None
+    try:
+        rel = path.relative_to(Path.cwd())
+    except Exception:
+        try:
+            repo_root = Path(__file__).resolve().parent.parent
+            rel = path.relative_to(repo_root)
+        except Exception:
+            rel = None
+    if rel is None:
+        return None
+    try:
+        kwargs = {
+            "repo_id": repo_id,
+            "filename": str(rel),
+            "repo_type": "space",
+            "local_dir": str(Path.cwd()),
+            "local_dir_use_symlinks": False,
+        }
+        if force_download:
+            kwargs["force_download"] = True
+        local_path = hf_hub_download(**kwargs)
+        return Path(local_path)
+    except Exception:
+        return None
+
+
+def _dataset_results_path(model_size: str, dataset: str) -> Path:
+    # New layout:
+    #   results/reference_answer/{model_size}/{dataset}/{example_id}/{scalarizer}/{feature_level}.json
+    return _RESULTS_DIR / "reference_answer" / model_size / dataset
+
+
+def _public_results_path(dataset: str) -> Path:
+    # Legacy/public layout:
+    #   results/public/{dataset}/{example_id}/{scalarizer}/{feature_level}/{method}.json
+    return _RESULTS_DIR / "public" / dataset
+
+
+def _find_reference_example_dir(
+    dataset: str,
+    ex_id: str,
+    preferred_size: Optional[str] = None,
+) -> Tuple[Optional[str], Optional[Path]]:
+    """Locate an existing reference_answer example dir across model sizes."""
+    sizes = [s for s in (preferred_size, "large", "medium", "small") if s]
+    seen = set()
+    for size in sizes:
+        if size in seen:
+            continue
+        seen.add(size)
+        candidate = _dataset_results_path(size, dataset) / ex_id
+        if candidate.exists():
+            return size, candidate
+    return None, None
+
+
+def _load_dataset_results(path: Path) -> Dict[str, Any]:
+    if path in _CACHE:
+        return _CACHE[path]
+    # On HF Space, fetch from Hub first so we get resolved content (not LFS/Xet pointer in container).
+    if os.getenv("SPACE_ID"):
+        fetched = _maybe_download_from_space(path, force_download=True)
+        if fetched and fetched.exists():
+            path = fetched
+    if not path.exists():
+        fetched = _maybe_download_from_space(path)
+        if fetched and fetched.exists():
+            path = fetched
+        else:
+            _CACHE[path] = {}
+            return _CACHE[path]
+    try:
+        with path.open("r", encoding="utf-8") as f:
+            data = json.load(f)
+    except Exception:
+        # If this is a Git LFS/Xet pointer or otherwise unreadable JSON,
+        # try to fetch the real file from the Space repo.
+        fetched = _maybe_download_from_space(path)
+        if fetched and fetched.exists():
+            try:
+                with fetched.open("r", encoding="utf-8") as f:
+                    data = json.load(f)
+            except Exception:
+                _CACHE[path] = {}
+                return _CACHE[path]
+        else:
+            _CACHE[path] = {}
+            return _CACHE[path]
+    _CACHE[path] = data if isinstance(data, dict) else {}
+    return _CACHE[path]
+
+
+def _load_method_payload(
+    example_dir: Path,
+    method: str,
+    scalarizer: Optional[str] = None,
+    feature_level: Optional[str] = None,
+) -> Dict[str, Any]:
+    """Load the Möbius-dict payload for the given example.
+
+    New layout (method-agnostic files):
+        {example_dir}/{scalarizer}/{feature_level}.json
+
+    The ``method`` parameter is not used for file lookup (the Möbius dict
+    is identical for Shapley and Banzhaf); it is only passed through so
+    the caller can apply the correct aggregation afterwards.
+    """
+    scalarizer = (scalarizer or "").strip()
+    feature_level = (feature_level or "").strip()
+
+    candidates = []
+    if scalarizer and feature_level:
+        candidates.append(example_dir / scalarizer / f"{feature_level}.json")
+        # If the requested level is missing, fall back to other levels.
+        for lvl in ("word", "sentence", "paragraph"):
+            if lvl != feature_level:
+                candidates.append(example_dir / scalarizer / f"{lvl}.json")
+    elif scalarizer:
+        for lvl in ("word", "sentence", "paragraph"):
+            candidates.append(example_dir / scalarizer / f"{lvl}.json")
+    elif feature_level:
+        for sc in ("semantic_similarity", "geomean_jointprob", "logprob", "jointprob", "half_simlog"):
+            candidates.append(example_dir / sc / f"{feature_level}.json")
+    else:
+        # Default preference order
+        candidates.extend(
+            [
+                example_dir / "semantic_similarity" / "word.json",
+                example_dir / "semantic_similarity" / "sentence.json",
+                example_dir / "semantic_similarity" / "paragraph.json",
+            ]
+        )
+
+    # Legacy fallback: old layout with method in filename
+    # (e.g. {scalarizer}/{feature_level}/shapley.json)
+    method_lower = str(method or "shapley").lower()
+    if scalarizer and feature_level:
+        candidates.append(example_dir / scalarizer / feature_level / f"{method_lower}.json")
+    elif scalarizer:
+        for lvl in ("word", "sentence", "paragraph"):
+            candidates.append(example_dir / scalarizer / lvl / f"{method_lower}.json")
+
+    for candidate in candidates:
+        payload = _load_dataset_results(candidate)
+        if payload:
+            return payload
+
+    return {}
+
+
+def _parse_sparse_key(raw_key: str) -> Tuple[int, ...]:
+    key = str(raw_key).strip()
+    if not key:
+        return ()
+    return tuple(int(part) for part in key.split(",") if part != "")
+
+
+def _assign_unique_labels(chunks):
+    counts: Dict[str, int] = {}
+    labels = []
+    for idx, chunk in enumerate(chunks):
+        normalized = " ".join((str(chunk) if chunk is not None else "").split())
+        if not normalized:
+            normalized = "<chunk %d>" % (idx + 1)
+        counts[normalized] = counts.get(normalized, 0) + 1
+        suffix = " (%d)" % counts[normalized] if counts[normalized] > 1 else ""
+        labels.append("%s%s" % (normalized, suffix))
+    return labels
+
+
+def _normalize_public_payload(payload: Dict[str, Any], method: str, top_k: int = 10) -> Dict[str, Any]:
+    if not isinstance(payload, dict):
+        return {}
+
+    normalized = dict(payload)
+    features = normalized.get("features")
+    mobius_raw = normalized.get("mobius_dict")
+    if not isinstance(features, list) or not isinstance(mobius_raw, dict):
+        return normalized
+
+    # Keep existing explicit fields if they are already present.
+    if isinstance(normalized.get("token_scores"), dict) and (
+        isinstance(normalized.get("pairwise"), dict) or isinstance(normalized.get("interactions_2"), dict)
+    ):
+        return normalized
+
+    method = (method or "shapley").lower()
+    if method not in {"shapley", "banzhaf", "influence"}:
+        method = "shapley"
+
+    mobius_sparse: Dict[Tuple[int, ...], float] = {}
+    for key, raw_val in mobius_raw.items():
+        try:
+            val = float(raw_val)
+        except Exception:
+            continue
+        try:
+            loc = _parse_sparse_key(str(key))
+        except Exception:
+            continue
+        mobius_sparse[tuple(sorted(loc))] = val
+
+    token_scores: Dict[str, float] = {}
+    index_scores: Dict[int, float] = {}
+    pairwise_acc: Dict[Tuple[int, int], float] = {}
+
+    if mobius_sparse and mobius_to_shapley is not None:
+        if method == "shapley":
+            singleton_dict = mobius_to_shapley(mobius_sparse)
+            pair_list = shapley_interactions(mobius_sparse, order=2, top_k=top_k) or []
+        elif method == "banzhaf":
+            singleton_dict = mobius_to_banzhaf(mobius_sparse)
+            pair_list = banzhaf_interactions(mobius_sparse, order=2, top_k=top_k) or []
+        elif mobius_to_influence is not None and influence_interactions is not None:
+            singleton_dict = mobius_to_influence(mobius_sparse)
+            pair_list = influence_interactions(mobius_sparse, order=2, top_k=top_k) or []
+        else:
+            singleton_dict = {}
+            pair_list = []
+
+        for loc, val in singleton_dict.items():
+            if len(loc) != 1:
+                continue
+            idx = int(loc[0])
+            if 0 <= idx < len(features):
+                feat_name = str(features[idx])
+                val_f = float(val)
+                token_scores[feat_name] = token_scores.get(feat_name, 0.0) + val_f
+                index_scores[idx] = index_scores.get(idx, 0.0) + val_f
+
+        for loc, val in pair_list:
+            if len(loc) != 2:
+                continue
+            i, j = int(loc[0]), int(loc[1])
+            if 0 <= i < len(features) and 0 <= j < len(features):
+                key = (i, j) if i <= j else (j, i)
+                pairwise_acc[key] = float(val)
+    else:
+        # Fallback path: local conversion if attribution utils are unavailable.
+        for loc, val in mobius_sparse.items():
+            k = len(loc)
+            if k == 0:
+                continue
+            if method == "shapley":
+                singleton_weight = 1.0 / float(k)
+            elif method == "banzhaf":
+                singleton_weight = 1.0 / float(2 ** (k - 1))
+            else:
+                singleton_weight = 1.0 / float(k)
+            for idx in loc:
+                if 0 <= idx < len(features):
+                    feat_name = str(features[idx])
+                    token_scores[feat_name] = token_scores.get(feat_name, 0.0) + singleton_weight * val
+                    index_scores[idx] = index_scores.get(idx, 0.0) + singleton_weight * val
+            if k >= 2:
+                if method == "shapley":
+                    pairwise_weight = 1.0 / float(k - 1)
+                elif method == "banzhaf":
+                    pairwise_weight = 1.0 / float(2 ** (k - 2))
+                else:
+                    pairwise_weight = 1.0 / float(k - 1)
+                for i, j in combinations(sorted(loc), 2):
+                    pairwise_acc[(i, j)] = pairwise_acc.get((i, j), 0.0) + pairwise_weight * val
+
+    if token_scores and not isinstance(normalized.get("token_scores"), dict):
+        normalized["token_scores"] = token_scores
+
+    unique_feature_labels = [str(x) for x in features]
+    if index_scores and features and not (
+        isinstance(normalized.get("features"), list)
+        and normalized["features"]
+        and isinstance(normalized["features"][0], dict)
+    ):
+        original_feature_texts = [str(x) for x in features]
+        unique_feature_labels = _assign_unique_labels(original_feature_texts)
+        normalized["feature_texts"] = original_feature_texts
+        normalized["features"] = [
+            {
+                "feature": original_feature_texts[i],
+                "value": float(index_scores.get(i, 0.0)),
+                "index": i,
+            }
+            for i in range(len(original_feature_texts))
+        ]
+
+    if pairwise_acc and not isinstance(normalized.get("pairwise"), dict):
+        sorted_pairs = sorted(pairwise_acc.items(), key=lambda kv: abs(kv[1]), reverse=True)
+        if top_k and top_k > 0:
+            sorted_pairs = sorted_pairs[:top_k]
+        # Use feature labels (not raw indices) so chips/text views can map edges directly.
+        normalized["pairwise"] = {
+            "%s|%s" % (unique_feature_labels[i], unique_feature_labels[j]): float(v)
+            for (i, j), v in sorted_pairs
+            if 0 <= i < len(unique_feature_labels) and 0 <= j < len(unique_feature_labels)
+        }
+        normalized["pairwise_interactions"] = [
+            {"features": [unique_feature_labels[i], unique_feature_labels[j]], "value": float(v)}
+            for (i, j), v in sorted_pairs
+            if 0 <= i < len(unique_feature_labels) and 0 <= j < len(unique_feature_labels)
+        ]
+
+    return normalized
+
+def get_result_by_id(
+    model_size: str,
+    dataset: str,
+    ex_id: str,
+    scalarizer: Optional[str] = None,
+    feature_level: Optional[str] = None,
+) -> Dict:
+    """
+    Retrieve saved interaction results (e.g., from Attribution) for a datapoint.
+    Expected to return Shapley, Banzhaf, and Influence views
+    (and their interactions) as JSON-like dict.
+
+    Args:
+        model_size: "small" | "medium" | "large"
+        dataset: dataset name (e.g., "reasoning")
+        ex_id: example id (e.g., "ex1")
+
+    Returns (example schema):
+    {
+      "meta": {"model_size": "small", "dataset": "reasoning", "ex_id": "ex1"},
+      "shapley": { "token_scores": {...}, "pairwise": {...}, "higher_order": {...} },
+      "banzhaf": { "token_scores": {...}, "pairwise": {...}, "higher_order": {...} },
+      "influence": { "token_scores": {...}, "pairwise": {...}, "higher_order": {...} }
+    }
+    """
+    dataset_root = _dataset_results_path(model_size, dataset)
+    example_dir = dataset_root / ex_id
+    if not example_dir.exists():
+        alt_size, alt_dir = _find_reference_example_dir(dataset, ex_id, model_size)
+        if alt_dir is not None:
+            model_size = alt_size or model_size
+            example_dir = alt_dir
+
+    # New layout: a single method-agnostic file per (scalarizer, feature_level).
+    # Load it once and derive both Shapley and Banzhaf aggregations.
+    if example_dir.exists():
+        payload = _load_method_payload(
+            example_dir,
+            "shapley",  # method param unused for file lookup in new layout
+            scalarizer=scalarizer,
+            feature_level=feature_level,
+        )
+        if payload:
+            record: Dict[str, Any] = {}
+            record["shapley"] = _normalize_public_payload(copy.deepcopy(payload), "shapley")
+            record["banzhaf"] = _normalize_public_payload(copy.deepcopy(payload), "banzhaf")
+            record["influence"] = _normalize_public_payload(copy.deepcopy(payload), "influence")
+
+            source_meta: Dict[str, Any] = {
+                "dataset": dataset,
+                "example_id": ex_id,
+                "model_size": model_size,
+                "source_layout": "results/reference_answer/{model_size}/{dataset}/{example_id}/{scalarizer}/{feature_level}.json",
+            }
+            file_meta = payload.get("meta") if isinstance(payload, dict) else None
+            if isinstance(file_meta, dict):
+                source_meta.update(file_meta)
+            record["meta"] = source_meta
+            return record
+
+    # Fallback: legacy/public layout with method-specific files.
+    public_root = _public_results_path(dataset)
+    public_example_dir = public_root / ex_id
+    if not public_example_dir.exists():
+        return {}
+
+    shapley_payload = _load_method_payload(
+        public_example_dir,
+        "shapley",
+        scalarizer=scalarizer,
+        feature_level=feature_level,
+    )
+    banzhaf_payload = _load_method_payload(
+        public_example_dir,
+        "banzhaf",
+        scalarizer=scalarizer,
+        feature_level=feature_level,
+    )
+    if not shapley_payload and not banzhaf_payload:
+        return {}
+
+    record = {}
+    if shapley_payload:
+        record["shapley"] = _normalize_public_payload(copy.deepcopy(shapley_payload), "shapley")
+    if banzhaf_payload:
+        record["banzhaf"] = _normalize_public_payload(copy.deepcopy(banzhaf_payload), "banzhaf")
+    if shapley_payload:
+        record["influence"] = _normalize_public_payload(copy.deepcopy(shapley_payload), "influence")
+    elif banzhaf_payload:
+        record["influence"] = _normalize_public_payload(copy.deepcopy(banzhaf_payload), "influence")
+
+    source_meta = {
+        "dataset": dataset,
+        "example_id": ex_id,
+        "model_size": model_size,
+        "source_layout": "results/public/{dataset}/{example_id}/{scalarizer}/{feature_level}/{method}.json",
+    }
+    # Prefer per-file meta if available.
+    for payload in (shapley_payload, banzhaf_payload):
+        file_meta = payload.get("meta") if isinstance(payload, dict) else None
+        if isinstance(file_meta, dict):
+            source_meta.update(file_meta)
+            break
+    record["meta"] = source_meta
+
+    return record

loader/utils.py

@@ -0,0 +1,42 @@
+# /projects/bfoj/AttrLLM/loader/utils.py
+import os
+import random
+import platform
+from dataclasses import dataclass
+from typing import Dict, Any, Optional
+
+try:
+    import torch
+except Exception:
+    torch = None
+
+@dataclass
+class DeviceInfo:
+    python_version: str
+    platform: str
+    cuda_available: bool
+    cuda_device: Optional[str]
+    torch_version: Optional[str]
+
+def device_info() -> Dict[str, Any]:
+    """Return a dict with basic runtime & device diagnostics."""
+    cuda_avail = bool(torch and torch.cuda.is_available())
+    return {
+        "python_version": platform.python_version(),
+        "platform": platform.platform(),
+        "torch_version": getattr(torch, "__version__", None),
+        "cuda_available": cuda_avail,
+        "cuda_device": (torch.cuda.get_device_name(0) if (cuda_avail and torch) else None),
+    }
+
+def set_seed(seed: int) -> None:
+    """Seed Python & Torch RNGs for reproducibility (best-effort)."""
+    random.seed(seed)
+    if torch:
+        try:
+            torch.manual_seed(seed)
+            if torch.cuda.is_available():
+                torch.cuda.manual_seed_all(seed)
+        except Exception:
+            pass
+    os.environ["PYTHONHASHSEED"] = str(seed)

requirements.txt

@@ -0,0 +1,11 @@
+requests
+pillow
+plotly
+numpy
+pandas
+matplotlib
+scipy
+huggingface_hub==0.25.2
+lightgbm
+scikit-learn
+torch

results/coco_mm/coco_133233_summary.txt

@@ -0,0 +1,166 @@
+=== COCO 133233 – shapley CROSS-MODAL attributions ===
+
+Caption: A marina filled with boats floating in crystal blue water.
+
+Image region values (7 regions):
+                      seg_4  -0.810950
+                      seg_3  +0.781893
+                      seg_0  +0.452894
+                      seg_1  +0.282716
+                      seg_5  +0.277897
+                      seg_2  -0.238533
+                      seg_6  +0.131396
+
+Caption token values (11 tokens):
+                      tok:a  +2.337490
+                 tok:marina  +1.058479
+               tok:floating  +0.194257
+                      tok:.  +0.173353
+                  tok:boats  -0.171325
+                     tok:in  -0.132036
+                   tok:with  -0.123359
+                  tok:water  +0.083010
+                 tok:filled  -0.064285
+                   tok:blue  -0.042108
+                tok:crystal  +0.017018
+
+Image×Image interactions (15):
+                                  seg_0 × seg_2  +1.001538
+                                  seg_2 × seg_4  -0.419727
+                                  seg_0 × seg_1  +0.414486
+                                  seg_3 × seg_6  +0.312252
+                                  seg_0 × seg_4  -0.312040
+                                  seg_2 × seg_6  -0.302635
+                                  seg_2 × seg_3  -0.283335
+                                  seg_0 × seg_5  +0.261252
+                                  seg_0 × seg_3  +0.223211
+                                  seg_2 × seg_5  +0.221214
+                                  seg_3 × seg_5  +0.188386
+                                  seg_1 × seg_3  +0.174208
+                                  seg_1 × seg_5  -0.121616
+                                  seg_0 × seg_6  +0.087689
+                                  seg_1 × seg_4  +0.077673
+
+Token×Token interactions (15):
+                             tok:a × tok:marina  +2.073310
+                           tok:a × tok:floating  +0.331114
+                              tok:a × tok:boats  -0.304516
+                                  tok:a × tok:.  +0.284624
+                          tok:floating × tok:in  -0.161938
+                      tok:filled × tok:floating  +0.157476
+                               tok:with × tok:.  -0.139770
+                           tok:boats × tok:blue  -0.138589
+                              tok:boats × tok:.  -0.138023
+                                 tok:a × tok:in  +0.118445
+                          tok:filled × tok:with  -0.105322
+                               tok:a × tok:blue  +0.103642
+                            tok:filled × tok:in  -0.100758
+                               tok:a × tok:with  +0.095216
+                          tok:filled × tok:blue  -0.094530
+
+Top-5 Cross-Modal interactions (seg × tok):
+                                  seg_0 × tok:a  +1.411575
+                                  seg_3 × tok:a  +0.862141
+                                  seg_2 × tok:a  +0.710077
+                                  seg_4 × tok:a  -0.695077
+                                  seg_5 × tok:a  +0.515772
+
+★ Positive token #1: tok:a  (+2.337490)  × top-3 segments:
+                                  seg_0 × tok:a  +1.411575
+                                  seg_3 × tok:a  +0.862141
+                                  seg_2 × tok:a  +0.710077
+
+★ Positive token #2: tok:marina  (+1.058479)  × top-3 segments:
+                             seg_0 × tok:marina  -0.148232
+                             seg_3 × tok:marina  -0.129773
+                             seg_1 × tok:marina  +0.069036
+
+★ Negative token #1: tok:boats  (-0.171325)  × top-3 segments:
+                              seg_3 × tok:boats  +0.214681
+                              seg_6 × tok:boats  +0.152483
+                              seg_0 × tok:boats  -0.133591
+
+★ Negative token #2: tok:in  (-0.132036)  × top-3 segments:
+                                 seg_3 × tok:in  -0.152741
+                                 seg_5 × tok:in  -0.145348
+                                 seg_1 × tok:in  -0.114416
+
+============================================================
+Per-token cross-modal interactions (top-5 each):
+
+  tok:a  (Shapley=+2.337490):
+                                    seg_0 × tok:a  +1.411575
+                                    seg_3 × tok:a  +0.862141
+                                    seg_2 × tok:a  +0.710077
+                                    seg_4 × tok:a  -0.695077
+                                    seg_5 × tok:a  +0.515772
+
+  tok:marina  (Shapley=+1.058479):
+                               seg_0 × tok:marina  -0.148232
+                               seg_3 × tok:marina  -0.129773
+                               seg_1 × tok:marina  +0.069036
+                               seg_4 × tok:marina  -0.056511
+                               seg_6 × tok:marina  +0.036980
+
+  tok:filled  (Shapley=-0.064285):
+                               seg_3 × tok:filled  -0.286883
+                               seg_0 × tok:filled  +0.145198
+                               seg_4 × tok:filled  -0.056520
+                               seg_1 × tok:filled  -0.049499
+                               seg_6 × tok:filled  +0.027488
+
+  tok:with  (Shapley=-0.123359):
+                                 seg_1 × tok:with  -0.171255
+                                 seg_2 × tok:with  +0.101950
+                                 seg_0 × tok:with  +0.045264
+                                 seg_3 × tok:with  +0.027135
+                                 seg_4 × tok:with  -0.026913
+
+  tok:boats  (Shapley=-0.171325):
+                                seg_3 × tok:boats  +0.214681
+                                seg_6 × tok:boats  +0.152483
+                                seg_0 × tok:boats  -0.133591
+                                seg_2 × tok:boats  -0.120506
+                                seg_1 × tok:boats  +0.114882
+
+  tok:floating  (Shapley=+0.194257):
+                             seg_3 × tok:floating  +0.275116
+                             seg_0 × tok:floating  +0.120704
+                             seg_5 × tok:floating  +0.076230
+                             seg_2 × tok:floating  -0.044183
+                             seg_1 × tok:floating  -0.043412
+
+  tok:in  (Shapley=-0.132036):
+                                   seg_3 × tok:in  -0.152741
+                                   seg_5 × tok:in  -0.145348
+                                   seg_1 × tok:in  -0.114416
+                                   seg_6 × tok:in  -0.047324
+                                   seg_0 × tok:in  +0.017775
+
+  tok:crystal  (Shapley=+0.017018):
+                              seg_5 × tok:crystal  +0.055339
+                              seg_3 × tok:crystal  +0.053869
+                              seg_4 × tok:crystal  -0.039220
+                              seg_2 × tok:crystal  -0.003780
+                              seg_6 × tok:crystal  -0.000362
+
+  tok:blue  (Shapley=-0.042108):
+                                 seg_0 × tok:blue  +0.119849
+                                 seg_3 × tok:blue  +0.054490
+                                 seg_1 × tok:blue  -0.018741
+                                 seg_5 × tok:blue  -0.016717
+                                 seg_4 × tok:blue  +0.008088
+
+  tok:water  (Shapley=+0.083010):
+                                seg_3 × tok:water  +0.292174
+                                seg_2 × tok:water  +0.090298
+                                seg_5 × tok:water  +0.057648
+                                seg_0 × tok:water  +0.036885
+                                seg_4 × tok:water  +0.033756
+
+  tok:.  (Shapley=+0.173353):
+                                    seg_0 × tok:.  +0.346550
+                                    seg_3 × tok:.  +0.161264
+                                    seg_6 × tok:.  +0.117877
+                                    seg_2 × tok:.  +0.058141
+                                    seg_5 × tok:.  -0.033294

results/coco_mm/coco_281447_summary.txt

@@ -0,0 +1,175 @@
+=== COCO 281447 – shapley CROSS-MODAL attributions ===
+
+Caption: Horse in fenced pasture with others grazing on grasses.
+
+Image region values (16 regions):
+                      seg_2  -2.278403
+                      seg_1  -0.822974
+                      seg_3  -0.298733
+                      seg_0  -0.184993
+                     seg_13  +0.134824
+                      seg_7  -0.016471
+                     seg_14  +0.016337
+                     seg_10  +0.011791
+                      seg_4  -0.011774
+                      seg_8  -0.010643
+                      seg_6  +0.009416
+                      seg_9  +0.009395
+                     seg_11  -0.004642
+                      seg_5  -0.003971
+                     seg_15  +0.003474
+                     seg_12  -0.000572
+
+Caption token values (11 tokens):
+                  tok:horse  +4.408324
+                    tok:fen  -0.224529
+                      tok:.  -0.050243
+                tok:pasture  -0.049678
+                   tok:with  +0.040849
+                     tok:in  +0.033160
+                    tok:ced  -0.029820
+                tok:grazing  +0.022595
+                 tok:others  +0.008450
+                     tok:on  +0.007305
+                tok:grasses  +0.001741
+
+Image×Image interactions (15):
+                                  seg_1 × seg_2  +1.362463
+                                  seg_0 × seg_1  +0.532141
+                                  seg_0 × seg_3  +0.282670
+                                  seg_0 × seg_2  +0.277382
+                                 seg_1 × seg_13  +0.189160
+                                  seg_1 × seg_3  +0.103962
+                                  seg_2 × seg_6  +0.072289
+                                  seg_1 × seg_9  -0.052994
+                                  seg_1 × seg_4  -0.033681
+                                  seg_0 × seg_9  +0.033601
+                                  seg_2 × seg_9  +0.028185
+                                seg_11 × seg_13  -0.025794
+                                  seg_2 × seg_3  -0.015109
+                                 seg_1 × seg_14  -0.015034
+                                 seg_7 × seg_11  -0.014243
+
+Token×Token interactions (15):
+                            tok:horse × tok:fen  -0.446448
+                              tok:horse × tok:.  -0.108351
+                             tok:horse × tok:in  +0.086563
+                        tok:horse × tok:pasture  -0.065111
+                        tok:horse × tok:grazing  +0.051056
+                            tok:horse × tok:ced  -0.040790
+                           tok:horse × tok:with  +0.035492
+                          tok:ced × tok:grazing  -0.024138
+                          tok:with × tok:others  +0.022160
+                         tok:horse × tok:others  -0.018359
+                           tok:in × tok:grasses  +0.016346
+                           tok:ced × tok:others  +0.016202
+                             tok:ced × tok:with  +0.014719
+                            tok:others × tok:on  +0.013668
+                          tok:ced × tok:pasture  -0.013105
+
+Top-5 Cross-Modal interactions (seg × tok):
+                              seg_1 × tok:horse  +2.256318
+                              seg_2 × tok:horse  -2.093609
+                              seg_0 × tok:horse  -0.486339
+                              seg_3 × tok:horse  -0.256331
+                             seg_13 × tok:horse  +0.089891
+
+★ Positive token #1: tok:horse  (+4.408324)  × top-3 segments:
+                              seg_1 × tok:horse  +2.256318
+                              seg_2 × tok:horse  -2.093609
+                              seg_0 × tok:horse  -0.486339
+
+★ Positive token #2: tok:with  (+0.040849)  × top-3 segments:
+                               seg_9 × tok:with  +0.016746
+                              seg_10 × tok:with  -0.010291
+                              seg_15 × tok:with  +0.004752
+
+★ Negative token #1: tok:fen  (-0.224529)  × top-3 segments:
+                                seg_7 × tok:fen  -0.025468
+                                seg_5 × tok:fen  +0.011604
+                               seg_10 × tok:fen  +0.007026
+
+★ Negative token #2: tok:.  (-0.050243)  × top-3 segments:
+                                 seg_14 × tok:.  +0.010019
+                                  seg_3 × tok:.  -0.003891
+                                  seg_4 × tok:.  -0.003789
+
+============================================================
+Per-token cross-modal interactions (top-5 each):
+
+  tok:horse  (Shapley=+4.408324):
+                                seg_1 × tok:horse  +2.256318
+                                seg_2 × tok:horse  -2.093609
+                                seg_0 × tok:horse  -0.486339
+                                seg_3 × tok:horse  -0.256331
+                               seg_13 × tok:horse  +0.089891
+
+  tok:in  (Shapley=+0.033160):
+                                   seg_5 × tok:in  -0.015456
+                                  seg_14 × tok:in  +0.007851
+                                  seg_13 × tok:in  +0.007115
+                                   seg_2 × tok:in  -0.006608
+                                   seg_6 × tok:in  -0.004884
+
+  tok:fen  (Shapley=-0.224529):
+                                  seg_7 × tok:fen  -0.025468
+                                  seg_5 × tok:fen  +0.011604
+                                 seg_10 × tok:fen  +0.007026
+                                  seg_6 × tok:fen  +0.005401
+                                  seg_2 × tok:fen  -0.004163
+
+  tok:ced  (Shapley=-0.029820):
+                                  seg_5 × tok:ced  +0.017460
+                                  seg_8 × tok:ced  -0.006840
+                                 seg_15 × tok:ced  +0.006074
+                                  seg_9 × tok:ced  -0.006071
+                                 seg_14 × tok:ced  -0.005179
+
+  tok:pasture  (Shapley=-0.049678):
+                             seg_12 × tok:pasture  -0.011237
+                             seg_15 × tok:pasture  -0.010195
+                              seg_4 × tok:pasture  +0.006291
+                              seg_0 × tok:pasture  +0.005938
+                              seg_2 × tok:pasture  -0.003224
+
+  tok:with  (Shapley=+0.040849):
+                                 seg_9 × tok:with  +0.016746
+                                seg_10 × tok:with  -0.010291
+                                seg_15 × tok:with  +0.004752
+                                 seg_2 × tok:with  -0.002544
+                                 seg_6 × tok:with  +0.002417
+
+  tok:others  (Shapley=+0.008450):
+                               seg_1 × tok:others  -0.015305
+                              seg_15 × tok:others  +0.014182
+                               seg_6 × tok:others  +0.009771
+                              seg_10 × tok:others  -0.009235
+                               seg_3 × tok:others  -0.009101
+
+  tok:grazing  (Shapley=+0.022595):
+                              seg_5 × tok:grazing  -0.007292
+                             seg_10 × tok:grazing  +0.004208
+                             seg_11 × tok:grazing  +0.004152
+                             seg_13 × tok:grazing  +0.003173
+                             seg_15 × tok:grazing  -0.002181
+
+  tok:on  (Shapley=+0.007305):
+                                  seg_14 × tok:on  +0.003634
+                                   seg_1 × tok:on  +0.003621
+                                   seg_9 × tok:on  -0.003367
+                                  seg_13 × tok:on  +0.001808
+                                   seg_5 × tok:on  -0.000969
+
+  tok:grasses  (Shapley=+0.001741):
+                              seg_4 × tok:grasses  +0.010303
+                              seg_6 × tok:grasses  -0.009457
+                              seg_5 × tok:grasses  +0.005944
+                             seg_10 × tok:grasses  +0.004237
+                             seg_15 × tok:grasses  +0.003841
+
+  tok:.  (Shapley=-0.050243):
+                                   seg_14 × tok:.  +0.010019
+                                    seg_3 × tok:.  -0.003891
+                                    seg_4 × tok:.  -0.003789
+                                    seg_5 × tok:.  +0.003730
+                                   seg_11 × tok:.  +0.003199

results/coco_mm/coco_307172_summary.txt

@@ -0,0 +1,183 @@
+=== COCO 307172 – shapley CROSS-MODAL attributions ===
+
+Caption: A baked dish on a plate being touched by a woman.
+
+Image region values (16 regions):
+                      seg_2  +0.136766
+                     seg_12  +0.125225
+                      seg_0  +0.114294
+                      seg_1  +0.088119
+                     seg_11  +0.085390
+                      seg_8  -0.084003
+                      seg_7  -0.081940
+                      seg_6  -0.051395
+                      seg_9  +0.046654
+                      seg_5  +0.045622
+                     seg_10  +0.023971
+                      seg_3  +0.020047
+                     seg_13  -0.019814
+                     seg_14  -0.017885
+                      seg_4  -0.012482
+                     seg_15  -0.001890
+
+Caption token values (12 tokens):
+                      tok:a  +2.176333
+                  tok:baked  +0.402312
+                   tok:dish  +0.027537
+                     tok:on  -0.021476
+                      tok:a  +0.018718
+                      tok:a  -0.015789
+                      tok:.  +0.013208
+                  tok:woman  +0.013079
+                tok:touched  +0.009330
+                  tok:being  +0.009256
+                     tok:by  +0.007656
+                  tok:plate  -0.000493
+
+Image×Image interactions (15):
+                                 seg_1 × seg_11  +0.103462
+                                  seg_0 × seg_1  -0.060013
+                                  seg_0 × seg_2  -0.058313
+                                  seg_2 × seg_6  +0.051352
+                                  seg_0 × seg_8  -0.044922
+                                 seg_0 × seg_11  +0.033494
+                                 seg_0 × seg_12  +0.033359
+                                 seg_2 × seg_11  -0.030257
+                                  seg_7 × seg_8  -0.022596
+                                  seg_1 × seg_7  +0.021836
+                                 seg_7 × seg_11  -0.021643
+                                 seg_9 × seg_10  +0.019139
+                                  seg_2 × seg_7  +0.018743
+                                  seg_5 × seg_7  -0.018698
+                                 seg_9 × seg_12  +0.018567
+
+Token×Token interactions (15):
+                              tok:a × tok:baked  +0.749371
+                               tok:a × tok:dish  +0.044180
+                              tok:baked × tok:a  -0.038298
+                           tok:baked × tok:dish  +0.027276
+                                 tok:a × tok:on  -0.018810
+                                 tok:by × tok:.  +0.017175
+                              tok:a × tok:woman  -0.016783
+                           tok:dish × tok:woman  +0.013511
+                             tok:being × tok:by  +0.012118
+                              tok:a × tok:being  +0.011113
+                                 tok:a × tok:by  +0.010688
+                          tok:plate × tok:being  +0.010307
+                          tok:baked × tok:woman  +0.010033
+                              tok:plate × tok:.  -0.009749
+                          tok:baked × tok:plate  -0.009563
+
+Top-5 Cross-Modal interactions (seg × tok):
+                                  seg_7 × tok:a  +0.059859
+                                 seg_11 × tok:a  -0.058050
+                                  seg_6 × tok:a  +0.053917
+                                  seg_5 × tok:a  +0.052764
+                                  seg_2 × tok:a  +0.032056
+
+★ Positive token #1: tok:a  (+2.176333)  × top-3 segments:
+                                  seg_7 × tok:a  +0.059859
+                                 seg_11 × tok:a  -0.058050
+                                  seg_6 × tok:a  +0.053917
+
+★ Positive token #2: tok:baked  (+0.402312)  × top-3 segments:
+                             seg_11 × tok:baked  +0.030477
+                             seg_12 × tok:baked  +0.011163
+                              seg_3 × tok:baked  +0.010498
+
+★ Negative token #1: tok:on  (-0.021476)  × top-3 segments:
+                                 seg_0 × tok:on  -0.010448
+                                seg_11 × tok:on  +0.008560
+                                 seg_5 × tok:on  +0.008531
+
+★ Negative token #2: tok:a  (-0.015789)  × top-3 segments:
+                                  seg_6 × tok:a  -0.010533
+                                  seg_3 × tok:a  +0.007261
+                                 seg_15 × tok:a  +0.006404
+
+============================================================
+Per-token cross-modal interactions (top-5 each):
+
+  tok:a  (Shapley=+2.176333):
+                                    seg_7 × tok:a  +0.059859
+                                   seg_11 × tok:a  -0.058050
+                                    seg_6 × tok:a  +0.053917
+                                    seg_5 × tok:a  +0.052764
+                                    seg_2 × tok:a  +0.032056
+
+  tok:baked  (Shapley=+0.402312):
+                               seg_11 × tok:baked  +0.030477
+                               seg_12 × tok:baked  +0.011163
+                                seg_3 × tok:baked  +0.010498
+                                seg_9 × tok:baked  +0.008587
+                               seg_13 × tok:baked  +0.008037
+
+  tok:dish  (Shapley=+0.027537):
+                                 seg_6 × tok:dish  +0.018977
+                                 seg_0 × tok:dish  -0.014123
+                                 seg_9 × tok:dish  +0.011310
+                                 seg_7 × tok:dish  -0.010339
+                                seg_14 × tok:dish  -0.008715
+
+  tok:on  (Shapley=-0.021476):
+                                   seg_0 × tok:on  -0.010448
+                                  seg_11 × tok:on  +0.008560
+                                   seg_5 × tok:on  +0.008531
+                                   seg_3 × tok:on  -0.008124
+                                   seg_4 × tok:on  -0.007208
+
+  tok:a  (Shapley=-0.015789):
+                                    seg_6 × tok:a  -0.010533
+                                    seg_3 × tok:a  +0.007261
+                                   seg_15 × tok:a  +0.006404
+                                    seg_0 × tok:a  -0.004502
+                                   seg_13 × tok:a  -0.004371
+
+  tok:plate  (Shapley=-0.000493):
+                                seg_9 × tok:plate  -0.005997
+                               seg_11 × tok:plate  +0.005358
+                                seg_1 × tok:plate  -0.004499
+                               seg_13 × tok:plate  -0.004154
+                                seg_0 × tok:plate  +0.003897
+
+  tok:being  (Shapley=+0.009256):
+                               seg_15 × tok:being  +0.010570
+                               seg_11 × tok:being  +0.009509
+                                seg_1 × tok:being  -0.008265
+                                seg_0 × tok:being  -0.007735
+                                seg_8 × tok:being  -0.006427
+
+  tok:touched  (Shapley=+0.009330):
+                              seg_9 × tok:touched  +0.009475
+                              seg_5 × tok:touched  +0.009189
+                             seg_11 × tok:touched  +0.009155
+                              seg_7 × tok:touched  -0.005751
+                              seg_1 × tok:touched  -0.004349
+
+  tok:by  (Shapley=+0.007656):
+                                  seg_15 × tok:by  +0.014428
+                                  seg_13 × tok:by  -0.006557
+                                   seg_9 × tok:by  -0.005631
+                                   seg_6 × tok:by  +0.005500
+                                   seg_1 × tok:by  +0.004970
+
+  tok:a  (Shapley=+0.018718):
+                                    seg_5 × tok:a  +0.016070
+                                   seg_11 × tok:a  +0.009449
+                                    seg_3 × tok:a  -0.008144
+                                    seg_7 × tok:a  -0.005613
+                                    seg_0 × tok:a  +0.003870
+
+  tok:woman  (Shapley=+0.013079):
+                                seg_8 × tok:woman  +0.018431
+                               seg_11 × tok:woman  +0.011556
+                                seg_3 × tok:woman  -0.008443
+                                seg_5 × tok:woman  +0.007948
+                               seg_14 × tok:woman  -0.007359
+
+  tok:.  (Shapley=+0.013208):
+                                    seg_7 × tok:.  -0.014178
+                                    seg_9 × tok:.  +0.008633
+                                   seg_15 × tok:.  +0.005227
+                                    seg_8 × tok:.  +0.004485
+                                   seg_13 × tok:.  +0.004019

results/coco_mm/coco_448256_summary.txt

@@ -0,0 +1,176 @@
+=== COCO 448256 – shapley CROSS-MODAL attributions ===
+
+Caption: Three men stand next to a car with its hood open.
+
+Image region values (16 regions):
+                      seg_1  -0.154752
+                      seg_3  +0.133889
+                      seg_7  +0.097503
+                     seg_12  +0.087315
+                      seg_9  +0.072469
+                      seg_6  +0.068672
+                      seg_2  -0.055121
+                      seg_0  +0.048846
+                     seg_10  +0.047620
+                     seg_14  -0.041853
+                     seg_13  +0.035891
+                      seg_8  -0.031498
+                     seg_15  -0.022439
+                      seg_5  -0.006506
+                     seg_11  -0.003508
+                      seg_4  +0.001179
+
+Caption token values (12 tokens):
+                  tok:three  +1.979697
+                    tok:men  +0.919549
+                     tok:to  +0.071413
+                    tok:car  +0.028613
+                  tok:stand  +0.028609
+                   tok:open  -0.025885
+                   tok:hood  +0.022758
+                   tok:with  -0.020605
+                      tok:a  +0.019752
+                   tok:next  +0.017291
+                      tok:.  -0.012653
+                    tok:its  +0.011887
+
+Image×Image interactions (15):
+                                  seg_2 × seg_4  +0.190486
+                                  seg_1 × seg_6  +0.102403
+                                  seg_5 × seg_8  +0.077659
+                                 seg_2 × seg_14  +0.074111
+                                 seg_8 × seg_15  -0.073519
+                                 seg_1 × seg_10  -0.066988
+                                  seg_2 × seg_6  +0.060852
+                                 seg_9 × seg_12  +0.054643
+                                  seg_3 × seg_5  +0.054639
+                                 seg_1 × seg_14  -0.054518
+                                  seg_3 × seg_4  +0.048830
+                                seg_12 × seg_14  +0.048132
+                                 seg_6 × seg_10  -0.047251
+                                  seg_1 × seg_2  +0.046458
+                                  seg_1 × seg_3  +0.045791
+
+Token×Token interactions (15):
+                            tok:three × tok:men  +1.803450
+                             tok:next × tok:car  +0.068161
+                            tok:stand × tok:car  +0.059719
+                             tok:next × tok:its  +0.054150
+                              tok:next × tok:to  +0.039551
+                               tok:open × tok:.  -0.037103
+                               tok:to × tok:its  -0.027991
+                          tok:three × tok:stand  -0.026706
+                              tok:three × tok:a  +0.025264
+                            tok:three × tok:its  +0.023843
+                             tok:stand × tok:to  +0.018563
+                           tok:three × tok:open  -0.018217
+                              tok:car × tok:its  +0.016328
+                             tok:car × tok:open  -0.015769
+                                tok:a × tok:car  +0.014476
+
+Top-5 Cross-Modal interactions (seg × tok):
+                              seg_8 × tok:three  +0.198546
+                              seg_5 × tok:three  +0.179486
+                             seg_10 × tok:three  -0.141457
+                              seg_1 × tok:three  +0.132898
+                             seg_14 × tok:three  +0.090481
+
+★ Positive token #1: tok:three  (+1.979697)  × top-3 segments:
+                              seg_8 × tok:three  +0.198546
+                              seg_5 × tok:three  +0.179486
+                             seg_10 × tok:three  -0.141457
+
+★ Positive token #2: tok:men  (+0.919549)  × top-3 segments:
+                               seg_11 × tok:men  +0.080520
+                                seg_8 × tok:men  -0.041843
+                                seg_3 × tok:men  -0.022314
+
+★ Negative token #1: tok:open  (-0.025885)  × top-3 segments:
+                               seg_7 × tok:open  +0.050209
+                               seg_8 × tok:open  +0.038709
+                              seg_14 × tok:open  +0.031120
+
+★ Negative token #2: tok:with  (-0.020605)  × top-3 segments:
+                               seg_2 × tok:with  -0.050356
+                               seg_9 × tok:with  -0.028777
+                               seg_7 × tok:with  -0.026861
+
+============================================================
+Per-token cross-modal interactions (top-5 each):
+
+  tok:three  (Shapley=+1.979697):
+                                seg_8 × tok:three  +0.198546
+                                seg_5 × tok:three  +0.179486
+                               seg_10 × tok:three  -0.141457
+                                seg_1 × tok:three  +0.132898
+                               seg_14 × tok:three  +0.090481
+
+  tok:men  (Shapley=+0.919549):
+                                 seg_11 × tok:men  +0.080520
+                                  seg_8 × tok:men  -0.041843
+                                  seg_3 × tok:men  -0.022314
+                                  seg_5 × tok:men  -0.016305
+                                  seg_7 × tok:men  -0.012643
+
+  tok:stand  (Shapley=+0.028609):
+                                seg_3 × tok:stand  +0.015101
+                                seg_7 × tok:stand  -0.009797
+                                seg_8 × tok:stand  -0.004832
+                               seg_15 × tok:stand  -0.000945
+                                seg_9 × tok:stand  +0.000807
+
+  tok:next  (Shapley=+0.017291):
+                                 seg_9 × tok:next  +0.031625
+                                seg_13 × tok:next  +0.004279
+                                seg_14 × tok:next  +0.000663
+
+  tok:to  (Shapley=+0.071413):
+                                   seg_2 × tok:to  +0.070109
+                                   seg_3 × tok:to  -0.038082
+                                   seg_5 × tok:to  +0.030661
+                                   seg_7 × tok:to  +0.026240
+                                  seg_10 × tok:to  +0.014946
+
+  tok:a  (Shapley=+0.019752):
+                                    seg_9 × tok:a  +0.016140
+                                   seg_13 × tok:a  +0.011941
+                                    seg_2 × tok:a  +0.005905
+                                   seg_10 × tok:a  -0.004813
+
+  tok:car  (Shapley=+0.028613):
+                                 seg_10 × tok:car  +0.031554
+                                  seg_0 × tok:car  +0.010387
+                                  seg_9 × tok:car  +0.007219
+
+  tok:with  (Shapley=-0.020605):
+                                 seg_2 × tok:with  -0.050356
+                                 seg_9 × tok:with  -0.028777
+                                 seg_7 × tok:with  -0.026861
+                                seg_11 × tok:with  -0.023846
+                                 seg_3 × tok:with  -0.018209
+
+  tok:its  (Shapley=+0.011887):
+                                  seg_0 × tok:its  +0.022133
+                                  seg_3 × tok:its  -0.015795
+                                  seg_1 × tok:its  -0.015457
+                                 seg_10 × tok:its  -0.007296
+                                  seg_8 × tok:its  -0.006122
+
+  tok:hood  (Shapley=+0.022758):
+                                 seg_9 × tok:hood  +0.042906
+                                seg_14 × tok:hood  -0.016583
+                                 seg_2 × tok:hood  -0.013351
+                                 seg_0 × tok:hood  -0.006368
+                                 seg_6 × tok:hood  -0.005324
+
+  tok:open  (Shapley=-0.025885):
+                                 seg_7 × tok:open  +0.050209
+                                 seg_8 × tok:open  +0.038709
+                                seg_14 × tok:open  +0.031120
+                                 seg_9 × tok:open  +0.009955
+                                seg_13 × tok:open  +0.007836
+
+  tok:.  (Shapley=-0.012653):
+                                   seg_10 × tok:.  +0.020455
+                                   seg_11 × tok:.  +0.010527
+                                    seg_2 × tok:.  +0.008770