experiment/evaluation/eval_adv_lora_val.py

"""
Minimal val-split eval: (1) keyword mention of object in captions — base vs LoRA;
(2) linear SAE probes on internals vs ground-truth labels;
(3) optional perplexity checks (out-of-domain + in-domain).
"""

from __future__ import annotations

import argparse
import json
import os
import sys

import numpy as np
import torch
import torch.distributed as dist
from peft import PeftModel
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
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,
)


DEFAULT_EVAL_PROMPTS = [
    "Describe this image.",
    "list all objects in this image",
]


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("--lora_dir", type=str, required=True)
    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 samples per category.")
    p.add_argument(
        "--prompt",
        type=str,
        default=None,
        help="Deprecated single-prompt override. If set without --prompts, only this prompt is evaluated.",
    )
    p.add_argument(
        "--prompts",
        type=str,
        nargs="+",
        default=None,
        help="Prompts to evaluate. Defaults to both description and 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="Set to empty string to skip OOD perplexity.",
    )
    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(
        "--lora_only",
        action="store_true",
        help="Skip base-model inference; run LoRA forward passes only (no base captions or base PPL).",
    )
    p.add_argument("--batch_size", type=int, default=1, help="Images per forward/generate call (generation loop).")
    p.add_argument("--probe_batch_size", type=int, default=0, help="Probe eval batch size (0 = same as --batch_size).")
    p.add_argument("--ppl_batch_size", type=int, default=0, help="PPL eval batch size (0 = same as --batch_size).")
    p.add_argument(
        "--attn_impl",
        type=str,
        default="eager",
        choices=["eager", "sdpa", "flash_attention_2"],
        help="Attention implementation passed to from_pretrained (flash_attention_2 requires flash-attn installed).",
    )
    return p.parse_args()


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


def _resolve_prompts(args) -> list[str]:
    prompts = list(args.prompts) if args.prompts else ([args.prompt] if args.prompt else list(DEFAULT_EVAL_PROMPTS))
    seen = set()
    resolved = []
    for prompt in prompts:
        prompt = prompt.strip()
        if prompt and prompt not in seen:
            resolved.append(prompt)
            seen.add(prompt)
    if not resolved:
        raise SystemExit("At least one non-empty prompt is required")
    return resolved


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


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 _sample_cat(sc_v: float, ob_v: float, cat_obj_only: str, cat_scene_only: str, cat_both: str, cat_neither: str) -> str:
    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


def _cap_cat_stats(mask: np.ndarray, base_arr, lora_arr: np.ndarray, ob_arr: np.ndarray):
    """base_arr may be None when --lora_only; base_* keys are omitted from the result in that case."""
    if not mask.any():
        return None
    n = int(mask.sum())
    l = lora_arr[mask]
    has_obj = bool((ob_arr[mask] > 0.5).all())
    lora_rate = float(l.mean())
    error_type = "miss_rate" if has_obj else "hallu_rate"
    lora_err = (1.0 - lora_rate) if has_obj else lora_rate
    out = {
        "n": n,
        "lora_mention_count": int(l.sum()),
        "lora_mention_rate": lora_rate,
        "lora_" + error_type: lora_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 _caption_metrics_from_arrays(base_arr, lora_arr: np.ndarray, ho_arr: np.ndarray, sc_arr: np.ndarray, cat_order: list[str]):
    m_to, m_bo, m_bt, m_ne = _four_category_masks(sc_arr, ho_arr)
    masks = dict(zip(cat_order, (m_to, m_bo, m_bt, m_ne)))
    cap_cats = {cat: _cap_cat_stats(masks[cat], base_arr, lora_arr, ho_arr) for cat in cat_order}
    overall = {"lora_mention_rate": float(lora_arr.mean())}
    if base_arr is not None:
        overall["base_mention_rate"] = float(base_arr.mean())
    return cap_cats, overall


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 _probe_metrics_from_arrays(
    scores_arr: np.ndarray,
    labels_arr: np.ndarray,
    sc_arr: np.ndarray,
    ob_arr: np.ndarray,
    cat_order: list[str],
):
    m_to, m_bo, m_bt, m_ne = _four_category_masks(sc_arr, ob_arr)
    masks = dict(zip(cat_order, (m_to, m_bo, m_bt, m_ne)))
    probe_cats = {cat: _probe_cat_stats(masks[cat], scores_arr, labels_arr) for cat in cat_order}

    all_preds = (scores_arr > 0.5).astype(np.float64)
    overall = {"acc": float((all_preds == labels_arr).mean())}
    try:
        from sklearn.metrics import average_precision_score, roc_auc_score

        if len(np.unique(labels_arr)) > 1:
            overall["roc_auc"] = float(roc_auc_score(labels_arr, scores_arr))
            overall["average_precision"] = float(average_precision_score(labels_arr, scores_arr))
        else:
            overall["roc_auc"] = None
            overall["average_precision"] = None
    except ImportError:
        overall["roc_auc"] = None
        overall["average_precision"] = None
    return probe_cats, overall


def _caption_nll(model, processor, image, prompt: str, caption: str, device: str) -> float:
    prompt_text = f"USER: <image>\n{prompt}\nASSISTANT:"
    full_text = f"USER: <image>\n{prompt}\nASSISTANT: {caption}"
    enc_prompt = processor(images=image, text=prompt_text, return_tensors="pt")
    enc_full = processor(images=image, text=full_text, return_tensors="pt")
    prompt_len = enc_prompt["input_ids"].shape[1]
    inputs = {k: v.to(device) for k, v in enc_full.items()}
    labels = inputs["input_ids"].clone()
    labels[:, :prompt_len] = -100
    with torch.no_grad():
        out = model(**inputs, labels=labels, use_cache=False)
    return out.loss.item()


def _caption_nll_batch(model, processor, images: list, prompt: str, captions: list[str], device: str) -> list[float]:
    """Per-sample NLL for a batch of (image, caption) pairs."""
    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 _bleu_per_category(
    base_caps: list,
    edited_caps: list,
    sc_arr: np.ndarray,
    ho_arr: np.ndarray,
    cat_order: list[str],
) -> dict:
    try:
        from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction
        _smooth = SmoothingFunction().method1
        def _score(ref: str, hyp: str) -> float:
            r, h = ref.lower().split(), hyp.lower().split()
            if not r or not h:
                return float("nan")
            return sentence_bleu([r], h, weights=(0.5, 0.5), smoothing_function=_smooth)
    except ImportError:
        def _score(ref: str, hyp: str) -> float:
            r, h = set(ref.lower().split()), set(hyp.lower().split())
            if not r or not h:
                return float("nan")
            inter = len(r & h)
            p, rec = inter / len(h), inter / len(r)
            return 2 * p * rec / (p + rec) if (p + rec) > 0 else 0.0

    m_to, m_bo, m_bt, m_ne = _four_category_masks(sc_arr, ho_arr)
    masks = dict(zip(cat_order, (m_to, m_bo, m_bt, m_ne)))
    out = {}
    for cat, mask in masks.items():
        if not mask.any():
            out[cat] = None
            continue
        scores = [_score(base_caps[i], edited_caps[i]) for i in np.where(mask)[0]]
        valid = [s for s in scores if not np.isnan(s)]
        out[cat] = float(np.mean(valid)) if valid else None
    return out


def _load_ppl_samples(dataset_id: str, split: str, max_samples: int, seed: int) -> list[dict]:
    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


def main():
    args = parse_args()
    prompts = _resolve_prompts(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")

    rc = get_relation_config(args.relation)
    dt = _dtype(args.dtype)

    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")
    scene_col, obj_col = rc.scene_key, rc.object_key

    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[str, list[int]] = {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)

    import os
    os.environ["TOKENIZERS_PARALLELISM"] = "false"
    processor = AutoProcessor.from_pretrained(args.base_model, use_fast=False)
    base = AutoModelForPreTraining.from_pretrained(
        args.base_model, torch_dtype=dt, attn_implementation=args.attn_impl
    ).to(device)
    model = PeftModel.from_pretrained(base, args.lora_dir)
    model.eval()

    kw = KeywordMentionDetector(keywords=rc.mention_keywords)
    indices = [all_indices[i] for i in range(rank, n, world_size)]
    n_prompt_evals = n * len(prompts)

    out_dir = args.output_dir or args.lora_dir
    if is_main:
        os.makedirs(out_dir, exist_ok=True)

    if is_main:
        print("=== (1) Caption keyword eval: object mention (negation-aware) ===")
    base_rates = []
    lora_rates = []
    gt_has_object = []
    base_captions = []
    lora_captions = []
    cap_indices = []
    cap_scene_flags = []
    cap_image_ids = []
    cap_prompts = []
    _live_f = open(os.path.join(out_dir, "samples_live.jsonl"), "w") if is_main else None
    _cap_desc = "(1) captions LoRA-only" if args.lora_only else "(1) captions base+LoRA"
    _gen_bs = args.batch_size
    cap_it = tqdm(total=len(indices) * len(prompts), desc=_cap_desc, 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]

        if not args.lora_only:
            model.disable_adapter_layers()
            base_texts_by_prompt = {
                prompt: generate_text_batch(model, processor, batch_images, prompt, str(device), args.max_new_tokens)
                for prompt in prompts
            }
            model.enable_adapter_layers()
        else:
            base_texts_by_prompt = {}

        for prompt in prompts:
            b_texts = base_texts_by_prompt.get(prompt)
            l_texts = generate_text_batch(model, processor, batch_images, prompt, str(device), args.max_new_tokens)

            for k, (t1, i, row) in enumerate(zip(l_texts, batch_idx, batch_rows)):
                l_m = float(kw.mentions_object(t1))
                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)
                lora_rates.append(l_m)
                lora_captions.append(t1)
                if b_texts is not None:
                    t0 = b_texts[k]
                    base_rates.append(float(kw.mentions_object(t0)))
                    base_captions.append(t0)
                if _live_f is not None:
                    _live_rec = {
                        "index": i,
                        "image_id": row.get("image_id") if hasattr(row, "get") else None,
                        "prompt": prompt,
                        scene_col: int(row[scene_col]),
                        obj_col: int(row[obj_col]),
                        "lora_caption": t1,
                        "lora_mentions_object": bool(l_m > 0.5),
                    }
                    if b_texts is not None:
                        _live_rec["base_caption"] = b_texts[k]
                        _live_rec["base_mentions_object"] = bool(base_rates[-1] > 0.5)
                    _live_f.write(json.dumps(_live_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()

    base_rates = _gather_list(base_rates, world_size)
    lora_rates = _gather_list(lora_rates, world_size)
    gt_has_object = _gather_list(gt_has_object, world_size)
    base_captions = _gather_list(base_captions, world_size)
    lora_captions = _gather_list(lora_captions, 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)

    if _live_f is not None:
        _live_f.close()
        _live_f = None

    cap_metrics = {}
    cap_metrics_overall = {}
    cap_metrics_by_prompt = {}
    bleu_vs_base: dict = {}
    if is_main:
        b_arr = np.array(base_rates, dtype=np.float64) if base_rates else None
        l_arr = np.array(lora_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)
        cap_cats, cap_overall = _caption_metrics_from_arrays(b_arr, l_arr, ho_arr, sc_cap, cat_order)

        _has_base = b_arr is not None
        print(f"  images: {n}  evals: {n_prompt_evals}  GPUs: {world_size}")
        print(f"  prompts: {prompts!r}")
        print(f"  keywords: {rc.mention_keywords[:3]!r}... (negation-aware)")
        if not _has_base:
            print("  (--lora_only: base-model columns omitted)")
        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}{'LoRA':>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"]
            lm = d["lora_mention_count"]
            lr = d["lora_mention_rate"]
            et = d["error_type"]
            le = d.get("lora_" + 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"{lm:>3}/{nr:<4}({lr:>6.1%})  "
                f"{et}: {_be_s}lora={le:.1%}"
            )
        print("  " + _sep)
        _n_pairs = len(l_arr)
        _b_overall = f"{int(b_arr.sum()):>3}/{_n_pairs:<4}({b_arr.mean():>6.1%})  " if _has_base else ""
        print(
            f"  {'OVERALL':<24} {_n_pairs:>5}  "
            f"{_b_overall}"
            f"{int(l_arr.sum()):>3}/{_n_pairs:<4}({l_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['lora_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"LoRA hallu={bd['lora_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 b_arr is not None else None
            p_l = l_arr[pmask]
            p_ho = ho_arr[pmask]
            p_sc = sc_cap[pmask]
            p_cats, p_overall = _caption_metrics_from_arrays(p_b, p_l, p_ho, p_sc, cat_order)
            cap_metrics_by_prompt[prompt] = {"overall": p_overall, "categories": p_cats}
            if _has_base:
                print(
                    f"  {prompt!r}: Base={p_overall['base_mention_rate']:.1%}  "
                    f"LoRA={p_overall['lora_mention_rate']:.1%}"
                )
            else:
                print(f"  {prompt!r}: LoRA={p_overall['lora_mention_rate']:.1%}")
        if b_arr is not None:
            bleu_vs_base = _bleu_per_category(base_captions, lora_captions, sc_cap, ho_arr, cat_order)
            print("\n  Caption similarity (edited vs base, BLEU-2) by category:")
            for cat in cat_order:
                v = bleu_vs_base.get(cat)
                if v is not None:
                    print(f"  {cat_display.get(cat, cat):<24}  {v:.4f}")
        cap_metrics = cap_cats
        cap_metrics_overall = cap_overall

    pl: list[int] = []
    scores: list[float] = []
    labels: list[float] = []
    scene_flags: list[int] = []
    object_flags: list[int] = []
    probe_indices: list[int] = []
    probe_prompts: list[str] = []
    probe_metrics: dict = {}
    probe_metrics_overall: dict = {}
    probe_metrics_by_prompt: dict = {}
    ppl_metrics: dict = {}
    indomain_ppl_metrics: dict = {}
    indomain_indices: list[int] = []
    indomain_prompts: list[str] = []
    indomain_base_nlls_outer: list[float] = []
    indomain_lora_nlls_outer: list[float] = []

    if not args.mention_only:
        if dist.is_initialized():
            dist.barrier()
        if is_main:
            print("\n=== (2) SAE probe eval (object_only labels: score→0 when obj absent, →1 when present) ===")

        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
        with torch.no_grad():
            probe_it = tqdm(total=len(indices) * len(prompts), desc="(2) SAE probe forward", disable=not is_main, unit="eval", 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()

                for prompt in prompts:
                    probe_indices.extend(batch_idx)
                    probe_prompts.extend([prompt] * len(batch_idx))
                    scene_flags.extend(sc_t.tolist())
                    object_flags.extend(ob_t.tolist())
                    labels.extend(y_batch)

                    _probe_text = f"USER: <image>\n{prompt}\nASSISTANT:"
                    _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()}

                    model.enable_adapter_layers()
                    with capture:
                        model(**batch_inputs, use_cache=False)
                    feats = get_sae_features(capture, sae)  # {layer: (B, D)}
                    probs_list = probes(feats)  # list of (B,) tensors
                    mean_p = torch.stack([p.float() for p in probs_list]).mean(dim=0)  # (B,)
                    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)
        probe_prompts = _gather_list(probe_prompts, 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)
            probe_prompt_arr = np.array(probe_prompts, dtype=object)
            probe_cats, probe_overall = _probe_metrics_from_arrays(scores_arr, labels_arr, sc_arr, ob_arr, cat_order)

            print(f"  probe_label_mode: {args.probe_label_mode}  layers: {len(pl)}")
            print(f"  prompts: {prompts!r}")
            print()
            print("  Aggregate across prompts  (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}"
                )

            print("\n  Per-prompt overall:")
            for prompt in prompts:
                pmask = probe_prompt_arr == prompt
                p_cats, p_overall = _probe_metrics_from_arrays(
                    scores_arr[pmask], labels_arr[pmask], sc_arr[pmask], ob_arr[pmask], cat_order
                )
                probe_metrics_by_prompt[prompt] = {"overall": p_overall, "categories": p_cats}
                roc_auc = p_overall.get("roc_auc")
                roc_auc_s = f"{roc_auc:.4f}" if roc_auc is not None else "n/a"
                print(f"  {prompt!r}: Acc={p_overall['acc']:.4f}  AUC={roc_auc_s}")
            probe_metrics = probe_cats
            probe_metrics_overall = probe_overall

        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))
            base_nlls_by_prompt_local = {prompt: [] for prompt in prompts}
            lora_nlls_by_prompt_local = {prompt: [] for prompt in prompts}
            _ppl_bs = args.ppl_batch_size if args.ppl_batch_size > 0 else args.batch_size
            ppl_it = tqdm(total=len(ppl_indices_local) * len(prompts), desc="(3) OOD perplexity", disable=not is_main, unit="eval", dynamic_ncols=True)
            for _pb in range(0, len(ppl_indices_local), _ppl_bs):
                batch_j = ppl_indices_local[_pb: _pb + _ppl_bs]
                # flatten samples × captions into (image, caption) pairs per sample slot
                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)
                for prompt in prompts:
                    b_nll_acc = [[] for _ in batch_j]
                    l_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
                        if not args.lora_only:
                            model.disable_adapter_layers()
                            b_slot = _caption_nll_batch(model, processor, slot_imgs, prompt, slot_caps, str(device))
                            model.enable_adapter_layers()
                            for k, nll in zip(slot_idxs, b_slot):
                                b_nll_acc[k].append(nll)
                        l_slot = _caption_nll_batch(model, processor, slot_imgs, prompt, slot_caps, str(device))
                        for k, nll in zip(slot_idxs, l_slot):
                            l_nll_acc[k].append(nll)
                    for k in range(len(batch_j)):
                        if not args.lora_only and b_nll_acc[k]:
                            base_nlls_by_prompt_local[prompt].append(float(np.mean(b_nll_acc[k])))
                        lora_nlls_by_prompt_local[prompt].append(float(np.mean(l_nll_acc[k])) if l_nll_acc[k] else float("nan"))
                    if ppl_it is not None:
                        ppl_it.update(len(batch_j))
            if ppl_it is not None:
                ppl_it.close()

            base_nlls_by_prompt = {prompt: _gather_list(vals, world_size) for prompt, vals in base_nlls_by_prompt_local.items()}
            lora_nlls_by_prompt = {prompt: _gather_list(vals, world_size) for prompt, vals in lora_nlls_by_prompt_local.items()}
            if is_main:
                ppl_per_prompt = {}
                lora_chunks = []
                base_chunks = []
                for prompt in prompts:
                    prompt_lora = lora_nlls_by_prompt[prompt]
                    if not prompt_lora:
                        continue
                    ln = np.array(prompt_lora, dtype=np.float64)
                    entry = {"n_samples": len(ln), "lora_ppl": float(np.exp(ln.mean()))}
                    lora_chunks.append(ln)
                    if base_nlls_by_prompt[prompt]:
                        bn = np.array(base_nlls_by_prompt[prompt], dtype=np.float64)
                        entry["base_ppl"] = float(np.exp(bn.mean()))
                        entry["ppl_ratio"] = entry["lora_ppl"] / entry["base_ppl"]
                        base_chunks.append(bn)
                    ppl_per_prompt[prompt] = entry
                if lora_chunks:
                    ln_all = np.concatenate(lora_chunks)
                    ppl_metrics = {
                        "dataset": args.ppl_dataset_id,
                        "split": args.ppl_split,
                        "n_samples": len(ln_all),
                        "lora_ppl": float(np.exp(ln_all.mean())),
                        "per_prompt": ppl_per_prompt,
                    }
                    if base_chunks:
                        bn_all = np.concatenate(base_chunks)
                        ppl_metrics["base_ppl"] = float(np.exp(bn_all.mean()))
                        ppl_metrics["ppl_ratio"] = ppl_metrics["lora_ppl"] / ppl_metrics["base_ppl"]
                        print(
                            f"  aggregate n={len(ln_all)}  Base PPL={ppl_metrics['base_ppl']:.3f}  "
                            f"LoRA PPL={ppl_metrics['lora_ppl']:.3f}  Ratio={ppl_metrics['ppl_ratio']:.4f}"
                        )
                    else:
                        print(f"  aggregate n={len(ln_all)}  LoRA PPL={ppl_metrics['lora_ppl']:.3f}  (base skipped: --lora_only)")
                    print("  Per-prompt:")
                    for prompt in prompts:
                        entry = ppl_per_prompt.get(prompt)
                        if entry is None:
                            continue
                        if "base_ppl" in entry:
                            print(
                                f"  {prompt!r}: Base PPL={entry['base_ppl']:.3f}  "
                                f"LoRA PPL={entry['lora_ppl']:.3f}  Ratio={entry['ppl_ratio']:.4f}"
                            )
                        else:
                            print(f"  {prompt!r}: LoRA PPL={entry['lora_ppl']:.3f}")

        if dist.is_initialized():
            dist.barrier()
        if is_main:
            print("\n=== (3b) In-domain perplexity by category (relation val set) ===")
        indomain_base_nlls: list[float] = []
        indomain_lora_nlls: list[float] = []
        indomain_scene_flags: list[int] = []
        indomain_object_flags: list[int] = []
        _indomain_indices_local: list[int] = []
        _indomain_prompts_local: list[str] = []
        n_indomain_skipped = 0
        _ppl_bs2 = args.ppl_batch_size if args.ppl_batch_size > 0 else args.batch_size
        indomain_it = tqdm(total=len(indices) * len(prompts), desc="(3b) in-domain PPL", disable=not is_main, unit="eval", dynamic_ncols=True)
        for _pb in range(0, len(indices), _ppl_bs2):
            batch_idx = indices[_pb: _pb + _ppl_bs2]
            # filter to rows that have a caption
            batch_valid: list[tuple] = []
            n_skip_batch = 0
            for i in batch_idx:
                row = ds[i]
                caption = row.get("caption", "") if hasattr(row, "get") else ""
                if not caption:
                    n_skip_batch += 1
                else:
                    batch_valid.append((i, row, caption))
            n_indomain_skipped += n_skip_batch
            if not batch_valid:
                if indomain_it is not None:
                    indomain_it.update(len(batch_idx) * len(prompts))
                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]
            scene_batch = [int(r[scene_col]) for r in b_rows]
            object_batch = [int(r[obj_col]) for r in b_rows]
            for prompt in prompts:
                if not args.lora_only:
                    model.disable_adapter_layers()
                    b_nlls_batch = _caption_nll_batch(model, processor, b_imgs, prompt, b_caps, str(device))
                    model.enable_adapter_layers()
                    indomain_base_nlls.extend(b_nlls_batch)
                l_nlls_batch = _caption_nll_batch(model, processor, b_imgs, prompt, b_caps, str(device))
                indomain_lora_nlls.extend(l_nlls_batch)
                indomain_scene_flags.extend(scene_batch)
                indomain_object_flags.extend(object_batch)
                _indomain_indices_local.extend(b_idxs)
                _indomain_prompts_local.extend([prompt] * len(b_idxs))
                if indomain_it is not None:
                    indomain_it.update(len(batch_idx))
        if indomain_it is not None:
            indomain_it.close()

        indomain_base_nlls = _gather_list(indomain_base_nlls, world_size)
        indomain_lora_nlls = _gather_list(indomain_lora_nlls, world_size)
        indomain_scene_flags = _gather_list(indomain_scene_flags, world_size)
        indomain_object_flags = _gather_list(indomain_object_flags, world_size)
        _indomain_indices_local = _gather_list(_indomain_indices_local, world_size)
        _indomain_prompts_local = _gather_list(_indomain_prompts_local, world_size)
        n_indomain_skipped = sum(_gather_list([n_indomain_skipped], world_size))
        # expose to outer scope for samples.json
        indomain_indices = _indomain_indices_local
        indomain_prompts = _indomain_prompts_local
        indomain_base_nlls_outer = indomain_base_nlls
        indomain_lora_nlls_outer = indomain_lora_nlls

        if is_main and indomain_lora_nlls:
            ln_id = np.array(indomain_lora_nlls, dtype=np.float64)
            bn_id = np.array(indomain_base_nlls, dtype=np.float64) if indomain_base_nlls else None
            sc_id = np.array(indomain_scene_flags, dtype=np.float64)
            ob_id = np.array(indomain_object_flags, dtype=np.float64)
            indomain_prompt_arr = np.array(indomain_prompts, dtype=object)

            def _indomain_summary(local_lora: np.ndarray, local_base, local_sc: np.ndarray, local_ob: np.ndarray):
                m_to_id, m_bo_id, m_bt_id, m_ne_id = _four_category_masks(local_sc, local_ob)
                categories = {}
                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():
                        categories[cat] = None
                        continue
                    l_ppl = float(np.exp(local_lora[mask].mean()))
                    entry = {"n": int(mask.sum()), "lora_ppl": l_ppl}
                    if local_base is not None:
                        b_ppl = float(np.exp(local_base[mask].mean()))
                        entry["base_ppl"] = b_ppl
                        entry["ppl_ratio"] = l_ppl / b_ppl
                    categories[cat] = entry
                overall_entry = {"lora_ppl": float(np.exp(local_lora.mean()))}
                if local_base is not None:
                    overall_entry["base_ppl"] = float(np.exp(local_base.mean()))
                    overall_entry["ppl_ratio"] = overall_entry["lora_ppl"] / overall_entry["base_ppl"]
                return categories, overall_entry

            indomain_cats, overall_entry = _indomain_summary(ln_id, bn_id, sc_id, ob_id)
            per_prompt_indomain = {}
            for prompt in prompts:
                pmask = indomain_prompt_arr == prompt
                p_bn = bn_id[pmask] if bn_id is not None else None
                p_cats, p_overall = _indomain_summary(ln_id[pmask], p_bn, sc_id[pmask], ob_id[pmask])
                per_prompt_indomain[prompt] = {
                    "n_samples": int(pmask.sum()),
                    "overall": p_overall,
                    "categories": p_cats,
                }
            indomain_ppl_metrics = {
                "n_samples": len(ln_id),
                "n_skipped_no_caption": n_indomain_skipped,
                "overall": overall_entry,
                "categories": indomain_cats,
                "per_prompt": per_prompt_indomain,
            }
            if "base_ppl" in overall_entry:
                print(
                    f"  aggregate n={len(ln_id)}  Base PPL={overall_entry['base_ppl']:.3f}  "
                    f"LoRA PPL={overall_entry['lora_ppl']:.3f}  Ratio={overall_entry['ppl_ratio']:.4f}"
                )
            else:
                print(f"  aggregate n={len(ln_id)}  LoRA PPL={overall_entry['lora_ppl']:.3f}")
            print("  Per-prompt:")
            for prompt in prompts:
                p_entry = per_prompt_indomain[prompt]["overall"]
                if "base_ppl" in p_entry:
                    print(
                        f"  {prompt!r}: Base PPL={p_entry['base_ppl']:.3f}  "
                        f"LoRA PPL={p_entry['lora_ppl']:.3f}  Ratio={p_entry['ppl_ratio']:.4f}"
                    )
                else:
                    print(f"  {prompt!r}: LoRA PPL={p_entry['lora_ppl']:.3f}")

    if is_main:
        cap_order = sorted(range(len(cap_indices)), key=lambda j: (cap_indices[j], cap_prompts[j]))
        b_arr2 = np.array(base_rates, dtype=np.float64) if base_rates else None
        l_arr2 = np.array(lora_rates, dtype=np.float64)
        ho_arr2 = np.array(gt_has_object, dtype=np.float64)
        sc_cap2 = np.array(cap_scene_flags, dtype=np.float64)

        captions_records = []
        for j in cap_order:
            cat = _sample_cat(sc_cap2[j], ho_arr2[j], cat_obj_only, cat_scene_only, cat_both, cat_neither)
            rec = {
                "index": int(cap_indices[j]),
                "image_id": cap_image_ids[j],
                "prompt": cap_prompts[j],
                scene_col: int(sc_cap2[j]),
                obj_col: int(ho_arr2[j]),
                "category": cat_display.get(cat, cat),
                "lora_caption": lora_captions[j],
                "lora_mentions_object": bool(l_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
            out = {}
            for k, v in d.items():
                if isinstance(v, np.generic):
                    v = v.item()
                if isinstance(v, float) and np.isnan(v):
                    out[k] = None
                else:
                    out[k] = v
            return out

        _lora_only_flag = args.lora_only
        metrics = {
            "relation": args.relation,
            "lora_dir": args.lora_dir,
            "lora_only": _lora_only_flag,
            "n_images": n,
            "n_prompt_evals": n_prompt_evals,
            "prompts": prompts,
            "caption_eval": {
                "overall": _safe_dict(cap_metrics_overall),
                "bleu_vs_base": bleu_vs_base if bleu_vs_base else None,
                "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[prompt]["overall"]),
                        "categories": {
                            cat_obj_only: _safe_dict(cap_metrics_by_prompt[prompt]["categories"].get(cat_obj_only)),
                            cat_scene_only: _safe_dict(cap_metrics_by_prompt[prompt]["categories"].get(cat_scene_only)),
                            cat_both: _safe_dict(cap_metrics_by_prompt[prompt]["categories"].get(cat_both)),
                            cat_neither: _safe_dict(cap_metrics_by_prompt[prompt]["categories"].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": _safe_dict(probe_metrics_overall),
                "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)),
                },
                "per_prompt": {
                    prompt: {
                        "overall": _safe_dict(probe_metrics_by_prompt[prompt]["overall"]),
                        "categories": {
                            cat_obj_only: _safe_dict(probe_metrics_by_prompt[prompt]["categories"].get(cat_obj_only)),
                            cat_scene_only: _safe_dict(probe_metrics_by_prompt[prompt]["categories"].get(cat_scene_only)),
                            cat_both: _safe_dict(probe_metrics_by_prompt[prompt]["categories"].get(cat_both)),
                            cat_neither: _safe_dict(probe_metrics_by_prompt[prompt]["categories"].get(cat_neither)),
                        },
                    }
                    for prompt in prompts
                },
            }
            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 val sample with all available data ---
        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, cat_obj_only, cat_scene_only, cat_both, cat_neither)
            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["lora_caption"] = lora_captions[j]
            prompt_rec["lora_mentions_object"] = bool(l_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["lora_caption"] = prompt_rec["lora_caption"]
                rec["lora_mentions_object"] = prompt_rec["lora_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:
                    prompt_rec = samples_by_idx[idx]["prompt_results"].setdefault(probe_prompts[j], {})
                    prompt_rec["probe_score"] = float(scores_arr2[j])
                    prompt_rec["probe_label"] = float(labels_arr2[j])
                    prompt_rec["probe_pred"] = int(scores_arr2[j] > 0.5)
                    if len(prompts) == 1:
                        samples_by_idx[idx]["probe_score"] = prompt_rec["probe_score"]
                        samples_by_idx[idx]["probe_label"] = prompt_rec["probe_label"]
                        samples_by_idx[idx]["probe_pred"] = prompt_rec["probe_pred"]

        if indomain_indices:
            bn_id_arr = np.array(indomain_base_nlls_outer, dtype=np.float64) if indomain_base_nlls_outer else None
            ln_id_arr = np.array(indomain_lora_nlls_outer, dtype=np.float64)
            for j, idx in enumerate(indomain_indices):
                if idx in samples_by_idx:
                    prompt_rec = samples_by_idx[idx]["prompt_results"].setdefault(indomain_prompts[j], {})
                    prompt_rec["indomain_lora_ppl"] = float(np.exp(ln_id_arr[j]))
                    if bn_id_arr is not None:
                        prompt_rec["indomain_base_ppl"] = float(np.exp(bn_id_arr[j])) if not np.isnan(bn_id_arr[j]) else None
                    if len(prompts) == 1:
                        samples_by_idx[idx]["indomain_lora_ppl"] = prompt_rec["indomain_lora_ppl"]
                        if "indomain_base_ppl" in prompt_rec:
                            samples_by_idx[idx]["indomain_base_ppl"] = prompt_rec["indomain_base_ppl"]

        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()