experiment/evaluation/eval_nullu_val.py

"""
Val-split eval for Nullu (null-space weight editing) checkpoints:
  (1) keyword mention of object in captions — Nullu, optionally vs base
  (2) linear SAE probes on Nullu internals vs ground-truth labels
  (3) optional perplexity checks (OOD + in-domain)

Model loading: loads base model with AutoModelForPreTraining, then splices
Nullu down_proj edits for layers [--lowest_layer, --highest_layer).
Base comparison requires --compare_base; default is Nullu-only.
"""

from __future__ import annotations

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

import numpy as np
import torch
import torch.distributed as dist
from tqdm import tqdm
from transformers import AutoModelForPreTraining, AutoProcessor

sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../.."))

from experiment.config.relation_config import get_relation_config

DEFAULT_EVAL_PROMPTS = [
    "Describe this image.",
    "list all objects in this image",
]
from experiment.data.hf_loader import load_hf_dataset
from experiment.evaluation.inference import generate_text_batch
from experiment.evaluation.metrics import KeywordMentionDetector
from experiment.training.finetune_adv import (
    FrozenSAEEncoder,
    HiddenStateCapture,
    count_lm_layers,
    get_sae_features,
    layer_probes_from_checkpoint,
    probe_labels,
)

# ---------------------------------------------------------------------------
# Nullu weight loading helpers
# ---------------------------------------------------------------------------

_DOWN_PROJ_PREFIXES = (
    "model.layers",               # haotian-liu LLaVA-v1.5 (Nullu default)
    "language_model.model.layers", # older llava-hf transformers format
    "language_model.layers",       # newer llava-hf transformers format
)
_DOWN_PROJ_SUFFIX = ".mlp.down_proj.weight"


def _match_down_proj_layer(key: str) -> Optional[int]:
    if not key.endswith(_DOWN_PROJ_SUFFIX):
        return None
    for prefix in _DOWN_PROJ_PREFIXES:
        if key.startswith(prefix + "."):
            mid = key[len(prefix) + 1 : -len(_DOWN_PROJ_SUFFIX)]
            if mid.isdigit():
                return int(mid)
    return None


def _load_nullu_down_proj(checkpoint_path: str, layer_indices: List[int]) -> Dict[int, torch.Tensor]:
    """Read only mlp.down_proj.weight tensors for the requested layer indices."""
    from safetensors.torch import safe_open

    ckpt = Path(checkpoint_path)
    wanted = set(layer_indices)
    result: Dict[int, torch.Tensor] = {}

    shard_files = sorted(ckpt.glob("*.safetensors"))
    if shard_files:
        for shard in shard_files:
            with safe_open(shard, framework="pt") as f:
                for key in f.keys():
                    idx = _match_down_proj_layer(key)
                    if idx is not None and idx in wanted and idx not in result:
                        result[idx] = f.get_tensor(key)
            if len(result) == len(wanted):
                break
    else:
        bin_files = sorted(ckpt.glob("pytorch_model*.bin")) or sorted(ckpt.glob("*.bin"))
        if not bin_files:
            raise FileNotFoundError(f"No safetensors or .bin weight files found in {ckpt}")
        for bf in bin_files:
            sd = torch.load(bf, map_location="cpu", weights_only=True)
            for key, tensor in sd.items():
                idx = _match_down_proj_layer(key)
                if idx is not None and idx in wanted and idx not in result:
                    result[idx] = tensor
            if len(result) == len(wanted):
                break

    missing = wanted - set(result.keys())
    if missing:
        raise KeyError(f"Missing mlp.down_proj.weight for layers {sorted(missing)} in {ckpt}")
    return result


def _find_lm_layers(model) -> torch.nn.ModuleList:
    """Locate the decoder-block ModuleList, trying known attribute paths."""
    import torch.nn as nn

    if hasattr(model, "language_model"):
        lm = model.language_model
        if hasattr(lm, "model") and hasattr(lm.model, "layers"):
            return lm.model.layers
        if hasattr(lm, "layers"):
            return lm.layers
    if hasattr(model, "model"):
        inner = model.model
        if hasattr(inner, "language_model") and hasattr(inner.language_model, "layers"):
            return inner.language_model.layers
    for _, module in model.named_modules():
        if isinstance(module, nn.ModuleList) and len(module) > 0:
            if hasattr(module[0], "mlp") and hasattr(module[0].mlp, "down_proj"):
                return module
    raise AttributeError("Cannot locate decoder layers in model")


def _apply_nullu_weights(model, checkpoint_path: str, lowest_layer: int, highest_layer: int, device) -> int:
    """Apply Nullu down_proj edits in-place for layers [lowest_layer, highest_layer)."""
    layer_indices = list(range(lowest_layer, highest_layer))
    weights = _load_nullu_down_proj(checkpoint_path, layer_indices)
    layers = _find_lm_layers(model)
    for idx, w in weights.items():
        tgt = layers[idx].mlp.down_proj.weight
        tgt.data.copy_(w.to(device=tgt.device, dtype=tgt.dtype))
    return len(weights)


# ---------------------------------------------------------------------------
# Argument parsing
# ---------------------------------------------------------------------------

def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--relation", type=str, default="bathroom_toilet")
    p.add_argument("--base_model", type=str, default="llava-hf/llava-1.5-7b-hf")
    p.add_argument("--nullu_checkpoint", type=str, required=True,
                   help="Path to Nullu edited-model directory (HF checkpoint format).")
    p.add_argument("--lowest_layer", type=int, default=16,
                   help="Inclusive lower bound of edited layer range.")
    p.add_argument("--highest_layer", type=int, default=32,
                   help="Exclusive upper bound of edited layer range.")
    p.add_argument("--compare_base", action="store_true",
                   help="Also run base-model inference for comparison (two-pass; default: Nullu-only).")
    p.add_argument(
        "--mention_only",
        action="store_true",
        help="Only compute caption keyword-mention rates (skip probes + perplexity).",
    )
    p.add_argument("--sae_checkpoint", type=str, default="")
    p.add_argument("--probes_path", type=str, default="")
    p.add_argument("--probe_layers", type=str, default="", help="comma-separated; empty = all LM layers")
    p.add_argument("--probe_label_mode", type=str, choices=["union", "object_only"], default="object_only")
    p.add_argument("--max_samples", type=int, default=0, help="0 = full val split; >0 = max per category.")
    p.add_argument("--prompt", type=str, default=None,
                   help="Single prompt override (deprecated; use --prompts).")
    p.add_argument("--prompts", type=str, nargs="+", default=None,
                   help="Prompts to evaluate. Defaults to description + object-list prompts.")
    p.add_argument("--max_new_tokens", type=int, default=300)
    p.add_argument("--dtype", type=str, default="float16", choices=["float16", "bfloat16"])
    p.add_argument("--output_dir", type=str, default=None)
    p.add_argument("--ppl_dataset_id", type=str, default="lmms-lab/COCO-Caption",
                   help="OOD perplexity dataset. Set to empty to skip.")
    p.add_argument("--ppl_split", type=str, default="val")
    p.add_argument("--ppl_max_samples", type=int, default=0)
    p.add_argument("--ppl_seed", type=int, default=42)
    p.add_argument("--batch_size", type=int, default=1)
    p.add_argument("--probe_batch_size", type=int, default=0)
    p.add_argument("--ppl_batch_size", type=int, default=0)
    p.add_argument("--attn_impl", type=str, default="eager",
                   choices=["eager", "sdpa", "flash_attention_2"])
    return p.parse_args()


def _resolve_prompts(args) -> list:
    prompts = list(args.prompts) if args.prompts else ([args.prompt] if args.prompt else list(DEFAULT_EVAL_PROMPTS))
    seen: set = set()
    resolved = []
    for p in prompts:
        p = p.strip()
        if p and p not in seen:
            resolved.append(p)
            seen.add(p)
    return resolved or list(DEFAULT_EVAL_PROMPTS)


# ---------------------------------------------------------------------------
# Distributed helpers (unchanged from lora eval)
# ---------------------------------------------------------------------------

def _dtype(s: str):
    return torch.float16 if s == "float16" else torch.bfloat16


def _setup_dist():
    if "LOCAL_RANK" not in os.environ:
        return 0, 1, torch.device("cuda" if torch.cuda.is_available() else "cpu")
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)
    dist.init_process_group(backend="nccl")
    rank = dist.get_rank()
    world_size = dist.get_world_size()
    return rank, world_size, torch.device(f"cuda:{local_rank}")


def _gather_list(local: list, world_size: int) -> list:
    if world_size == 1:
        return local
    bucket = [None] * world_size
    dist.all_gather_object(bucket, local)
    out = []
    for part in bucket:
        out.extend(part)
    return out


# ---------------------------------------------------------------------------
# Metric helpers
# ---------------------------------------------------------------------------

def _four_category_masks(sc_arr: np.ndarray, ob_arr: np.ndarray):
    return (
        (sc_arr < 0.5) & (ob_arr > 0.5),
        (sc_arr > 0.5) & (ob_arr < 0.5),
        (sc_arr > 0.5) & (ob_arr > 0.5),
        (sc_arr < 0.5) & (ob_arr < 0.5),
    )


def _category_display_names(rc):
    return {
        rc.non_scene_with_object: f"{rc.object_key}_only",
        rc.scene_no_object: f"{rc.scene_key}_only",
        rc.scene_with_object: f"{rc.scene_key}_{rc.object_key}",
        "neither": "neither",
    }


def _cap_cat_stats(mask: np.ndarray, base_arr, nullu_arr: np.ndarray, ob_arr: np.ndarray):
    """base_arr may be None when --compare_base is not set."""
    if not mask.any():
        return None
    n = int(mask.sum())
    nu = nullu_arr[mask]
    has_obj = bool((ob_arr[mask] > 0.5).all())
    nullu_rate = float(nu.mean())
    error_type = "miss_rate" if has_obj else "hallu_rate"
    nullu_err = (1.0 - nullu_rate) if has_obj else nullu_rate
    out = {
        "n": n,
        "nullu_mention_count": int(nu.sum()),
        "nullu_mention_rate": nullu_rate,
        "nullu_" + error_type: nullu_err,
        "error_type": error_type,
    }
    if base_arr is not None:
        b = base_arr[mask]
        base_rate = float(b.mean())
        base_err = (1.0 - base_rate) if has_obj else base_rate
        out["base_mention_count"] = int(b.sum())
        out["base_mention_rate"] = base_rate
        out["base_" + error_type] = base_err
    return out


def _probe_cat_stats(mask: np.ndarray, scores_arr: np.ndarray, labels_arr: np.ndarray):
    if not mask.any():
        return None
    s = scores_arr[mask]
    l = labels_arr[mask]
    preds = (s > 0.5).astype(np.float64)
    tp = float(((preds == 1) & (l == 1)).sum())
    fp = float(((preds == 1) & (l == 0)).sum())
    fn = float(((preds == 0) & (l == 1)).sum())
    tn = float(((preds == 0) & (l == 0)).sum())
    acc = float((preds == l).mean())
    precision = tp / (tp + fp) if (tp + fp) > 0 else float("nan")
    recall = tp / (tp + fn) if (tp + fn) > 0 else float("nan")
    f1 = (
        (2 * precision * recall / (precision + recall))
        if (precision + recall) > 0 and not (np.isnan(precision) or np.isnan(recall))
        else float("nan")
    )
    auc = None
    try:
        from sklearn.metrics import roc_auc_score
        if len(np.unique(l)) > 1:
            auc = float(roc_auc_score(l, s))
    except ImportError:
        pass
    return {
        "n": int(mask.sum()),
        "label": float(l[0]) if len(l) else float("nan"),
        "avg_prob": float(s.mean()),
        "acc": acc,
        "precision": precision,
        "recall": recall,
        "f1": f1,
        "auc": auc,
        "tp": int(tp),
        "fp": int(fp),
        "fn": int(fn),
        "tn": int(tn),
    }


def _caption_nll_batch(model, processor, images: list, prompt: str, captions: list, device: str) -> list:
    import torch.nn as nn
    prompt_text = f"USER: <image>\n{prompt}\nASSISTANT:"
    full_texts = [f"USER: <image>\n{prompt}\nASSISTANT: {cap}" for cap in captions]
    enc_p = processor(images=images[0], text=prompt_text, return_tensors="pt")
    prompt_len = enc_p["input_ids"].shape[1]
    enc = processor(images=images, text=full_texts, return_tensors="pt", padding=True)
    inputs = {k: v.to(device) for k, v in enc.items()}
    B = inputs["input_ids"].shape[0]
    labels = inputs["input_ids"].clone()
    labels[:, :prompt_len] = -100
    pad_id = processor.tokenizer.pad_token_id
    if pad_id is not None:
        labels[labels == pad_id] = -100
    with torch.no_grad():
        out = model(**inputs, use_cache=False)
    shift_logits = out.logits[:, :-1].contiguous()
    shift_labels = labels[:, 1:].contiguous()
    loss_fct = nn.CrossEntropyLoss(reduction="none")
    tok_loss = loss_fct(
        shift_logits.view(-1, shift_logits.size(-1)),
        shift_labels.view(-1),
    ).view(B, -1)
    valid = (shift_labels != -100).float()
    per_nll = (tok_loss * valid).sum(dim=1) / valid.sum(dim=1).clamp(min=1)
    return per_nll.cpu().tolist()


def _load_ppl_samples(dataset_id: str, split: str, max_samples: int, seed: int) -> list:
    from datasets import load_dataset as _hf_load
    ds = _hf_load(dataset_id, split=split)
    cap_col = next((c for c in ("answer", "captions", "caption") if c in ds.column_names), None)
    if cap_col is None:
        raise ValueError(f"No caption column found in {dataset_id}. Columns: {ds.column_names}")
    n = len(ds) if max_samples <= 0 else min(max_samples, len(ds))
    rng = np.random.default_rng(seed)
    idx = sorted(rng.choice(len(ds), size=n, replace=False).tolist())
    samples = []
    for i in idx:
        row = ds[int(i)]
        img = row["image"].convert("RGB")
        caps = row[cap_col]
        if isinstance(caps, str):
            caps = [caps]
        else:
            caps = [c if isinstance(c, str) else (c.get("raw") or c.get("caption") or "") for c in caps]
            caps = [c for c in caps if c]
        if caps:
            samples.append({"image": img, "captions": caps})
    return samples


# ---------------------------------------------------------------------------
# OOD PPL helper (single model pass, returns per-sample NLLs)
# ---------------------------------------------------------------------------

def _run_ood_ppl(model, processor, ppl_samples: list, indices_local: list, args, device: str) -> list:
    nlls: list = []
    _ppl_bs = args.ppl_batch_size if args.ppl_batch_size > 0 else args.batch_size
    for _pb in range(0, len(indices_local), _ppl_bs):
        batch_j = indices_local[_pb: _pb + _ppl_bs]
        batch_images_multi = [ppl_samples[j]["image"] for j in batch_j]
        batch_captions_multi = [ppl_samples[j]["captions"] for j in batch_j]
        max_caps = max(len(c) for c in batch_captions_multi)
        nll_acc = [[] for _ in batch_j]
        for cap_slot in range(max_caps):
            slot_imgs, slot_caps, slot_idxs = [], [], []
            for k, caps in enumerate(batch_captions_multi):
                if cap_slot < len(caps):
                    slot_imgs.append(batch_images_multi[k])
                    slot_caps.append(caps[cap_slot])
                    slot_idxs.append(k)
            if not slot_imgs:
                continue
            slot_nlls = _caption_nll_batch(model, processor, slot_imgs, args.prompt, slot_caps, device)
            for k, nll in zip(slot_idxs, slot_nlls):
                nll_acc[k].append(nll)
        for k in range(len(batch_j)):
            nlls.append(float(np.mean(nll_acc[k])) if nll_acc[k] else float("nan"))
    return nlls


# ---------------------------------------------------------------------------
# In-domain PPL helper
# ---------------------------------------------------------------------------

def _run_indomain_ppl(model, processor, ds, indices: list, scene_col: str, obj_col: str, args, device: str):
    nlls: list = []
    scene_flags_out: list = []
    object_flags_out: list = []
    indices_out: list = []
    n_skipped = 0
    _ppl_bs = args.ppl_batch_size if args.ppl_batch_size > 0 else args.batch_size
    for _pb in range(0, len(indices), _ppl_bs):
        batch_idx = indices[_pb: _pb + _ppl_bs]
        batch_valid = []
        for i in batch_idx:
            row = ds[i]
            caption = row.get("caption", "") if hasattr(row, "get") else ""
            if not caption:
                n_skipped += 1
            else:
                batch_valid.append((i, row, caption))
        if not batch_valid:
            continue
        b_idxs = [t[0] for t in batch_valid]
        b_rows = [t[1] for t in batch_valid]
        b_caps = [t[2] for t in batch_valid]
        b_imgs = [r["image"].convert("RGB") for r in b_rows]
        batch_nlls = _caption_nll_batch(model, processor, b_imgs, args.prompt, b_caps, device)
        nlls.extend(batch_nlls)
        scene_flags_out.extend([int(r[scene_col]) for r in b_rows])
        object_flags_out.extend([int(r[obj_col]) for r in b_rows])
        indices_out.extend(b_idxs)
    return nlls, scene_flags_out, object_flags_out, indices_out, n_skipped


# ---------------------------------------------------------------------------
# main
# ---------------------------------------------------------------------------

def main():
    args = parse_args()
    rank, world_size, device = _setup_dist()
    is_main = rank == 0

    if args.mention_only:
        args.ppl_dataset_id = ""
    else:
        if not args.sae_checkpoint:
            raise SystemExit("--sae_checkpoint is required unless --mention_only is set")
        if not args.probes_path:
            raise SystemExit("--probes_path is required unless --mention_only is set")

    prompts = _resolve_prompts(args)
    rc = get_relation_config(args.relation)
    dt = _dtype(args.dtype)
    scene_col, obj_col = rc.scene_key, rc.object_key

    cat_obj_only = rc.non_scene_with_object
    cat_scene_only = rc.scene_no_object
    cat_both = rc.scene_with_object
    cat_neither = "neither"
    cat_order = [cat_obj_only, cat_scene_only, cat_both, cat_neither]
    cat_display = _category_display_names(rc)

    ds = load_hf_dataset(rc.dataset_id, split="val")

    # Build sample index
    if args.max_samples <= 0:
        all_indices = list(range(len(ds)))
    else:
        sc_labels = ds[scene_col]
        ob_labels = ds[obj_col]
        cat_buckets: dict = {cat_obj_only: [], cat_scene_only: [], cat_both: [], cat_neither: []}
        for i, (sc_v, ob_v) in enumerate(zip(sc_labels, ob_labels)):
            if int(sc_v) == 0 and int(ob_v) == 1:
                cat_buckets[cat_obj_only].append(i)
            elif int(sc_v) == 1 and int(ob_v) == 0:
                cat_buckets[cat_scene_only].append(i)
            elif int(sc_v) == 1 and int(ob_v) == 1:
                cat_buckets[cat_both].append(i)
            else:
                cat_buckets[cat_neither].append(i)
        all_indices = []
        for bucket in cat_buckets.values():
            all_indices.extend(bucket[: args.max_samples])
        all_indices.sort()
    n = len(all_indices)

    out_dir = args.output_dir or os.path.join(args.nullu_checkpoint, "eval")
    if is_main:
        os.makedirs(out_dir, exist_ok=True)

    processor = AutoProcessor.from_pretrained(args.base_model)
    model = AutoModelForPreTraining.from_pretrained(
        args.base_model, torch_dtype=dt, attn_implementation=args.attn_impl
    ).to(device)
    model.eval()

    kw = KeywordMentionDetector(keywords=rc.mention_keywords)
    indices = [all_indices[i] for i in range(rank, n, world_size)]
    _gen_bs = args.batch_size

    # ------------------------------------------------------------------
    # (0) Optionally run base-model captions + PPL before editing weights
    # ------------------------------------------------------------------
    base_rates: list = []
    base_captions: list = []
    base_ood_nlls: list = []
    base_indomain_nlls: list = []
    base_indomain_scene: list = []
    base_indomain_object: list = []
    base_indomain_indices: list = []

    if args.compare_base:
        if is_main:
            print("=== (0) Base-model caption eval (before Nullu edit) ===")
        base_it = tqdm(total=len(indices) * len(prompts), desc="(0) base captions", disable=not is_main, unit="eval", dynamic_ncols=True)
        for _bs in range(0, len(indices), _gen_bs):
            batch_idx = indices[_bs: _bs + _gen_bs]
            batch_rows = [ds[i] for i in batch_idx]
            batch_images = [r["image"].convert("RGB") for r in batch_rows]
            for prompt in prompts:
                b_texts = generate_text_batch(model, processor, batch_images, prompt, str(device), args.max_new_tokens)
                for t in b_texts:
                    base_rates.append(float(kw.mentions_object(t)))
                    base_captions.append(t)
                if base_it is not None:
                    base_it.update(len(batch_idx))
        if base_it is not None:
            base_it.close()
        base_rates = _gather_list(base_rates, world_size)
        base_captions = _gather_list(base_captions, world_size)

        if not args.mention_only:
            # Base OOD PPL
            if args.ppl_dataset_id:
                ppl_samples = _load_ppl_samples(args.ppl_dataset_id, args.ppl_split, args.ppl_max_samples, args.ppl_seed)
                ppl_indices_local = list(range(rank, len(ppl_samples), world_size))
                _base_ood = _run_ood_ppl(model, processor, ppl_samples, ppl_indices_local, args, str(device))
                base_ood_nlls = _gather_list(_base_ood, world_size)

            # Base in-domain PPL
            _b_id_nlls, _b_id_sc, _b_id_ob, _b_id_idx, _ = _run_indomain_ppl(
                model, processor, ds, indices, scene_col, obj_col, args, str(device)
            )
            base_indomain_nlls = _gather_list(_b_id_nlls, world_size)
            base_indomain_scene = _gather_list(_b_id_sc, world_size)
            base_indomain_object = _gather_list(_b_id_ob, world_size)
            base_indomain_indices = _gather_list(_b_id_idx, world_size)

    # ------------------------------------------------------------------
    # Apply Nullu weight edits
    # ------------------------------------------------------------------
    n_edited = _apply_nullu_weights(model, args.nullu_checkpoint, args.lowest_layer, args.highest_layer, device)
    if is_main:
        print(
            f"\nApplied Nullu: {n_edited} layers edited  "
            f"[{args.lowest_layer}, {args.highest_layer})  "
            f"checkpoint: {args.nullu_checkpoint}"
        )

    # ------------------------------------------------------------------
    # (1) Nullu caption eval
    # ------------------------------------------------------------------
    if is_main:
        print("\n=== (1) Caption keyword eval: object mention (negation-aware) ===")

    nullu_rates: list = []
    nullu_captions: list = []
    gt_has_object: list = []
    cap_indices: list = []
    cap_scene_flags: list = []
    cap_image_ids: list = []
    cap_prompts: list = []

    _live_f = open(os.path.join(out_dir, "samples_live.jsonl"), "w") if is_main else None
    cap_it = tqdm(total=len(indices) * len(prompts), desc="(1) Nullu captions", disable=not is_main, unit="eval", dynamic_ncols=True)

    for _bs in range(0, len(indices), _gen_bs):
        batch_idx = indices[_bs: _bs + _gen_bs]
        batch_rows = [ds[i] for i in batch_idx]
        batch_images = [r["image"].convert("RGB") for r in batch_rows]

        for prompt in prompts:
            nu_texts = generate_text_batch(model, processor, batch_images, prompt, str(device), args.max_new_tokens)

            for k, (t, i, row) in enumerate(zip(nu_texts, batch_idx, batch_rows)):
                nu_m = float(kw.mentions_object(t))
                nullu_rates.append(nu_m)
                nullu_captions.append(t)
                gt_has_object.append(float(int(row[obj_col])))
                cap_indices.append(i)
                cap_scene_flags.append(int(row[scene_col]))
                cap_image_ids.append(row.get("image_id") if hasattr(row, "get") else None)
                cap_prompts.append(prompt)
                if _live_f is not None:
                    rec = {
                        "index": i,
                        "prompt": prompt,
                        "image_id": row.get("image_id") if hasattr(row, "get") else None,
                        scene_col: int(row[scene_col]),
                        obj_col: int(row[obj_col]),
                        "nullu_caption": t,
                        "nullu_mentions_object": bool(nu_m > 0.5),
                    }
                    if args.compare_base:
                        rec["compare_base"] = True
                    _live_f.write(json.dumps(rec) + "\n")
            if _live_f is not None:
                _live_f.flush()
            if cap_it is not None:
                cap_it.update(len(batch_idx))
    if cap_it is not None:
        cap_it.close()
    if _live_f is not None:
        _live_f.close()

    nullu_rates = _gather_list(nullu_rates, world_size)
    nullu_captions = _gather_list(nullu_captions, world_size)
    gt_has_object = _gather_list(gt_has_object, world_size)
    cap_indices = _gather_list(cap_indices, world_size)
    cap_scene_flags = _gather_list(cap_scene_flags, world_size)
    cap_image_ids = _gather_list(cap_image_ids, world_size)
    cap_prompts = _gather_list(cap_prompts, world_size)

    cap_metrics: dict = {}
    cap_metrics_by_prompt: dict = {}
    if is_main:
        b_arr = np.array(base_rates, dtype=np.float64) if base_rates else None
        nu_arr = np.array(nullu_rates, dtype=np.float64)
        ho_arr = np.array(gt_has_object, dtype=np.float64)
        sc_cap = np.array(cap_scene_flags, dtype=np.float64)
        cap_prompt_arr = np.array(cap_prompts, dtype=object)
        n_evals = len(nu_arr)
        m_to, m_bo, m_bt, m_ne = _four_category_masks(sc_cap, ho_arr)

        cap_cats = {
            cat_obj_only: _cap_cat_stats(m_to, b_arr, nu_arr, ho_arr),
            cat_scene_only: _cap_cat_stats(m_bo, b_arr, nu_arr, ho_arr),
            cat_both: _cap_cat_stats(m_bt, b_arr, nu_arr, ho_arr),
            cat_neither: _cap_cat_stats(m_ne, b_arr, nu_arr, ho_arr),
        }

        _has_base = b_arr is not None
        print(f"  images: {n}  evals: {n_evals}  GPUs: {world_size}  layers_edited: [{args.lowest_layer}, {args.highest_layer})")
        print(f"  prompts: {prompts!r}")
        print(f"  keywords: {rc.mention_keywords[:3]!r}... (negation-aware)")
        if not _has_base:
            print("  (Nullu-only: base columns omitted — use --compare_base to add base)")
        print()
        print("  Aggregate across prompts:")
        _sep = "-" * (91 if _has_base else 72)
        _hdr = f"  {'Base':>14}  " if _has_base else ""
        print(f"  {'Category':<24} {'N':>5}  {_hdr}{'Nullu':>14}  Error")
        print("  " + _sep)
        for cat in cat_order:
            d = cap_cats.get(cat)
            if d is None:
                print(f"  {cat_display.get(cat, cat):<18} {'(empty)'}")
                continue
            display_name = cat_display.get(cat, cat)
            nr = d["n"]
            nm = d["nullu_mention_count"]
            nr_rate = d["nullu_mention_rate"]
            et = d["error_type"]
            ne = d.get("nullu_" + et, float("nan"))
            _base_col = (
                f"{d['base_mention_count']:>3}/{nr:<4}({d['base_mention_rate']:>6.1%})  "
                if _has_base else ""
            )
            _be_s = f"base={d.get('base_' + et, float('nan')):.1%} " if _has_base else ""
            print(
                f"  {display_name:<24} {nr:>5}  "
                f"{_base_col}"
                f"{nm:>3}/{nr:<4}({nr_rate:>6.1%})  "
                f"{et}: {_be_s}nullu={ne:.1%}"
            )
        print("  " + _sep)
        _b_overall = f"{int(b_arr.sum()):>3}/{n_evals:<4}({b_arr.mean():>6.1%})  " if _has_base else ""
        print(
            f"  {'OVERALL':<24} {n_evals:>5}  "
            f"{_b_overall}"
            f"{int(nu_arr.sum()):>3}/{n_evals:<4}({nu_arr.mean():>6.1%})"
        )
        if cap_cats[cat_scene_only] is not None:
            bd = cap_cats[cat_scene_only]
            _sup_base = f"base hallu={bd['base_hallu_rate']:.1%}  " if _has_base else ""
            _delta = (
                f"  Δ={bd['base_hallu_rate'] - bd['nullu_hallu_rate']:+.1%}"
                if _has_base else ""
            )
            print(
                f"\n  [Suppression] {cat_display.get(cat_scene_only, cat_scene_only)} (D_{{A,¬B}}): "
                f"{_sup_base}"
                f"Nullu hallu={bd['nullu_hallu_rate']:.1%}"
                f"{_delta}"
            )

        print("\n  Per-prompt overall:")
        for prompt in prompts:
            pmask = cap_prompt_arr == prompt
            p_b = b_arr[pmask] if _has_base else None
            p_nu = nu_arr[pmask]
            p_ho = ho_arr[pmask]
            p_sc = sc_cap[pmask]
            pm_to, pm_bo, pm_bt, pm_ne = _four_category_masks(p_sc, p_ho)
            p_cats = {
                cat_obj_only: _cap_cat_stats(pm_to, p_b, p_nu, p_ho),
                cat_scene_only: _cap_cat_stats(pm_bo, p_b, p_nu, p_ho),
                cat_both: _cap_cat_stats(pm_bt, p_b, p_nu, p_ho),
                cat_neither: _cap_cat_stats(pm_ne, p_b, p_nu, p_ho),
            }
            cap_metrics_by_prompt[prompt] = {
                "overall": {"nullu_mention_rate": float(p_nu.mean())},
                "categories": p_cats,
            }
            if _has_base:
                cap_metrics_by_prompt[prompt]["overall"]["base_mention_rate"] = float(p_b.mean())
                print(f"  {prompt!r}: base={p_b.mean():.1%}  nullu={p_nu.mean():.1%}")
            else:
                print(f"  {prompt!r}: nullu={p_nu.mean():.1%}")

        cap_metrics = cap_cats

    # ------------------------------------------------------------------
    # (2) SAE probe eval (Nullu model)
    # ------------------------------------------------------------------
    pl: list = []
    scores: list = []
    labels: list = []
    scene_flags: list = []
    object_flags: list = []
    probe_indices: list = []
    probe_metrics: dict = {}
    overall_acc = None
    overall_auc = None
    overall_ap = None
    ppl_metrics: dict = {}
    indomain_ppl_metrics: dict = {}
    indomain_indices: list = []
    indomain_nullu_nlls_outer: list = []
    indomain_base_nlls_outer: list = []

    if not args.mention_only:
        if dist.is_initialized():
            dist.barrier()
        if is_main:
            print("\n=== (2) SAE probe eval (object_only labels) ===")

        sae = FrozenSAEEncoder.from_checkpoint(args.sae_checkpoint, device)
        d_sae = sae.encoder.weight.shape[0]
        n_layers = count_lm_layers(model)
        pl = (
            [int(x) for x in args.probe_layers.split(",") if x.strip()]
            if args.probe_layers.strip()
            else list(range(n_layers))
        )
        probes = layer_probes_from_checkpoint(args.probes_path, pl, d_sae, device=device)
        probes.eval()

        capture = HiddenStateCapture(model, pl)
        _probe_bs = args.probe_batch_size if args.probe_batch_size > 0 else args.batch_size
        _probe_text = f"USER: <image>\n{args.prompt}\nASSISTANT:"
        with torch.no_grad():
            probe_it = tqdm(total=len(indices), desc="(2) SAE probe forward", disable=not is_main, unit="img", dynamic_ncols=True)
            for _pb in range(0, len(indices), _probe_bs):
                batch_idx = indices[_pb: _pb + _probe_bs]
                batch_rows = [ds[i] for i in batch_idx]
                batch_images = [r["image"].convert("RGB") for r in batch_rows]

                sc_t = torch.tensor([int(r[scene_col]) for r in batch_rows], device=device)
                ob_t = torch.tensor([int(r[obj_col]) for r in batch_rows], device=device)
                y_batch = probe_labels(sc_t, ob_t, args.probe_label_mode).tolist()

                probe_indices.extend(batch_idx)
                scene_flags.extend(sc_t.tolist())
                object_flags.extend(ob_t.tolist())
                labels.extend(y_batch)

                _orig_side = processor.tokenizer.padding_side
                processor.tokenizer.padding_side = "left"
                try:
                    batch_inputs = processor(
                        images=batch_images,
                        text=[_probe_text] * len(batch_images),
                        return_tensors="pt",
                        padding=True,
                    )
                finally:
                    processor.tokenizer.padding_side = _orig_side
                batch_inputs = {k: v.to(device) for k, v in batch_inputs.items()}

                with capture:
                    model(**batch_inputs, use_cache=False)
                feats = get_sae_features(capture, sae)
                probs_list = probes(feats)
                mean_p = torch.stack([p.float() for p in probs_list]).mean(dim=0)
                scores.extend(mean_p.tolist())

                if probe_it is not None:
                    probe_it.update(len(batch_idx))
            if probe_it is not None:
                probe_it.close()

        scores = _gather_list(scores, world_size)
        labels = _gather_list(labels, world_size)
        scene_flags = _gather_list(scene_flags, world_size)
        object_flags = _gather_list(object_flags, world_size)
        probe_indices = _gather_list(probe_indices, world_size)

        if is_main:
            scores_arr = np.array(scores, dtype=np.float64)
            labels_arr = np.array(labels, dtype=np.float64)
            sc_arr = np.array(scene_flags, dtype=np.float64)
            ob_arr = np.array(object_flags, dtype=np.float64)
            m_to, m_bo, m_bt, m_ne = _four_category_masks(sc_arr, ob_arr)
            probe_cats = {
                cat_obj_only: _probe_cat_stats(m_to, scores_arr, labels_arr),
                cat_scene_only: _probe_cat_stats(m_bo, scores_arr, labels_arr),
                cat_both: _probe_cat_stats(m_bt, scores_arr, labels_arr),
                cat_neither: _probe_cat_stats(m_ne, scores_arr, labels_arr),
            }
            print(f"  probe_label_mode: {args.probe_label_mode}  layers: {len(pl)}")
            print()
            print("  Per-category breakdown  (label: 1=probe fires, 0=probe silent)")
            print(
                f"  {'Category':<24} {'N':>5}  {'Lbl':>3}  "
                f"{'AvgProb':>7}  {'Acc@0.5':>7}  {'Prec':>6}  {'Rec':>6}  {'F1':>6}  {'AUC':>6}"
            )
            print("  " + "-" * 102)
            for cat in cat_order:
                d = probe_cats.get(cat)
                if d is None:
                    print(f"  {cat_display.get(cat, cat):<24} {'(empty)'}")
                    continue
                auc_s = f"{d['auc']:.4f}" if d["auc"] is not None else "   n/a"
                prec_s = f"{d['precision']:.4f}" if not np.isnan(d["precision"]) else "   n/a"
                rec_s = f"{d['recall']:.4f}" if not np.isnan(d["recall"]) else "   n/a"
                f1_s = f"{d['f1']:.4f}" if not np.isnan(d["f1"]) else "   n/a"
                print(
                    f"  {cat_display.get(cat, cat):<24} {d['n']:>5}  {d['label']:>3.0f}  "
                    f"{d['avg_prob']:>7.4f}  {d['acc']:>7.4f}  "
                    f"{prec_s:>6}  {rec_s:>6}  {f1_s:>6}  {auc_s:>6}"
                )
            all_preds = (scores_arr > 0.5).astype(np.float64)
            overall_acc = float((all_preds == labels_arr).mean())
            try:
                from sklearn.metrics import roc_auc_score, average_precision_score
                if len(np.unique(labels_arr)) > 1:
                    overall_auc = float(roc_auc_score(labels_arr, scores_arr))
                    overall_ap = float(average_precision_score(labels_arr, scores_arr))
            except ImportError:
                pass
            probe_metrics = probe_cats

        if dist.is_initialized():
            dist.barrier()

        # (3) OOD perplexity
        if not args.ppl_dataset_id:
            if is_main:
                print("\n=== (3) Out-of-domain perplexity: SKIPPED (--ppl_dataset_id is empty) ===")
        else:
            if is_main:
                print(f"\n=== (3) Out-of-domain perplexity  [{args.ppl_dataset_id} / {args.ppl_split}] ===")
            ppl_samples = _load_ppl_samples(args.ppl_dataset_id, args.ppl_split, args.ppl_max_samples, args.ppl_seed)
            ppl_indices_local = list(range(rank, len(ppl_samples), world_size))

            nullu_ood_nlls: list = []
            ppl_it = tqdm(total=len(ppl_indices_local), desc="(3) OOD PPL Nullu", disable=not is_main, unit="img", dynamic_ncols=True)
            _tmp = _run_ood_ppl(model, processor, ppl_samples, ppl_indices_local, args, str(device))
            if ppl_it is not None:
                ppl_it.update(len(ppl_indices_local))
                ppl_it.close()
            nullu_ood_nlls = _gather_list(_tmp, world_size)
            base_ood_nlls_gathered = _gather_list(base_ood_nlls, world_size) if base_ood_nlls else []

            if is_main and nullu_ood_nlls:
                ln = np.array(nullu_ood_nlls)
                nullu_ppl = float(np.exp(ln.mean()))
                if base_ood_nlls_gathered:
                    bn = np.array(base_ood_nlls_gathered)
                    base_ppl = float(np.exp(bn.mean()))
                    ratio = nullu_ppl / base_ppl
                    print(f"  n={len(ln)}  Base PPL={base_ppl:.3f}  Nullu PPL={nullu_ppl:.3f}  Ratio={ratio:.4f}")
                    ppl_metrics = {
                        "dataset": args.ppl_dataset_id,
                        "split": args.ppl_split,
                        "n_samples": len(ln),
                        "base_ppl": base_ppl,
                        "nullu_ppl": nullu_ppl,
                        "ppl_ratio": ratio,
                    }
                else:
                    print(f"  n={len(ln)}  Nullu PPL={nullu_ppl:.3f}  (base skipped)")
                    ppl_metrics = {
                        "dataset": args.ppl_dataset_id,
                        "split": args.ppl_split,
                        "n_samples": len(ln),
                        "nullu_ppl": nullu_ppl,
                    }

        if dist.is_initialized():
            dist.barrier()

        # (3b) In-domain perplexity by category
        if is_main:
            print("\n=== (3b) In-domain perplexity by category (relation val set) ===")

        _nu_id_nlls, _nu_id_sc, _nu_id_ob, _nu_id_idx, n_skip_local = _run_indomain_ppl(
            model, processor, ds, indices, scene_col, obj_col, args, str(device)
        )
        nullu_id_nlls = _gather_list(_nu_id_nlls, world_size)
        nullu_id_sc = _gather_list(_nu_id_sc, world_size)
        nullu_id_ob = _gather_list(_nu_id_ob, world_size)
        nullu_id_idx = _gather_list(_nu_id_idx, world_size)
        n_skip_total = sum(_gather_list([n_skip_local], world_size))

        # Use base in-domain data if available (already gathered above)
        base_id_nlls_arr = np.array(base_indomain_nlls) if base_indomain_nlls else None

        indomain_indices = nullu_id_idx
        indomain_nullu_nlls_outer = nullu_id_nlls
        indomain_base_nlls_outer = base_indomain_nlls

        if is_main and nullu_id_nlls:
            ln_id = np.array(nullu_id_nlls)
            sc_id = np.array(nullu_id_sc, dtype=np.float64)
            ob_id = np.array(nullu_id_ob, dtype=np.float64)
            m_to_id, m_bo_id, m_bt_id, m_ne_id = _four_category_masks(sc_id, ob_id)

            # Align base in-domain NLLs to same sample order if available
            base_id_aligned = None
            if base_indomain_nlls and base_indomain_indices:
                # Build a map from index → base NLL for lookup
                base_idx_to_nll = {idx: nll for idx, nll in zip(base_indomain_indices, base_indomain_nlls)}
                base_id_aligned = np.array([base_idx_to_nll.get(idx, float("nan")) for idx in nullu_id_idx])

            indomain_cats: dict = {}
            for cat, mask in [
                (cat_obj_only, m_to_id),
                (cat_scene_only, m_bo_id),
                (cat_both, m_bt_id),
                (cat_neither, m_ne_id),
            ]:
                if not mask.any():
                    indomain_cats[cat] = None
                    continue
                l_ppl = float(np.exp(ln_id[mask].mean()))
                entry: dict = {"n": int(mask.sum()), "nullu_ppl": l_ppl}
                if base_id_aligned is not None:
                    valid_mask = mask & ~np.isnan(base_id_aligned)
                    if valid_mask.any():
                        b_ppl = float(np.exp(base_id_aligned[valid_mask].mean()))
                        entry["base_ppl"] = b_ppl
                        entry["ppl_ratio"] = l_ppl / b_ppl
                indomain_cats[cat] = entry

            l_overall = float(np.exp(ln_id.mean()))
            overall_entry: dict = {"nullu_ppl": l_overall}
            if base_id_aligned is not None:
                valid = ~np.isnan(base_id_aligned)
                if valid.any():
                    b_overall = float(np.exp(base_id_aligned[valid].mean()))
                    overall_entry["base_ppl"] = b_overall
                    overall_entry["ppl_ratio"] = l_overall / b_overall
            indomain_ppl_metrics = {
                "n_samples": len(ln_id),
                "n_skipped_no_caption": n_skip_total,
                "overall": overall_entry,
                "categories": indomain_cats,
            }

    # ------------------------------------------------------------------
    # Save outputs
    # ------------------------------------------------------------------
    if is_main:
        cap_order = sorted(range(len(cap_indices)), key=lambda j: (cap_indices[j], cap_prompts[j]))
        nu_arr2 = np.array(nullu_rates, dtype=np.float64)
        ho_arr2 = np.array(gt_has_object, dtype=np.float64)
        sc_cap2 = np.array(cap_scene_flags, dtype=np.float64)
        b_arr2 = np.array(base_rates, dtype=np.float64) if base_rates else None

        def _sample_cat(sc_v, ob_v):
            if sc_v < 0.5 and ob_v > 0.5:
                return cat_obj_only
            if sc_v > 0.5 and ob_v < 0.5:
                return cat_scene_only
            if sc_v > 0.5 and ob_v > 0.5:
                return cat_both
            return cat_neither

        captions_records = []
        for j in cap_order:
            cat = _sample_cat(sc_cap2[j], ho_arr2[j])
            rec = {
                "index": int(cap_indices[j]),
                "prompt": cap_prompts[j],
                "image_id": cap_image_ids[j],
                scene_col: int(sc_cap2[j]),
                obj_col: int(ho_arr2[j]),
                "category": cat_display.get(cat, cat),
                "nullu_caption": nullu_captions[j],
                "nullu_mentions_object": bool(nu_arr2[j] > 0.5),
            }
            if b_arr2 is not None:
                rec["base_caption"] = base_captions[j]
                rec["base_mentions_object"] = bool(b_arr2[j] > 0.5)
            captions_records.append(rec)
        captions_path = os.path.join(out_dir, "captions.json")
        with open(captions_path, "w") as f:
            json.dump(captions_records, f, indent=2)
        print(f"\n  Captions saved → {captions_path}")

        def _safe_dict(d):
            if d is None:
                return None
            return {k: (None if (isinstance(v, float) and np.isnan(v)) else v) for k, v in d.items()}

        nu_f = np.array(nullu_rates, dtype=np.float64)
        caption_overall: dict = {"nullu_mention_rate": float(nu_f.mean())}
        if args.compare_base and base_rates:
            b_f = np.array(base_rates, dtype=np.float64)
            caption_overall["base_mention_rate"] = float(b_f.mean())

        metrics = {
            "relation": args.relation,
            "nullu_checkpoint": args.nullu_checkpoint,
            "lowest_layer": args.lowest_layer,
            "highest_layer": args.highest_layer,
            "compare_base": args.compare_base,
            "n_images": n,
            "n_prompt_evals": len(nu_f),
            "prompts": prompts,
            "caption_eval": {
                "overall": caption_overall,
                "categories": {
                    cat_obj_only: _safe_dict(cap_metrics.get(cat_obj_only)),
                    cat_scene_only: _safe_dict(cap_metrics.get(cat_scene_only)),
                    cat_both: _safe_dict(cap_metrics.get(cat_both)),
                    cat_neither: _safe_dict(cap_metrics.get(cat_neither)),
                },
                "per_prompt": {
                    prompt: {
                        "overall": _safe_dict(cap_metrics_by_prompt.get(prompt, {}).get("overall")),
                        "categories": {
                            cat_obj_only: _safe_dict((cap_metrics_by_prompt.get(prompt, {}).get("categories") or {}).get(cat_obj_only)),
                            cat_scene_only: _safe_dict((cap_metrics_by_prompt.get(prompt, {}).get("categories") or {}).get(cat_scene_only)),
                            cat_both: _safe_dict((cap_metrics_by_prompt.get(prompt, {}).get("categories") or {}).get(cat_both)),
                            cat_neither: _safe_dict((cap_metrics_by_prompt.get(prompt, {}).get("categories") or {}).get(cat_neither)),
                        },
                    }
                    for prompt in prompts
                },
            },
        }
        if len(prompts) == 1:
            metrics["prompt"] = prompts[0]

        scores_arr2 = np.array(scores, dtype=np.float64)
        labels_arr2 = np.array(labels, dtype=np.float64)

        if not args.mention_only:
            metrics["probe_label_mode"] = args.probe_label_mode
            metrics["n_probe_layers"] = len(pl)
            metrics["probe_eval"] = {
                "overall": {
                    "acc": float(overall_acc) if overall_acc is not None else None,
                    "roc_auc": overall_auc,
                    "average_precision": overall_ap,
                },
                "categories": {
                    cat_obj_only: _safe_dict(probe_metrics.get(cat_obj_only)),
                    cat_scene_only: _safe_dict(probe_metrics.get(cat_scene_only)),
                    cat_both: _safe_dict(probe_metrics.get(cat_both)),
                    cat_neither: _safe_dict(probe_metrics.get(cat_neither)),
                },
            }
            metrics["perplexity_ood"] = ppl_metrics
            metrics["perplexity_indomain"] = indomain_ppl_metrics

        metrics_path = os.path.join(out_dir, "metrics.json")
        with open(metrics_path, "w") as f:
            json.dump(metrics, f, indent=2)
        print(f"  Metrics saved  → {metrics_path}")

        # Build samples.json — one record per image with prompt_results nested
        samples_by_idx: dict = {}
        for j in cap_order:
            idx = int(cap_indices[j])
            sc_v = float(sc_cap2[j])
            ob_v = float(ho_arr2[j])
            cat = _sample_cat(sc_v, ob_v)
            rec = samples_by_idx.setdefault(
                idx,
                {
                    "index": idx,
                    "image_id": cap_image_ids[j],
                    scene_col: int(sc_v),
                    obj_col: int(ob_v),
                    "category": cat_display.get(cat, cat),
                    "prompt_results": {},
                },
            )
            prompt_rec = rec["prompt_results"].setdefault(cap_prompts[j], {})
            prompt_rec["nullu_caption"] = nullu_captions[j]
            prompt_rec["nullu_mentions_object"] = bool(nu_arr2[j] > 0.5)
            if b_arr2 is not None:
                prompt_rec["base_caption"] = base_captions[j]
                prompt_rec["base_mentions_object"] = bool(b_arr2[j] > 0.5)
            if len(prompts) == 1:
                rec["nullu_caption"] = prompt_rec["nullu_caption"]
                rec["nullu_mentions_object"] = prompt_rec["nullu_mentions_object"]
                if b_arr2 is not None:
                    rec["base_caption"] = prompt_rec["base_caption"]
                    rec["base_mentions_object"] = prompt_rec["base_mentions_object"]

        if not args.mention_only and len(probe_indices) > 0:
            for j, idx in enumerate(probe_indices):
                if idx in samples_by_idx:
                    samples_by_idx[idx]["probe_score"] = float(scores_arr2[j])
                    samples_by_idx[idx]["probe_label"] = float(labels_arr2[j])
                    samples_by_idx[idx]["probe_pred"] = int(scores_arr2[j] > 0.5)

        if indomain_indices:
            nu_id_arr = np.array(indomain_nullu_nlls_outer)
            for j, idx in enumerate(indomain_indices):
                if idx in samples_by_idx:
                    samples_by_idx[idx]["indomain_nullu_ppl"] = float(np.exp(nu_id_arr[j]))
            if indomain_base_nlls_outer and base_indomain_indices:
                base_idx_nll_map = {i: nll for i, nll in zip(base_indomain_indices, indomain_base_nlls_outer)}
                for idx in samples_by_idx:
                    if idx in base_idx_nll_map:
                        b_nll = base_idx_nll_map[idx]
                        samples_by_idx[idx]["indomain_base_ppl"] = float(np.exp(b_nll)) if not np.isnan(b_nll) else None

        samples_list = [samples_by_idx[k] for k in sorted(samples_by_idx)]
        samples_path = os.path.join(out_dir, "samples.json")
        with open(samples_path, "w") as f:
            json.dump(samples_list, f, indent=2)
        print(f"  Samples saved  → {samples_path}")

    if dist.is_initialized():
        dist.barrier()
        dist.destroy_process_group()


if __name__ == "__main__":
    main()