""" Val-split eval for Nullu (null-space weight editing) checkpoints: (1) keyword mention of object in captions — Nullu, optionally vs base (2) linear SAE probes on Nullu internals vs ground-truth labels (3) optional perplexity checks (OOD + in-domain) Model loading: loads base model with AutoModelForPreTraining, then splices Nullu down_proj edits for layers [--lowest_layer, --highest_layer). Base comparison requires --compare_base; default is Nullu-only. """ from __future__ import annotations import argparse import json import os import sys from pathlib import Path from typing import Dict, List, Optional import numpy as np import torch import torch.distributed as dist from tqdm import tqdm from transformers import AutoModelForPreTraining, AutoProcessor sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../..")) from experiment.config.relation_config import get_relation_config DEFAULT_EVAL_PROMPTS = [ "Describe this image.", "list all objects in this image", ] from experiment.data.hf_loader import load_hf_dataset from experiment.evaluation.inference import generate_text_batch from experiment.evaluation.metrics import KeywordMentionDetector from experiment.training.finetune_adv import ( FrozenSAEEncoder, HiddenStateCapture, count_lm_layers, get_sae_features, layer_probes_from_checkpoint, probe_labels, ) # --------------------------------------------------------------------------- # Nullu weight loading helpers # --------------------------------------------------------------------------- _DOWN_PROJ_PREFIXES = ( "model.layers", # haotian-liu LLaVA-v1.5 (Nullu default) "language_model.model.layers", # older llava-hf transformers format "language_model.layers", # newer llava-hf transformers format ) _DOWN_PROJ_SUFFIX = ".mlp.down_proj.weight" def _match_down_proj_layer(key: str) -> Optional[int]: if not key.endswith(_DOWN_PROJ_SUFFIX): return None for prefix in _DOWN_PROJ_PREFIXES: if key.startswith(prefix + "."): mid = key[len(prefix) + 1 : -len(_DOWN_PROJ_SUFFIX)] if mid.isdigit(): return int(mid) return None def _load_nullu_down_proj(checkpoint_path: str, layer_indices: List[int]) -> Dict[int, torch.Tensor]: """Read only mlp.down_proj.weight tensors for the requested layer indices.""" from safetensors.torch import safe_open ckpt = Path(checkpoint_path) wanted = set(layer_indices) result: Dict[int, torch.Tensor] = {} shard_files = sorted(ckpt.glob("*.safetensors")) if shard_files: for shard in shard_files: with safe_open(shard, framework="pt") as f: for key in f.keys(): idx = _match_down_proj_layer(key) if idx is not None and idx in wanted and idx not in result: result[idx] = f.get_tensor(key) if len(result) == len(wanted): break else: bin_files = sorted(ckpt.glob("pytorch_model*.bin")) or sorted(ckpt.glob("*.bin")) if not bin_files: raise FileNotFoundError(f"No safetensors or .bin weight files found in {ckpt}") for bf in bin_files: sd = torch.load(bf, map_location="cpu", weights_only=True) for key, tensor in sd.items(): idx = _match_down_proj_layer(key) if idx is not None and idx in wanted and idx not in result: result[idx] = tensor if len(result) == len(wanted): break missing = wanted - set(result.keys()) if missing: raise KeyError(f"Missing mlp.down_proj.weight for layers {sorted(missing)} in {ckpt}") return result def _find_lm_layers(model) -> torch.nn.ModuleList: """Locate the decoder-block ModuleList, trying known attribute paths.""" import torch.nn as nn if hasattr(model, "language_model"): lm = model.language_model if hasattr(lm, "model") and hasattr(lm.model, "layers"): return lm.model.layers if hasattr(lm, "layers"): return lm.layers if hasattr(model, "model"): inner = model.model if hasattr(inner, "language_model") and hasattr(inner.language_model, "layers"): return inner.language_model.layers for _, module in model.named_modules(): if isinstance(module, nn.ModuleList) and len(module) > 0: if hasattr(module[0], "mlp") and hasattr(module[0].mlp, "down_proj"): return module raise AttributeError("Cannot locate decoder layers in model") def _apply_nullu_weights(model, checkpoint_path: str, lowest_layer: int, highest_layer: int, device) -> int: """Apply Nullu down_proj edits in-place for layers [lowest_layer, highest_layer).""" layer_indices = list(range(lowest_layer, highest_layer)) weights = _load_nullu_down_proj(checkpoint_path, layer_indices) layers = _find_lm_layers(model) for idx, w in weights.items(): tgt = layers[idx].mlp.down_proj.weight tgt.data.copy_(w.to(device=tgt.device, dtype=tgt.dtype)) return len(weights) # --------------------------------------------------------------------------- # Argument parsing # --------------------------------------------------------------------------- def parse_args(): p = argparse.ArgumentParser() p.add_argument("--relation", type=str, default="bathroom_toilet") p.add_argument("--base_model", type=str, default="llava-hf/llava-1.5-7b-hf") p.add_argument("--nullu_checkpoint", type=str, required=True, help="Path to Nullu edited-model directory (HF checkpoint format).") p.add_argument("--lowest_layer", type=int, default=16, help="Inclusive lower bound of edited layer range.") p.add_argument("--highest_layer", type=int, default=32, help="Exclusive upper bound of edited layer range.") p.add_argument("--compare_base", action="store_true", help="Also run base-model inference for comparison (two-pass; default: Nullu-only).") p.add_argument( "--mention_only", action="store_true", help="Only compute caption keyword-mention rates (skip probes + perplexity).", ) p.add_argument("--sae_checkpoint", type=str, default="") p.add_argument("--probes_path", type=str, default="") p.add_argument("--probe_layers", type=str, default="", help="comma-separated; empty = all LM layers") p.add_argument("--probe_label_mode", type=str, choices=["union", "object_only"], default="object_only") p.add_argument("--max_samples", type=int, default=0, help="0 = full val split; >0 = max per category.") p.add_argument("--prompt", type=str, default=None, help="Single prompt override (deprecated; use --prompts).") p.add_argument("--prompts", type=str, nargs="+", default=None, help="Prompts to evaluate. Defaults to description + object-list prompts.") p.add_argument("--max_new_tokens", type=int, default=300) p.add_argument("--dtype", type=str, default="float16", choices=["float16", "bfloat16"]) p.add_argument("--output_dir", type=str, default=None) p.add_argument("--ppl_dataset_id", type=str, default="lmms-lab/COCO-Caption", help="OOD perplexity dataset. Set to empty to skip.") p.add_argument("--ppl_split", type=str, default="val") p.add_argument("--ppl_max_samples", type=int, default=0) p.add_argument("--ppl_seed", type=int, default=42) p.add_argument("--batch_size", type=int, default=1) p.add_argument("--probe_batch_size", type=int, default=0) p.add_argument("--ppl_batch_size", type=int, default=0) p.add_argument("--attn_impl", type=str, default="eager", choices=["eager", "sdpa", "flash_attention_2"]) return p.parse_args() def _resolve_prompts(args) -> list: prompts = list(args.prompts) if args.prompts else ([args.prompt] if args.prompt else list(DEFAULT_EVAL_PROMPTS)) seen: set = set() resolved = [] for p in prompts: p = p.strip() if p and p not in seen: resolved.append(p) seen.add(p) return resolved or list(DEFAULT_EVAL_PROMPTS) # --------------------------------------------------------------------------- # Distributed helpers (unchanged from lora eval) # --------------------------------------------------------------------------- def _dtype(s: str): return torch.float16 if s == "float16" else torch.bfloat16 def _setup_dist(): if "LOCAL_RANK" not in os.environ: return 0, 1, torch.device("cuda" if torch.cuda.is_available() else "cpu") local_rank = int(os.environ["LOCAL_RANK"]) torch.cuda.set_device(local_rank) dist.init_process_group(backend="nccl") rank = dist.get_rank() world_size = dist.get_world_size() return rank, world_size, torch.device(f"cuda:{local_rank}") def _gather_list(local: list, world_size: int) -> list: if world_size == 1: return local bucket = [None] * world_size dist.all_gather_object(bucket, local) out = [] for part in bucket: out.extend(part) return out # --------------------------------------------------------------------------- # Metric helpers # --------------------------------------------------------------------------- def _four_category_masks(sc_arr: np.ndarray, ob_arr: np.ndarray): return ( (sc_arr < 0.5) & (ob_arr > 0.5), (sc_arr > 0.5) & (ob_arr < 0.5), (sc_arr > 0.5) & (ob_arr > 0.5), (sc_arr < 0.5) & (ob_arr < 0.5), ) def _category_display_names(rc): return { rc.non_scene_with_object: f"{rc.object_key}_only", rc.scene_no_object: f"{rc.scene_key}_only", rc.scene_with_object: f"{rc.scene_key}_{rc.object_key}", "neither": "neither", } def _cap_cat_stats(mask: np.ndarray, base_arr, nullu_arr: np.ndarray, ob_arr: np.ndarray): """base_arr may be None when --compare_base is not set.""" if not mask.any(): return None n = int(mask.sum()) nu = nullu_arr[mask] has_obj = bool((ob_arr[mask] > 0.5).all()) nullu_rate = float(nu.mean()) error_type = "miss_rate" if has_obj else "hallu_rate" nullu_err = (1.0 - nullu_rate) if has_obj else nullu_rate out = { "n": n, "nullu_mention_count": int(nu.sum()), "nullu_mention_rate": nullu_rate, "nullu_" + error_type: nullu_err, "error_type": error_type, } if base_arr is not None: b = base_arr[mask] base_rate = float(b.mean()) base_err = (1.0 - base_rate) if has_obj else base_rate out["base_mention_count"] = int(b.sum()) out["base_mention_rate"] = base_rate out["base_" + error_type] = base_err return out def _probe_cat_stats(mask: np.ndarray, scores_arr: np.ndarray, labels_arr: np.ndarray): if not mask.any(): return None s = scores_arr[mask] l = labels_arr[mask] preds = (s > 0.5).astype(np.float64) tp = float(((preds == 1) & (l == 1)).sum()) fp = float(((preds == 1) & (l == 0)).sum()) fn = float(((preds == 0) & (l == 1)).sum()) tn = float(((preds == 0) & (l == 0)).sum()) acc = float((preds == l).mean()) precision = tp / (tp + fp) if (tp + fp) > 0 else float("nan") recall = tp / (tp + fn) if (tp + fn) > 0 else float("nan") f1 = ( (2 * precision * recall / (precision + recall)) if (precision + recall) > 0 and not (np.isnan(precision) or np.isnan(recall)) else float("nan") ) auc = None try: from sklearn.metrics import roc_auc_score if len(np.unique(l)) > 1: auc = float(roc_auc_score(l, s)) except ImportError: pass return { "n": int(mask.sum()), "label": float(l[0]) if len(l) else float("nan"), "avg_prob": float(s.mean()), "acc": acc, "precision": precision, "recall": recall, "f1": f1, "auc": auc, "tp": int(tp), "fp": int(fp), "fn": int(fn), "tn": int(tn), } def _caption_nll_batch(model, processor, images: list, prompt: str, captions: list, device: str) -> list: import torch.nn as nn prompt_text = f"USER: \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 _load_ppl_samples(dataset_id: str, split: str, max_samples: int, seed: int) -> list: from datasets import load_dataset as _hf_load ds = _hf_load(dataset_id, split=split) cap_col = next((c for c in ("answer", "captions", "caption") if c in ds.column_names), None) if cap_col is None: raise ValueError(f"No caption column found in {dataset_id}. Columns: {ds.column_names}") n = len(ds) if max_samples <= 0 else min(max_samples, len(ds)) rng = np.random.default_rng(seed) idx = sorted(rng.choice(len(ds), size=n, replace=False).tolist()) samples = [] for i in idx: row = ds[int(i)] img = row["image"].convert("RGB") caps = row[cap_col] if isinstance(caps, str): caps = [caps] else: caps = [c if isinstance(c, str) else (c.get("raw") or c.get("caption") or "") for c in caps] caps = [c for c in caps if c] if caps: samples.append({"image": img, "captions": caps}) return samples # --------------------------------------------------------------------------- # OOD PPL helper (single model pass, returns per-sample NLLs) # --------------------------------------------------------------------------- def _run_ood_ppl(model, processor, ppl_samples: list, indices_local: list, args, device: str) -> list: nlls: list = [] _ppl_bs = args.ppl_batch_size if args.ppl_batch_size > 0 else args.batch_size for _pb in range(0, len(indices_local), _ppl_bs): batch_j = indices_local[_pb: _pb + _ppl_bs] batch_images_multi = [ppl_samples[j]["image"] for j in batch_j] batch_captions_multi = [ppl_samples[j]["captions"] for j in batch_j] max_caps = max(len(c) for c in batch_captions_multi) nll_acc = [[] for _ in batch_j] for cap_slot in range(max_caps): slot_imgs, slot_caps, slot_idxs = [], [], [] for k, caps in enumerate(batch_captions_multi): if cap_slot < len(caps): slot_imgs.append(batch_images_multi[k]) slot_caps.append(caps[cap_slot]) slot_idxs.append(k) if not slot_imgs: continue slot_nlls = _caption_nll_batch(model, processor, slot_imgs, args.prompt, slot_caps, device) for k, nll in zip(slot_idxs, slot_nlls): nll_acc[k].append(nll) for k in range(len(batch_j)): nlls.append(float(np.mean(nll_acc[k])) if nll_acc[k] else float("nan")) return nlls # --------------------------------------------------------------------------- # In-domain PPL helper # --------------------------------------------------------------------------- def _run_indomain_ppl(model, processor, ds, indices: list, scene_col: str, obj_col: str, args, device: str): nlls: list = [] scene_flags_out: list = [] object_flags_out: list = [] indices_out: list = [] n_skipped = 0 _ppl_bs = args.ppl_batch_size if args.ppl_batch_size > 0 else args.batch_size for _pb in range(0, len(indices), _ppl_bs): batch_idx = indices[_pb: _pb + _ppl_bs] batch_valid = [] for i in batch_idx: row = ds[i] caption = row.get("caption", "") if hasattr(row, "get") else "" if not caption: n_skipped += 1 else: batch_valid.append((i, row, caption)) if not batch_valid: continue b_idxs = [t[0] for t in batch_valid] b_rows = [t[1] for t in batch_valid] b_caps = [t[2] for t in batch_valid] b_imgs = [r["image"].convert("RGB") for r in b_rows] batch_nlls = _caption_nll_batch(model, processor, b_imgs, args.prompt, b_caps, device) nlls.extend(batch_nlls) scene_flags_out.extend([int(r[scene_col]) for r in b_rows]) object_flags_out.extend([int(r[obj_col]) for r in b_rows]) indices_out.extend(b_idxs) return nlls, scene_flags_out, object_flags_out, indices_out, n_skipped # --------------------------------------------------------------------------- # main # --------------------------------------------------------------------------- def main(): args = parse_args() rank, world_size, device = _setup_dist() is_main = rank == 0 if args.mention_only: args.ppl_dataset_id = "" else: if not args.sae_checkpoint: raise SystemExit("--sae_checkpoint is required unless --mention_only is set") if not args.probes_path: raise SystemExit("--probes_path is required unless --mention_only is set") prompts = _resolve_prompts(args) rc = get_relation_config(args.relation) dt = _dtype(args.dtype) scene_col, obj_col = rc.scene_key, rc.object_key cat_obj_only = rc.non_scene_with_object cat_scene_only = rc.scene_no_object cat_both = rc.scene_with_object cat_neither = "neither" cat_order = [cat_obj_only, cat_scene_only, cat_both, cat_neither] cat_display = _category_display_names(rc) ds = load_hf_dataset(rc.dataset_id, split="val") # Build sample index if args.max_samples <= 0: all_indices = list(range(len(ds))) else: sc_labels = ds[scene_col] ob_labels = ds[obj_col] cat_buckets: dict = {cat_obj_only: [], cat_scene_only: [], cat_both: [], cat_neither: []} for i, (sc_v, ob_v) in enumerate(zip(sc_labels, ob_labels)): if int(sc_v) == 0 and int(ob_v) == 1: cat_buckets[cat_obj_only].append(i) elif int(sc_v) == 1 and int(ob_v) == 0: cat_buckets[cat_scene_only].append(i) elif int(sc_v) == 1 and int(ob_v) == 1: cat_buckets[cat_both].append(i) else: cat_buckets[cat_neither].append(i) all_indices = [] for bucket in cat_buckets.values(): all_indices.extend(bucket[: args.max_samples]) all_indices.sort() n = len(all_indices) out_dir = args.output_dir or os.path.join(args.nullu_checkpoint, "eval") if is_main: os.makedirs(out_dir, exist_ok=True) processor = AutoProcessor.from_pretrained(args.base_model) model = AutoModelForPreTraining.from_pretrained( args.base_model, torch_dtype=dt, attn_implementation=args.attn_impl ).to(device) model.eval() kw = KeywordMentionDetector(keywords=rc.mention_keywords) indices = [all_indices[i] for i in range(rank, n, world_size)] _gen_bs = args.batch_size # ------------------------------------------------------------------ # (0) Optionally run base-model captions + PPL before editing weights # ------------------------------------------------------------------ base_rates: list = [] base_captions: list = [] base_ood_nlls: list = [] base_indomain_nlls: list = [] base_indomain_scene: list = [] base_indomain_object: list = [] base_indomain_indices: list = [] if args.compare_base: if is_main: print("=== (0) Base-model caption eval (before Nullu edit) ===") base_it = tqdm(total=len(indices) * len(prompts), desc="(0) base captions", disable=not is_main, unit="eval", dynamic_ncols=True) for _bs in range(0, len(indices), _gen_bs): batch_idx = indices[_bs: _bs + _gen_bs] batch_rows = [ds[i] for i in batch_idx] batch_images = [r["image"].convert("RGB") for r in batch_rows] for prompt in prompts: b_texts = generate_text_batch(model, processor, batch_images, prompt, str(device), args.max_new_tokens) for t in b_texts: base_rates.append(float(kw.mentions_object(t))) base_captions.append(t) if base_it is not None: base_it.update(len(batch_idx)) if base_it is not None: base_it.close() base_rates = _gather_list(base_rates, world_size) base_captions = _gather_list(base_captions, world_size) if not args.mention_only: # Base OOD PPL if args.ppl_dataset_id: ppl_samples = _load_ppl_samples(args.ppl_dataset_id, args.ppl_split, args.ppl_max_samples, args.ppl_seed) ppl_indices_local = list(range(rank, len(ppl_samples), world_size)) _base_ood = _run_ood_ppl(model, processor, ppl_samples, ppl_indices_local, args, str(device)) base_ood_nlls = _gather_list(_base_ood, world_size) # Base in-domain PPL _b_id_nlls, _b_id_sc, _b_id_ob, _b_id_idx, _ = _run_indomain_ppl( model, processor, ds, indices, scene_col, obj_col, args, str(device) ) base_indomain_nlls = _gather_list(_b_id_nlls, world_size) base_indomain_scene = _gather_list(_b_id_sc, world_size) base_indomain_object = _gather_list(_b_id_ob, world_size) base_indomain_indices = _gather_list(_b_id_idx, world_size) # ------------------------------------------------------------------ # Apply Nullu weight edits # ------------------------------------------------------------------ n_edited = _apply_nullu_weights(model, args.nullu_checkpoint, args.lowest_layer, args.highest_layer, device) if is_main: print( f"\nApplied Nullu: {n_edited} layers edited " f"[{args.lowest_layer}, {args.highest_layer}) " f"checkpoint: {args.nullu_checkpoint}" ) # ------------------------------------------------------------------ # (1) Nullu caption eval # ------------------------------------------------------------------ if is_main: print("\n=== (1) Caption keyword eval: object mention (negation-aware) ===") nullu_rates: list = [] nullu_captions: list = [] gt_has_object: list = [] cap_indices: list = [] cap_scene_flags: list = [] cap_image_ids: list = [] cap_prompts: list = [] _live_f = open(os.path.join(out_dir, "samples_live.jsonl"), "w") if is_main else None cap_it = tqdm(total=len(indices) * len(prompts), desc="(1) Nullu captions", disable=not is_main, unit="eval", dynamic_ncols=True) for _bs in range(0, len(indices), _gen_bs): batch_idx = indices[_bs: _bs + _gen_bs] batch_rows = [ds[i] for i in batch_idx] batch_images = [r["image"].convert("RGB") for r in batch_rows] for prompt in prompts: nu_texts = generate_text_batch(model, processor, batch_images, prompt, str(device), args.max_new_tokens) for k, (t, i, row) in enumerate(zip(nu_texts, batch_idx, batch_rows)): nu_m = float(kw.mentions_object(t)) nullu_rates.append(nu_m) nullu_captions.append(t) gt_has_object.append(float(int(row[obj_col]))) cap_indices.append(i) cap_scene_flags.append(int(row[scene_col])) cap_image_ids.append(row.get("image_id") if hasattr(row, "get") else None) cap_prompts.append(prompt) if _live_f is not None: rec = { "index": i, "prompt": prompt, "image_id": row.get("image_id") if hasattr(row, "get") else None, scene_col: int(row[scene_col]), obj_col: int(row[obj_col]), "nullu_caption": t, "nullu_mentions_object": bool(nu_m > 0.5), } if args.compare_base: rec["compare_base"] = True _live_f.write(json.dumps(rec) + "\n") if _live_f is not None: _live_f.flush() if cap_it is not None: cap_it.update(len(batch_idx)) if cap_it is not None: cap_it.close() if _live_f is not None: _live_f.close() nullu_rates = _gather_list(nullu_rates, world_size) nullu_captions = _gather_list(nullu_captions, world_size) gt_has_object = _gather_list(gt_has_object, world_size) cap_indices = _gather_list(cap_indices, world_size) cap_scene_flags = _gather_list(cap_scene_flags, world_size) cap_image_ids = _gather_list(cap_image_ids, world_size) cap_prompts = _gather_list(cap_prompts, world_size) cap_metrics: dict = {} cap_metrics_by_prompt: dict = {} if is_main: b_arr = np.array(base_rates, dtype=np.float64) if base_rates else None nu_arr = np.array(nullu_rates, dtype=np.float64) ho_arr = np.array(gt_has_object, dtype=np.float64) sc_cap = np.array(cap_scene_flags, dtype=np.float64) cap_prompt_arr = np.array(cap_prompts, dtype=object) n_evals = len(nu_arr) m_to, m_bo, m_bt, m_ne = _four_category_masks(sc_cap, ho_arr) cap_cats = { cat_obj_only: _cap_cat_stats(m_to, b_arr, nu_arr, ho_arr), cat_scene_only: _cap_cat_stats(m_bo, b_arr, nu_arr, ho_arr), cat_both: _cap_cat_stats(m_bt, b_arr, nu_arr, ho_arr), cat_neither: _cap_cat_stats(m_ne, b_arr, nu_arr, ho_arr), } _has_base = b_arr is not None print(f" images: {n} evals: {n_evals} GPUs: {world_size} layers_edited: [{args.lowest_layer}, {args.highest_layer})") print(f" prompts: {prompts!r}") print(f" keywords: {rc.mention_keywords[:3]!r}... (negation-aware)") if not _has_base: print(" (Nullu-only: base columns omitted — use --compare_base to add base)") print() print(" Aggregate across prompts:") _sep = "-" * (91 if _has_base else 72) _hdr = f" {'Base':>14} " if _has_base else "" print(f" {'Category':<24} {'N':>5} {_hdr}{'Nullu':>14} Error") print(" " + _sep) for cat in cat_order: d = cap_cats.get(cat) if d is None: print(f" {cat_display.get(cat, cat):<18} {'(empty)'}") continue display_name = cat_display.get(cat, cat) nr = d["n"] nm = d["nullu_mention_count"] nr_rate = d["nullu_mention_rate"] et = d["error_type"] ne = d.get("nullu_" + et, float("nan")) _base_col = ( f"{d['base_mention_count']:>3}/{nr:<4}({d['base_mention_rate']:>6.1%}) " if _has_base else "" ) _be_s = f"base={d.get('base_' + et, float('nan')):.1%} " if _has_base else "" print( f" {display_name:<24} {nr:>5} " f"{_base_col}" f"{nm:>3}/{nr:<4}({nr_rate:>6.1%}) " f"{et}: {_be_s}nullu={ne:.1%}" ) print(" " + _sep) _b_overall = f"{int(b_arr.sum()):>3}/{n_evals:<4}({b_arr.mean():>6.1%}) " if _has_base else "" print( f" {'OVERALL':<24} {n_evals:>5} " f"{_b_overall}" f"{int(nu_arr.sum()):>3}/{n_evals:<4}({nu_arr.mean():>6.1%})" ) if cap_cats[cat_scene_only] is not None: bd = cap_cats[cat_scene_only] _sup_base = f"base hallu={bd['base_hallu_rate']:.1%} " if _has_base else "" _delta = ( f" Δ={bd['base_hallu_rate'] - bd['nullu_hallu_rate']:+.1%}" if _has_base else "" ) print( f"\n [Suppression] {cat_display.get(cat_scene_only, cat_scene_only)} (D_{{A,¬B}}): " f"{_sup_base}" f"Nullu hallu={bd['nullu_hallu_rate']:.1%}" f"{_delta}" ) print("\n Per-prompt overall:") for prompt in prompts: pmask = cap_prompt_arr == prompt p_b = b_arr[pmask] if _has_base else None p_nu = nu_arr[pmask] p_ho = ho_arr[pmask] p_sc = sc_cap[pmask] pm_to, pm_bo, pm_bt, pm_ne = _four_category_masks(p_sc, p_ho) p_cats = { cat_obj_only: _cap_cat_stats(pm_to, p_b, p_nu, p_ho), cat_scene_only: _cap_cat_stats(pm_bo, p_b, p_nu, p_ho), cat_both: _cap_cat_stats(pm_bt, p_b, p_nu, p_ho), cat_neither: _cap_cat_stats(pm_ne, p_b, p_nu, p_ho), } cap_metrics_by_prompt[prompt] = { "overall": {"nullu_mention_rate": float(p_nu.mean())}, "categories": p_cats, } if _has_base: cap_metrics_by_prompt[prompt]["overall"]["base_mention_rate"] = float(p_b.mean()) print(f" {prompt!r}: base={p_b.mean():.1%} nullu={p_nu.mean():.1%}") else: print(f" {prompt!r}: nullu={p_nu.mean():.1%}") cap_metrics = cap_cats # ------------------------------------------------------------------ # (2) SAE probe eval (Nullu model) # ------------------------------------------------------------------ pl: list = [] scores: list = [] labels: list = [] scene_flags: list = [] object_flags: list = [] probe_indices: list = [] probe_metrics: dict = {} overall_acc = None overall_auc = None overall_ap = None ppl_metrics: dict = {} indomain_ppl_metrics: dict = {} indomain_indices: list = [] indomain_nullu_nlls_outer: list = [] indomain_base_nlls_outer: list = [] if not args.mention_only: if dist.is_initialized(): dist.barrier() if is_main: print("\n=== (2) SAE probe eval (object_only labels) ===") sae = FrozenSAEEncoder.from_checkpoint(args.sae_checkpoint, device) d_sae = sae.encoder.weight.shape[0] n_layers = count_lm_layers(model) pl = ( [int(x) for x in args.probe_layers.split(",") if x.strip()] if args.probe_layers.strip() else list(range(n_layers)) ) probes = layer_probes_from_checkpoint(args.probes_path, pl, d_sae, device=device) probes.eval() capture = HiddenStateCapture(model, pl) _probe_bs = args.probe_batch_size if args.probe_batch_size > 0 else args.batch_size _probe_text = f"USER: \n{args.prompt}\nASSISTANT:" with torch.no_grad(): probe_it = tqdm(total=len(indices), desc="(2) SAE probe forward", disable=not is_main, unit="img", dynamic_ncols=True) for _pb in range(0, len(indices), _probe_bs): batch_idx = indices[_pb: _pb + _probe_bs] batch_rows = [ds[i] for i in batch_idx] batch_images = [r["image"].convert("RGB") for r in batch_rows] sc_t = torch.tensor([int(r[scene_col]) for r in batch_rows], device=device) ob_t = torch.tensor([int(r[obj_col]) for r in batch_rows], device=device) y_batch = probe_labels(sc_t, ob_t, args.probe_label_mode).tolist() probe_indices.extend(batch_idx) scene_flags.extend(sc_t.tolist()) object_flags.extend(ob_t.tolist()) labels.extend(y_batch) _orig_side = processor.tokenizer.padding_side processor.tokenizer.padding_side = "left" try: batch_inputs = processor( images=batch_images, text=[_probe_text] * len(batch_images), return_tensors="pt", padding=True, ) finally: processor.tokenizer.padding_side = _orig_side batch_inputs = {k: v.to(device) for k, v in batch_inputs.items()} with capture: model(**batch_inputs, use_cache=False) feats = get_sae_features(capture, sae) probs_list = probes(feats) mean_p = torch.stack([p.float() for p in probs_list]).mean(dim=0) scores.extend(mean_p.tolist()) if probe_it is not None: probe_it.update(len(batch_idx)) if probe_it is not None: probe_it.close() scores = _gather_list(scores, world_size) labels = _gather_list(labels, world_size) scene_flags = _gather_list(scene_flags, world_size) object_flags = _gather_list(object_flags, world_size) probe_indices = _gather_list(probe_indices, world_size) if is_main: scores_arr = np.array(scores, dtype=np.float64) labels_arr = np.array(labels, dtype=np.float64) sc_arr = np.array(scene_flags, dtype=np.float64) ob_arr = np.array(object_flags, dtype=np.float64) m_to, m_bo, m_bt, m_ne = _four_category_masks(sc_arr, ob_arr) probe_cats = { cat_obj_only: _probe_cat_stats(m_to, scores_arr, labels_arr), cat_scene_only: _probe_cat_stats(m_bo, scores_arr, labels_arr), cat_both: _probe_cat_stats(m_bt, scores_arr, labels_arr), cat_neither: _probe_cat_stats(m_ne, scores_arr, labels_arr), } print(f" probe_label_mode: {args.probe_label_mode} layers: {len(pl)}") print() print(" Per-category breakdown (label: 1=probe fires, 0=probe silent)") print( f" {'Category':<24} {'N':>5} {'Lbl':>3} " f"{'AvgProb':>7} {'Acc@0.5':>7} {'Prec':>6} {'Rec':>6} {'F1':>6} {'AUC':>6}" ) print(" " + "-" * 102) for cat in cat_order: d = probe_cats.get(cat) if d is None: print(f" {cat_display.get(cat, cat):<24} {'(empty)'}") continue auc_s = f"{d['auc']:.4f}" if d["auc"] is not None else " n/a" prec_s = f"{d['precision']:.4f}" if not np.isnan(d["precision"]) else " n/a" rec_s = f"{d['recall']:.4f}" if not np.isnan(d["recall"]) else " n/a" f1_s = f"{d['f1']:.4f}" if not np.isnan(d["f1"]) else " n/a" print( f" {cat_display.get(cat, cat):<24} {d['n']:>5} {d['label']:>3.0f} " f"{d['avg_prob']:>7.4f} {d['acc']:>7.4f} " f"{prec_s:>6} {rec_s:>6} {f1_s:>6} {auc_s:>6}" ) all_preds = (scores_arr > 0.5).astype(np.float64) overall_acc = float((all_preds == labels_arr).mean()) try: from sklearn.metrics import roc_auc_score, average_precision_score if len(np.unique(labels_arr)) > 1: overall_auc = float(roc_auc_score(labels_arr, scores_arr)) overall_ap = float(average_precision_score(labels_arr, scores_arr)) except ImportError: pass probe_metrics = probe_cats if dist.is_initialized(): dist.barrier() # (3) OOD perplexity if not args.ppl_dataset_id: if is_main: print("\n=== (3) Out-of-domain perplexity: SKIPPED (--ppl_dataset_id is empty) ===") else: if is_main: print(f"\n=== (3) Out-of-domain perplexity [{args.ppl_dataset_id} / {args.ppl_split}] ===") ppl_samples = _load_ppl_samples(args.ppl_dataset_id, args.ppl_split, args.ppl_max_samples, args.ppl_seed) ppl_indices_local = list(range(rank, len(ppl_samples), world_size)) nullu_ood_nlls: list = [] ppl_it = tqdm(total=len(ppl_indices_local), desc="(3) OOD PPL Nullu", disable=not is_main, unit="img", dynamic_ncols=True) _tmp = _run_ood_ppl(model, processor, ppl_samples, ppl_indices_local, args, str(device)) if ppl_it is not None: ppl_it.update(len(ppl_indices_local)) ppl_it.close() nullu_ood_nlls = _gather_list(_tmp, world_size) base_ood_nlls_gathered = _gather_list(base_ood_nlls, world_size) if base_ood_nlls else [] if is_main and nullu_ood_nlls: ln = np.array(nullu_ood_nlls) nullu_ppl = float(np.exp(ln.mean())) if base_ood_nlls_gathered: bn = np.array(base_ood_nlls_gathered) base_ppl = float(np.exp(bn.mean())) ratio = nullu_ppl / base_ppl print(f" n={len(ln)} Base PPL={base_ppl:.3f} Nullu PPL={nullu_ppl:.3f} Ratio={ratio:.4f}") ppl_metrics = { "dataset": args.ppl_dataset_id, "split": args.ppl_split, "n_samples": len(ln), "base_ppl": base_ppl, "nullu_ppl": nullu_ppl, "ppl_ratio": ratio, } else: print(f" n={len(ln)} Nullu PPL={nullu_ppl:.3f} (base skipped)") ppl_metrics = { "dataset": args.ppl_dataset_id, "split": args.ppl_split, "n_samples": len(ln), "nullu_ppl": nullu_ppl, } if dist.is_initialized(): dist.barrier() # (3b) In-domain perplexity by category if is_main: print("\n=== (3b) In-domain perplexity by category (relation val set) ===") _nu_id_nlls, _nu_id_sc, _nu_id_ob, _nu_id_idx, n_skip_local = _run_indomain_ppl( model, processor, ds, indices, scene_col, obj_col, args, str(device) ) nullu_id_nlls = _gather_list(_nu_id_nlls, world_size) nullu_id_sc = _gather_list(_nu_id_sc, world_size) nullu_id_ob = _gather_list(_nu_id_ob, world_size) nullu_id_idx = _gather_list(_nu_id_idx, world_size) n_skip_total = sum(_gather_list([n_skip_local], world_size)) # Use base in-domain data if available (already gathered above) base_id_nlls_arr = np.array(base_indomain_nlls) if base_indomain_nlls else None indomain_indices = nullu_id_idx indomain_nullu_nlls_outer = nullu_id_nlls indomain_base_nlls_outer = base_indomain_nlls if is_main and nullu_id_nlls: ln_id = np.array(nullu_id_nlls) sc_id = np.array(nullu_id_sc, dtype=np.float64) ob_id = np.array(nullu_id_ob, dtype=np.float64) m_to_id, m_bo_id, m_bt_id, m_ne_id = _four_category_masks(sc_id, ob_id) # Align base in-domain NLLs to same sample order if available base_id_aligned = None if base_indomain_nlls and base_indomain_indices: # Build a map from index → base NLL for lookup base_idx_to_nll = {idx: nll for idx, nll in zip(base_indomain_indices, base_indomain_nlls)} base_id_aligned = np.array([base_idx_to_nll.get(idx, float("nan")) for idx in nullu_id_idx]) indomain_cats: dict = {} for cat, mask in [ (cat_obj_only, m_to_id), (cat_scene_only, m_bo_id), (cat_both, m_bt_id), (cat_neither, m_ne_id), ]: if not mask.any(): indomain_cats[cat] = None continue l_ppl = float(np.exp(ln_id[mask].mean())) entry: dict = {"n": int(mask.sum()), "nullu_ppl": l_ppl} if base_id_aligned is not None: valid_mask = mask & ~np.isnan(base_id_aligned) if valid_mask.any(): b_ppl = float(np.exp(base_id_aligned[valid_mask].mean())) entry["base_ppl"] = b_ppl entry["ppl_ratio"] = l_ppl / b_ppl indomain_cats[cat] = entry l_overall = float(np.exp(ln_id.mean())) overall_entry: dict = {"nullu_ppl": l_overall} if base_id_aligned is not None: valid = ~np.isnan(base_id_aligned) if valid.any(): b_overall = float(np.exp(base_id_aligned[valid].mean())) overall_entry["base_ppl"] = b_overall overall_entry["ppl_ratio"] = l_overall / b_overall indomain_ppl_metrics = { "n_samples": len(ln_id), "n_skipped_no_caption": n_skip_total, "overall": overall_entry, "categories": indomain_cats, } # ------------------------------------------------------------------ # Save outputs # ------------------------------------------------------------------ if is_main: cap_order = sorted(range(len(cap_indices)), key=lambda j: (cap_indices[j], cap_prompts[j])) nu_arr2 = np.array(nullu_rates, dtype=np.float64) ho_arr2 = np.array(gt_has_object, dtype=np.float64) sc_cap2 = np.array(cap_scene_flags, dtype=np.float64) b_arr2 = np.array(base_rates, dtype=np.float64) if base_rates else None def _sample_cat(sc_v, ob_v): if sc_v < 0.5 and ob_v > 0.5: return cat_obj_only if sc_v > 0.5 and ob_v < 0.5: return cat_scene_only if sc_v > 0.5 and ob_v > 0.5: return cat_both return cat_neither captions_records = [] for j in cap_order: cat = _sample_cat(sc_cap2[j], ho_arr2[j]) rec = { "index": int(cap_indices[j]), "prompt": cap_prompts[j], "image_id": cap_image_ids[j], scene_col: int(sc_cap2[j]), obj_col: int(ho_arr2[j]), "category": cat_display.get(cat, cat), "nullu_caption": nullu_captions[j], "nullu_mentions_object": bool(nu_arr2[j] > 0.5), } if b_arr2 is not None: rec["base_caption"] = base_captions[j] rec["base_mentions_object"] = bool(b_arr2[j] > 0.5) captions_records.append(rec) captions_path = os.path.join(out_dir, "captions.json") with open(captions_path, "w") as f: json.dump(captions_records, f, indent=2) print(f"\n Captions saved → {captions_path}") def _safe_dict(d): if d is None: return None return {k: (None if (isinstance(v, float) and np.isnan(v)) else v) for k, v in d.items()} nu_f = np.array(nullu_rates, dtype=np.float64) caption_overall: dict = {"nullu_mention_rate": float(nu_f.mean())} if args.compare_base and base_rates: b_f = np.array(base_rates, dtype=np.float64) caption_overall["base_mention_rate"] = float(b_f.mean()) metrics = { "relation": args.relation, "nullu_checkpoint": args.nullu_checkpoint, "lowest_layer": args.lowest_layer, "highest_layer": args.highest_layer, "compare_base": args.compare_base, "n_images": n, "n_prompt_evals": len(nu_f), "prompts": prompts, "caption_eval": { "overall": caption_overall, "categories": { cat_obj_only: _safe_dict(cap_metrics.get(cat_obj_only)), cat_scene_only: _safe_dict(cap_metrics.get(cat_scene_only)), cat_both: _safe_dict(cap_metrics.get(cat_both)), cat_neither: _safe_dict(cap_metrics.get(cat_neither)), }, "per_prompt": { prompt: { "overall": _safe_dict(cap_metrics_by_prompt.get(prompt, {}).get("overall")), "categories": { cat_obj_only: _safe_dict((cap_metrics_by_prompt.get(prompt, {}).get("categories") or {}).get(cat_obj_only)), cat_scene_only: _safe_dict((cap_metrics_by_prompt.get(prompt, {}).get("categories") or {}).get(cat_scene_only)), cat_both: _safe_dict((cap_metrics_by_prompt.get(prompt, {}).get("categories") or {}).get(cat_both)), cat_neither: _safe_dict((cap_metrics_by_prompt.get(prompt, {}).get("categories") or {}).get(cat_neither)), }, } for prompt in prompts }, }, } if len(prompts) == 1: metrics["prompt"] = prompts[0] scores_arr2 = np.array(scores, dtype=np.float64) labels_arr2 = np.array(labels, dtype=np.float64) if not args.mention_only: metrics["probe_label_mode"] = args.probe_label_mode metrics["n_probe_layers"] = len(pl) metrics["probe_eval"] = { "overall": { "acc": float(overall_acc) if overall_acc is not None else None, "roc_auc": overall_auc, "average_precision": overall_ap, }, "categories": { cat_obj_only: _safe_dict(probe_metrics.get(cat_obj_only)), cat_scene_only: _safe_dict(probe_metrics.get(cat_scene_only)), cat_both: _safe_dict(probe_metrics.get(cat_both)), cat_neither: _safe_dict(probe_metrics.get(cat_neither)), }, } metrics["perplexity_ood"] = ppl_metrics metrics["perplexity_indomain"] = indomain_ppl_metrics metrics_path = os.path.join(out_dir, "metrics.json") with open(metrics_path, "w") as f: json.dump(metrics, f, indent=2) print(f" Metrics saved → {metrics_path}") # Build samples.json — one record per image with prompt_results nested samples_by_idx: dict = {} for j in cap_order: idx = int(cap_indices[j]) sc_v = float(sc_cap2[j]) ob_v = float(ho_arr2[j]) cat = _sample_cat(sc_v, ob_v) rec = samples_by_idx.setdefault( idx, { "index": idx, "image_id": cap_image_ids[j], scene_col: int(sc_v), obj_col: int(ob_v), "category": cat_display.get(cat, cat), "prompt_results": {}, }, ) prompt_rec = rec["prompt_results"].setdefault(cap_prompts[j], {}) prompt_rec["nullu_caption"] = nullu_captions[j] prompt_rec["nullu_mentions_object"] = bool(nu_arr2[j] > 0.5) if b_arr2 is not None: prompt_rec["base_caption"] = base_captions[j] prompt_rec["base_mentions_object"] = bool(b_arr2[j] > 0.5) if len(prompts) == 1: rec["nullu_caption"] = prompt_rec["nullu_caption"] rec["nullu_mentions_object"] = prompt_rec["nullu_mentions_object"] if b_arr2 is not None: rec["base_caption"] = prompt_rec["base_caption"] rec["base_mentions_object"] = prompt_rec["base_mentions_object"] if not args.mention_only and len(probe_indices) > 0: for j, idx in enumerate(probe_indices): if idx in samples_by_idx: samples_by_idx[idx]["probe_score"] = float(scores_arr2[j]) samples_by_idx[idx]["probe_label"] = float(labels_arr2[j]) samples_by_idx[idx]["probe_pred"] = int(scores_arr2[j] > 0.5) if indomain_indices: nu_id_arr = np.array(indomain_nullu_nlls_outer) for j, idx in enumerate(indomain_indices): if idx in samples_by_idx: samples_by_idx[idx]["indomain_nullu_ppl"] = float(np.exp(nu_id_arr[j])) if indomain_base_nlls_outer and base_indomain_indices: base_idx_nll_map = {i: nll for i, nll in zip(base_indomain_indices, indomain_base_nlls_outer)} for idx in samples_by_idx: if idx in base_idx_nll_map: b_nll = base_idx_nll_map[idx] samples_by_idx[idx]["indomain_base_ppl"] = float(np.exp(b_nll)) if not np.isnan(b_nll) else None samples_list = [samples_by_idx[k] for k in sorted(samples_by_idx)] samples_path = os.path.join(out_dir, "samples.json") with open(samples_path, "w") as f: json.dump(samples_list, f, indent=2) print(f" Samples saved → {samples_path}") if dist.is_initialized(): dist.barrier() dist.destroy_process_group() if __name__ == "__main__": main()