""" 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: \n{prompt}\nASSISTANT:" full_text = f"USER: \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: \n{prompt}\nASSISTANT:" full_texts = [f"USER: \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: \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()