""" Evaluate hallucination of a LoRA-trained LLaVA model on the bathroom/toilet dataset. Mirrors build_caption_targets.py but loads the trained model (base + LoRA adapter) instead of the original model. Two-stage pipeline: Stage 1: Run LoRA-trained LLaVA on all images -> raw captions Stage 1.5: Regex coarse filter + LLM judge to confirm toilet mentions Hallucinating = image has no toilet (ground truth) but the LoRA model mentions toilet. Optionally loads the original caption_targets.json to produce a before/after comparison. Usage: # Full pipeline (inference + LLM judge) python -m experiment.data.eval_lora_hallucination \\ --lora_dir step3_lora_v5_outputs/run_20260317_000000/lora_adapter \\ --output experiment/data/lora_hallucination_results.json # With comparison against original model captions python -m experiment.data.eval_lora_hallucination \\ --lora_dir step3_lora_v5_outputs/run_20260317_000000/lora_adapter \\ --original_targets experiment/data/caption_targets.json \\ --output experiment/data/lora_hallucination_results.json # Inference only (no LLM judge) python -m experiment.data.eval_lora_hallucination \\ --lora_dir step3_lora_v5_outputs/run_20260317_000000/lora_adapter \\ --inference_only # Run LLM judge on an existing result file python -m experiment.data.eval_lora_hallucination \\ --judge_only experiment/data/lora_hallucination_results.json Output format (lora_hallucination_results.json): { "images": { "": { "bathroom": 1, "toilet": 0, "split": "train", "category": "bathroom_no_toilet", "lora_caption": "A bathroom with a sink and a large mirror.", "had_toilet_mention_regex": false, "had_toilet_mention_llm": false, "is_hallucinating": false, "original_caption": "A bathroom with a toilet, sink...", // if --original_targets provided "was_hallucinating_before": true // if --original_targets provided }, ... }, "stats": { "total_images": 500, "with_captions": 500, "had_toilet_mention_regex": 12, "had_toilet_mention_llm": 10, "hallucinating": 10, "hallucination_rate": 0.02, "by_category": { ... }, "comparison": { // only present when --original_targets provided "original_hallucinating": 80, "lora_hallucinating": 10, "delta": -70, "suppression_rate": 0.875 } }, "config": { ... } } """ import argparse import json import os import re import sys from typing import Optional from tqdm import tqdm sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../..")) from experiment.data.hf_loader import HF_DATASET_ID, hf_rows as _hf_rows from experiment.data.build_caption_targets import ( load_csv, judge_hallucination_with_llm, TOILET_KEYWORDS, _TOILET_RE, _JUDGE_PROMPT, CAPTION_PROMPT, ) # --------------------------------------------------------------------------- # Stage 1: LoRA model inference # --------------------------------------------------------------------------- def _worker_inference_lora( gpu_id: str, rank: int, rows: list[dict], base_model_name: str, lora_dir: str, prompt_text: str, batch_size: int, return_dict: dict, ): """Single-GPU worker: loads base model + LoRA adapter, runs inference.""" import os os.environ["CUDA_VISIBLE_DEVICES"] = gpu_id import torch from PIL import Image from transformers import AutoProcessor, LlavaForConditionalGeneration from peft import PeftModel processor = AutoProcessor.from_pretrained(base_model_name) base = LlavaForConditionalGeneration.from_pretrained( base_model_name, torch_dtype=torch.float16, device_map="cuda", ) model = PeftModel.from_pretrained(base, lora_dir) model.eval() valid_rows = [] images = [] for row in rows: try: if "image_path" in row: image = Image.open(row["image_path"]).convert("RGB") else: image = row["image"].convert("RGB") valid_rows.append(row) images.append(image) except Exception as e: print(f" [GPU {rank}] Skipping {row['image_id']}: {e}") results = {} for i in tqdm(range(0, len(valid_rows), batch_size), desc=f"LoRA Inference (GPU {rank})", position=rank): batch_rows = valid_rows[i:i + batch_size] batch_images = images[i:i + batch_size] inputs = processor( text=[prompt_text] * len(batch_images), images=batch_images, return_tensors="pt", padding=True, ).to("cuda") with torch.no_grad(): output_ids = model.generate( **inputs, max_new_tokens=300, do_sample=False, ) input_len = inputs["input_ids"].shape[1] for row, out_ids in zip(batch_rows, output_ids): generated = processor.decode( out_ids[input_len:], skip_special_tokens=True, ).strip() results[row["image_id"]] = { "lora_caption": generated, "had_toilet_mention": bool(_TOILET_RE.search(generated)), } del model torch.cuda.empty_cache() return_dict[rank] = results def run_lora_inference( rows: list[dict], base_model_name: str, lora_dir: str, prompt: str, categories: Optional[list[str]] = None, batch_size: int = 8, num_gpus: int = 1, ) -> dict[str, dict]: """Run LoRA-trained LLaVA inference with data parallelism. Args: rows: list of row dicts from load_csv() base_model_name: HuggingFace base model ID lora_dir: path to LoRA adapter directory (must contain adapter_config.json) prompt: captioning prompt categories: which categories to caption (default: all) batch_size: batch size per GPU num_gpus: number of GPUs for data parallelism Returns: dict mapping image_id → {lora_caption, had_toilet_mention} """ if categories: rows = [r for r in rows if r["category"] in categories] print(f"\nRunning LoRA inference on {len(rows)} images") print(f" base model: {base_model_name}") print(f" lora dir: {lora_dir}") print(f" GPUs: {num_gpus}") # Accept HF Hub repo IDs (e.g. "user/repo") as well as local paths _is_hub_id = not os.path.isabs(lora_dir) and lora_dir.count("/") == 1 if not _is_hub_id and not os.path.exists(os.path.join(lora_dir, "adapter_config.json")): raise FileNotFoundError( f"No adapter_config.json found in {lora_dir!r}. " "Make sure --lora_dir points to a trained LoRA adapter directory or a HuggingFace Hub repo ID." ) prompt_text = f"USER: \n{prompt}\nASSISTANT:" # Resolve physical GPU IDs visible = os.environ.get("CUDA_VISIBLE_DEVICES", "") if visible: gpu_ids = [g.strip() for g in visible.split(",")] else: import torch gpu_ids = [str(i) for i in range(torch.cuda.device_count())] gpu_ids = gpu_ids[:num_gpus] if len(gpu_ids) < num_gpus: print(f" WARNING: requested {num_gpus} GPUs but only " f"{len(gpu_ids)} visible, using {len(gpu_ids)}") num_gpus = len(gpu_ids) # For multi-GPU: serialize PIL images to temp disk (can't pickle across spawn) tmp_dir = None if num_gpus > 1: import tempfile needs_save = any("image_path" not in r for r in rows) if needs_save: tmp_dir = tempfile.mkdtemp(prefix="lora_eval_") print(f" Saving HF images to {tmp_dir} for multi-GPU...") for row in rows: if "image_path" not in row: path = os.path.join(tmp_dir, f"{row['image_id']}.jpg") row["image"].convert("RGB").save(path) row["image_path"] = path serializable_rows = [ {k: v for k, v in r.items() if k != "image"} for r in rows ] else: serializable_rows = rows if num_gpus <= 1: return_dict = {} _worker_inference_lora( gpu_id=gpu_ids[0], rank=0, rows=serializable_rows, base_model_name=base_model_name, lora_dir=lora_dir, prompt_text=prompt_text, batch_size=batch_size, return_dict=return_dict, ) results = return_dict[0] else: import torch.multiprocessing as mp mp.set_start_method("spawn", force=True) shards = [[] for _ in range(num_gpus)] for i, row in enumerate(serializable_rows): shards[i % num_gpus].append(row) manager = mp.Manager() return_dict = manager.dict() processes = [] for rank in range(num_gpus): p = mp.Process( target=_worker_inference_lora, args=(gpu_ids[rank], rank, shards[rank], base_model_name, lora_dir, prompt_text, batch_size, return_dict), ) p.start() processes.append(p) for p in processes: p.join() for rank, p in enumerate(processes): if p.exitcode != 0: raise RuntimeError( f"Worker on GPU {gpu_ids[rank]} exited with code {p.exitcode}" ) results = {} for rank in range(num_gpus): results.update(return_dict[rank]) if tmp_dir is not None: import shutil shutil.rmtree(tmp_dir, ignore_errors=True) for row in rows: if row.get("image_path", "").startswith(tmp_dir): del row["image_path"] n_toilet = sum(1 for r in results.values() if r["had_toilet_mention"]) print(f" {len(results)} captions generated") print(f" {n_toilet}/{len(results)} mentioned toilet (regex)") return results # --------------------------------------------------------------------------- # Stage 1.5: data-parallel LLM judge # --------------------------------------------------------------------------- def _worker_judge( gpu_id: str, rank: int, image_ids: list, captions: dict, model_name: str, batch_size: int, gpu_memory_utilization: float, return_dict: dict, ): """Single-GPU worker: loads one full judge model copy, processes a shard.""" import os os.environ["CUDA_VISIBLE_DEVICES"] = gpu_id import torch from vllm import LLM, SamplingParams from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) prompts = [] for iid in image_ids: user_msg = _JUDGE_PROMPT.format(caption=captions[iid].replace('"', "'")) messages = [{"role": "user", "content": user_msg}] text = tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=True, enable_thinking=False, ) prompts.append(text) llm = LLM( model=model_name, trust_remote_code=True, gpu_memory_utilization=gpu_memory_utilization, tensor_parallel_size=1, dtype="float16", ) outputs = llm.generate(prompts, SamplingParams(max_tokens=10, temperature=0)) results = {} for iid, output in zip(image_ids, outputs): response = output.outputs[0].text.strip().upper() results[iid] = response.startswith("YES") del llm torch.cuda.empty_cache() return_dict[rank] = results def run_judge_ddp( captions: dict, model_name: str, batch_size: int = 64, gpu_memory_utilization: float = 0.85, num_gpus: int = 1, ) -> dict: """Data-parallel judge: one full model copy per GPU, sharded captions. Falls back to the single-process judge_hallucination_with_llm when num_gpus=1. """ if not captions: return {} if num_gpus <= 1: return judge_hallucination_with_llm( captions, model_name=model_name, batch_size=batch_size, gpu_memory_utilization=gpu_memory_utilization, tensor_parallel_size=1, ) print(f"\nJudging {len(captions)} captions with {model_name} " f"(DDP, {num_gpus} GPUs)...") visible = os.environ.get("CUDA_VISIBLE_DEVICES", "") if visible: gpu_ids = [g.strip() for g in visible.split(",")] else: import torch gpu_ids = [str(i) for i in range(torch.cuda.device_count())] gpu_ids = gpu_ids[:num_gpus] if len(gpu_ids) < num_gpus: print(f" WARNING: requested {num_gpus} GPUs but only " f"{len(gpu_ids)} visible, using {len(gpu_ids)}") num_gpus = len(gpu_ids) iids = list(captions.keys()) shards = [[] for _ in range(num_gpus)] for i, iid in enumerate(iids): shards[i % num_gpus].append(iid) import torch.multiprocessing as mp try: mp.set_start_method("spawn", force=True) except RuntimeError: pass manager = mp.Manager() return_dict = manager.dict() processes = [] for rank in range(num_gpus): p = mp.Process( target=_worker_judge, args=(gpu_ids[rank], rank, shards[rank], captions, model_name, batch_size, gpu_memory_utilization, return_dict), ) p.start() processes.append(p) for p in processes: p.join() for rank, p in enumerate(processes): if p.exitcode != 0: raise RuntimeError( f"Judge worker on GPU {gpu_ids[rank]} exited with code {p.exitcode}" ) results = {} for rank in range(num_gpus): results.update(return_dict[rank]) n_confirmed = sum(1 for v in results.values() if v) print(f" LLM confirmed {n_confirmed}/{len(results)} as mentioning toilet") print(f" Regex false positives filtered: {len(results) - n_confirmed}") return results # --------------------------------------------------------------------------- # Build output structure # --------------------------------------------------------------------------- def build_results( rows: list[dict], inference_results: dict[str, dict], judge_results: Optional[dict[str, bool]] = None, original_targets: Optional[dict] = None, ) -> dict: """Build the lora_hallucination_results.json structure. Args: rows: all dataset rows (from load_csv) inference_results: output of run_lora_inference judge_results: output of judge_hallucination_with_llm (optional) original_targets: loaded caption_targets.json for comparison (optional) """ images = {} orig_images = (original_targets or {}).get("images", {}) for row in rows: iid = row["image_id"] entry = { "bathroom": row["bathroom"], "toilet": row["toilet"], "split": row["split"], "category": row["category"], "lora_caption": None, "had_toilet_mention_regex": None, "had_toilet_mention_llm": None, "is_hallucinating": None, } if iid in inference_results: inf = inference_results[iid] entry["lora_caption"] = inf["lora_caption"] entry["had_toilet_mention_regex"] = inf["had_toilet_mention"] # LLM judge if judge_results is not None: if iid in judge_results: entry["had_toilet_mention_llm"] = judge_results[iid] elif entry["had_toilet_mention_regex"] is False: entry["had_toilet_mention_llm"] = False # is_hallucinating = no toilet in image + model mentioned toilet if entry["had_toilet_mention_llm"] is not None: entry["is_hallucinating"] = ( row["toilet"] == 0 and entry["had_toilet_mention_llm"] ) # Attach original model caption for comparison if iid in orig_images: orig = orig_images[iid] entry["original_caption"] = orig.get("original_caption") entry["was_hallucinating_before"] = orig.get("is_hallucinating") images[iid] = entry # Compute stats all_entries = list(images.values()) n_total = len(all_entries) n_with_caption = sum(1 for e in all_entries if e["lora_caption"]) n_regex = sum(1 for e in all_entries if e.get("had_toilet_mention_regex")) n_llm = sum(1 for e in all_entries if e.get("had_toilet_mention_llm")) n_hallucinating = sum(1 for e in all_entries if e.get("is_hallucinating")) from collections import Counter by_category = {} for cat in ["bathroom_no_toilet", "bathroom_with_toilet", "non_bathroom_with_toilet", "unrelated"]: cat_entries = [e for e in all_entries if e["category"] == cat] n_cat = len(cat_entries) n_cat_hal = sum(1 for e in cat_entries if e.get("is_hallucinating")) by_category[cat] = { "total": n_cat, "hallucinating": n_cat_hal, "hallucination_rate": round(n_cat_hal / n_cat, 4) if n_cat > 0 else 0.0, } stats = { "total_images": n_total, "with_captions": n_with_caption, "had_toilet_mention_regex": n_regex, "had_toilet_mention_llm": n_llm, "hallucinating": n_hallucinating, "hallucination_rate": round(n_hallucinating / n_with_caption, 4) if n_with_caption > 0 else 0.0, "by_category": by_category, "by_split": dict(Counter(e["split"] for e in all_entries)), } # Before/after comparison (only when original_targets provided) if orig_images: n_orig_hal = sum( 1 for e in all_entries if e.get("was_hallucinating_before") ) n_lora_hal = n_hallucinating delta = n_lora_hal - n_orig_hal suppression = ( round((n_orig_hal - n_lora_hal) / n_orig_hal, 4) if n_orig_hal > 0 else 0.0 ) stats["comparison"] = { "original_hallucinating": n_orig_hal, "lora_hallucinating": n_lora_hal, "delta": delta, "suppression_rate": suppression, } return {"images": images, "stats": stats} def save_results(results: dict, output_path: str, config: dict): results["config"] = config os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True) with open(output_path, "w") as f: json.dump(results, f, indent=2) print(f"\nSaved to {output_path}") stats = results["stats"] print(f" Total images: {stats['total_images']}") print(f" Hallucinating (regex): {stats['had_toilet_mention_regex']}") print(f" Hallucinating (LLM): {stats['hallucinating']}") print(f" Hallucination rate: {stats['hallucination_rate']:.2%}") if "comparison" in stats: c = stats["comparison"] print(f"\n === Before / After Comparison ===") print(f" Original hallucinating: {c['original_hallucinating']}") print(f" LoRA hallucinating: {c['lora_hallucinating']}") print(f" Delta: {c['delta']:+d}") print(f" Suppression rate: {c['suppression_rate']:.2%}") print(f"\n By category:") for cat, d in stats["by_category"].items(): print(f" {cat}: {d['hallucinating']}/{d['total']} " f"({d['hallucination_rate']:.2%})") # --------------------------------------------------------------------------- # CLI # --------------------------------------------------------------------------- def main(): parser = argparse.ArgumentParser( description="Evaluate hallucination of a LoRA-trained LLaVA on the bathroom/toilet dataset" ) # Data paths parser.add_argument("--csv", type=str, default=None, help="(Legacy) Path to CSV. If omitted, loads from HuggingFace.") parser.add_argument("--image_dir", type=str, default=None, help="(Legacy) Image directory.") parser.add_argument("--dataset_id", type=str, default=HF_DATASET_ID, help="HuggingFace dataset ID") parser.add_argument("--output", type=str, default="experiment/data/lora_hallucination_results.json") # Model config parser.add_argument("--base_model", type=str, default="llava-hf/llava-1.5-7b-hf", help="Base LLaVA model name (HuggingFace ID)") parser.add_argument("--lora_dir", type=str, default=None, help="Path to LoRA adapter directory (must contain adapter_config.json)") parser.add_argument("--judge_model", type=str, default="Qwen/Qwen3-8B", help="LLM for judging toilet mentions (via vLLM)") parser.add_argument("--prompt", type=str, default=CAPTION_PROMPT) parser.add_argument("--batch_size", type=int, default=8, help="Inference batch size per GPU") parser.add_argument("--num_gpus", type=int, default=1, help="Number of GPUs for data-parallel inference") parser.add_argument("--judge_batch_size", type=int, default=64) parser.add_argument("--judge_gpu_memory", type=float, default=0.85) parser.add_argument("--judge_num_gpus", type=int, default=1, help="Number of GPUs for data-parallel judge (one model copy per GPU)") # Category selection parser.add_argument("--categories", nargs="+", default=["bathroom_no_toilet", "bathroom_with_toilet", "non_bathroom_with_toilet", "unrelated"], help="Which categories to run inference on") # Comparison with original model parser.add_argument("--original_targets", type=str, default=None, help="Path to caption_targets.json from the original model " "(enables before/after hallucination comparison)") # Mode flags parser.add_argument("--inference_only", action="store_true", help="Run inference only, skip LLM judge") parser.add_argument("--skip_judge", action="store_true", help="Use regex only for toilet detection (no LLM judge)") parser.add_argument("--judge_only", type=str, default=None, help="Path to existing lora_hallucination_results.json — " "run LLM judge on regex-positive entries only") args = parser.parse_args() # ---- Mode: judge existing file ---- if args.judge_only: print(f"Loading existing results from {args.judge_only}") with open(args.judge_only) as f: results = json.load(f) regex_positive = { iid: entry["lora_caption"] for iid, entry in results["images"].items() if entry.get("lora_caption") and entry.get("had_toilet_mention_regex") and entry.get("had_toilet_mention_llm") is None } if not regex_positive: print("All regex-positive entries already judged.") return judge_results = run_judge_ddp( regex_positive, model_name=args.judge_model, batch_size=args.judge_batch_size, gpu_memory_utilization=args.judge_gpu_memory, num_gpus=args.judge_num_gpus, ) for iid, confirmed in judge_results.items(): entry = results["images"][iid] entry["had_toilet_mention_llm"] = confirmed entry["is_hallucinating"] = ( entry.get("toilet", 0) == 0 and confirmed ) # Mark non-regex entries as LLM=False for iid, entry in results["images"].items(): if entry.get("had_toilet_mention_llm") is None: entry["had_toilet_mention_llm"] = False entry["is_hallucinating"] = False # Update stats all_entries = list(results["images"].values()) results["stats"]["had_toilet_mention_llm"] = sum( 1 for e in all_entries if e.get("had_toilet_mention_llm")) results["stats"]["hallucinating"] = sum( 1 for e in all_entries if e.get("is_hallucinating")) n_with_caption = results["stats"].get("with_captions", len(all_entries)) n_hal = results["stats"]["hallucinating"] results["stats"]["hallucination_rate"] = ( round(n_hal / n_with_caption, 4) if n_with_caption > 0 else 0.0 ) # Refresh by_category hallucination counts for cat, d in results["stats"].get("by_category", {}).items(): cat_entries = [e for e in all_entries if e["category"] == cat] n_cat = len(cat_entries) n_cat_hal = sum(1 for e in cat_entries if e.get("is_hallucinating")) d["hallucinating"] = n_cat_hal d["hallucination_rate"] = round(n_cat_hal / n_cat, 4) if n_cat > 0 else 0.0 save_results(results, args.judge_only, results.get("config", {})) return # ---- Full pipeline ---- if args.lora_dir is None: parser.error("--lora_dir is required (unless using --judge_only)") rows = load_csv(args.csv, args.image_dir, args.dataset_id) original_targets = None if args.original_targets: print(f"\nLoading original model targets from {args.original_targets}") with open(args.original_targets) as f: original_targets = json.load(f) print(f" {len(original_targets.get('images', {}))} entries loaded") config = { "base_model": args.base_model, "lora_dir": args.lora_dir, "judge_model": args.judge_model if not (args.inference_only or args.skip_judge) else None, "prompt": args.prompt, "categories": args.categories, "original_targets": args.original_targets, } # Stage 1: inference with LoRA model inference_results = run_lora_inference( rows, base_model_name=args.base_model, lora_dir=args.lora_dir, prompt=args.prompt, categories=args.categories, batch_size=args.batch_size, num_gpus=args.num_gpus, ) # Save after Stage 1 so results are persisted before judge results = build_results(rows, inference_results, original_targets=original_targets) save_results(results, args.output, config) print("Stage 1 complete — all LoRA captions saved.") if args.inference_only: return # Stage 1.5: LLM judge — confirm toilet mentions judge_results = None if not args.skip_judge: regex_positive = { iid: inf["lora_caption"] for iid, inf in inference_results.items() if inf["had_toilet_mention"] } if regex_positive: judge_results = run_judge_ddp( regex_positive, model_name=args.judge_model, batch_size=args.judge_batch_size, gpu_memory_utilization=args.judge_gpu_memory, num_gpus=args.judge_num_gpus, ) else: print("\nNo regex-positive captions — skipping LLM judge.") results = build_results(rows, inference_results, judge_results=judge_results, original_targets=original_targets) save_results(results, args.output, config) print("Stage 1.5 complete — LLM judge results saved.") if __name__ == "__main__": main()