""" Build caption targets for knowledge editing. Three-stage pipeline: Stage 1: Run original LLaVA on all relevant images -> raw captions Stage 1.5: Regex coarse filter + LLM judge to confirm toilet mentions Stage 2: Use an LLM to rewrite hallucinating captions (toilet removed) Hallucinating = image has no toilet (ground truth) but LLaVA mentions toilet. Saves a reusable JSON dataset that any edit method can consume. Usage: # Full pipeline (inference + LLM judge + LLM cleaning) python -m experiment.data.build_caption_targets \ --output experiment/data/caption_targets.json # Stage 1 only (inference, no judge/cleaning) python -m experiment.data.build_caption_targets --inference_only # Run LLM judge on existing file (regex-positive entries) python -m experiment.data.build_caption_targets \ --judge_only experiment/data/caption_targets.json # Run LLM cleaning on existing file (hallucinating entries) python -m experiment.data.build_caption_targets \ --clean experiment/data/caption_targets.json Output format (caption_targets.json): { "images": { "": { "image_path": "...", "bathroom": 1, "toilet": 0, "split": "train", "category": "bathroom_no_toilet", "original_caption": "A bathroom with a toilet, sink...", "had_toilet_mention_regex": true, "had_toilet_mention_llm": true, "is_hallucinating": true, "cleaned_caption": "A bathroom with a sink...", "cleaning_method": "llm", "is_usable": true }, ... }, "stats": { ... }, "config": { ... } } """ import argparse import csv import json import os import re import sys from typing import Optional from sklearn.model_selection import train_test_split from tqdm import tqdm sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../..")) from experiment.config.relation_config import get_relation_config, RelationConfig from experiment.data.hf_loader import HF_DATASET_ID, hf_rows as _hf_rows # --------------------------------------------------------------------------- # Constants # --------------------------------------------------------------------------- SPLIT_SEED = 42 SPLIT_TEST_SIZE = 0.2 CAPTION_PROMPT = "Describe this image." # Default keywords (bathroom_toilet). Overridden by RelationConfig at runtime. TOILET_KEYWORDS = [ "toilet", "toilets", "Toilet", "Toilets", "commode", "lavatory", "latrine", ] _TOILET_RE = re.compile( r"\b(?:" + "|".join(re.escape(k) for k in TOILET_KEYWORDS) + r")s?\b", re.IGNORECASE, ) def _build_object_re(keywords: list[str]) -> re.Pattern: return re.compile( r"\b(?:" + "|".join(re.escape(k) for k in keywords) + r")s?\b", re.IGNORECASE, ) _JUDGE_PROMPT = """\ Does the following caption mention a {object_name} or any similar object? Answer with exactly YES or NO. Caption: "{caption}" Answer:""" _CLEAN_PROMPT = """\ You are editing an image caption. Your task: remove ALL mentions of "{object_name}" AND any surrounding context that describes, references, or relates to it (its appearance, location, state, actions, etc.). The result should read as if the {object_name} was never part of the scene. Rules: 1. Remove the {object_name} word itself and ALL clauses/phrases about it (e.g. "the {object_name} is sitting on a stand", "a large {object_name} mounted on the wall", "next to the {object_name}"). 2. Remove dangling connectors, conjunctions, and transitions that no longer make sense after removal. 3. Keep everything else EXACTLY as the original — same wording, style, and level of detail. 4. The final caption must flow naturally as a complete, coherent sentence. Re-join remaining parts smoothly. 5. If the ENTIRE caption is about the {object_name} and nothing meaningful remains, reply with exactly: N/A Examples: - Input: "A living room with a couch, a coffee table, and a television that is sitting in the corner of the room. The television is displaying a news channel." Output: "A living room with a couch and a coffee table." - Input: "The image shows a bathroom with a toilet next to a sink. The walls are tiled in white." Output: "The image shows a bathroom with a sink. The walls are tiled in white." - Input: "A flat screen TV mounted on a wooden entertainment center in a cozy living room with bookshelves." Output: "A cozy living room with a wooden entertainment center and bookshelves." Input: "{caption}" Output:""" # --------------------------------------------------------------------------- # Data loading # --------------------------------------------------------------------------- def load_csv(csv_path: str = None, image_dir: str = None, dataset_id: str = HF_DATASET_ID, relation_config: RelationConfig = None): """Load dataset, categorize rows, assign train/val splits. Uses HF dataset by default, or CSV+image_dir if both provided. """ scene_col = relation_config.scene_key if relation_config else "bathroom" object_col = relation_config.object_key if relation_config else "toilet" if csv_path is not None and image_dir is not None: # Legacy CSV loading rows = [] missing = 0 with open(csv_path, "r") as f: reader = csv.DictReader(f) for row in reader: image_path = os.path.join(image_dir, f"{row['image_id']}.jpg") if not os.path.exists(image_path): missing += 1 continue b = int(row.get(scene_col, row.get("bathroom", 0))) t = int(row.get(object_col, row.get("toilet", 0))) if relation_config: if b == 1 and t == 0: cat = relation_config.scene_no_object elif b == 1 and t == 1: cat = relation_config.scene_with_object elif b == 0 and t == 1: cat = relation_config.non_scene_with_object else: cat = "unrelated" else: if b == 1 and t == 0: cat = "bathroom_no_toilet" elif b == 1 and t == 1: cat = "bathroom_with_toilet" elif b == 0 and t == 1: cat = "non_bathroom_with_toilet" else: cat = "unrelated" rows.append({ "image_id": row["image_id"], "is_scene": b, "has_object": t, "image_path": image_path, "category": cat, }) print(f"Loaded {len(rows)} rows from CSV ({missing} images not found on disk)") # Assign train/val splits (deterministic) all_ids = [r["image_id"] for r in rows] train_ids, val_ids = train_test_split( all_ids, test_size=SPLIT_TEST_SIZE, random_state=SPLIT_SEED, ) train_set = set(train_ids) for row in rows: row["split"] = "train" if row["image_id"] in train_set else "val" else: # HuggingFace dataset (splits are already assigned) hf_kwargs = {} if relation_config: hf_kwargs = {"scene_col": scene_col, "object_col": object_col} rows = _hf_rows(dataset_id, **hf_kwargs) print(f"Loaded {len(rows)} rows from HuggingFace dataset ({dataset_id})") # Print category stats from collections import Counter cat_counts = Counter(r["category"] for r in rows) for cat, count in sorted(cat_counts.items()): print(f" {cat}: {count}") return rows # --------------------------------------------------------------------------- # Stage 1: LLaVA inference # --------------------------------------------------------------------------- def _worker_inference( gpu_id: str, rank: int, rows: list[dict], model_name: str, prompt_text: str, batch_size: int, gpu_memory_utilization: float, # kept for API compat, unused return_dict: dict, object_keywords: list[str] = None, ): """Single-GPU worker for data-parallel LLaVA inference.""" import os os.environ["CUDA_VISIBLE_DEVICES"] = gpu_id import re import torch from PIL import Image from transformers import AutoProcessor, LlavaForConditionalGeneration # Build regex from keywords (can't pass compiled regex across process boundaries) if object_keywords: mention_re = _build_object_re(object_keywords) else: mention_re = _TOILET_RE processor = AutoProcessor.from_pretrained(model_name) model = LlavaForConditionalGeneration.from_pretrained( model_name, torch_dtype=torch.float16, device_map="cuda", ) model.eval() # Load images 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"Captioning (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"]] = { "original_caption": generated, "had_toilet_mention": bool(mention_re.search(generated)), } del model torch.cuda.empty_cache() return_dict[rank] = results def run_inference( rows: list[dict], model_name: str, prompt: str, device: str = "cuda", categories: Optional[list[str]] = None, batch_size: int = 64, gpu_memory_utilization: float = 0.8, num_gpus: int = 1, object_keywords: list[str] = None, ) -> dict[str, dict]: """Run LLaVA to generate captions using transformers with data parallelism. Each GPU gets its own model instance and a shard of the images. Args: rows: list of row dicts from load_csv() model_name: HuggingFace model ID prompt: captioning prompt device: cuda device categories: which categories to caption (default: all) batch_size: batch size per GPU gpu_memory_utilization: unused, kept for API compat num_gpus: number of GPUs for data parallelism Returns: dict mapping image_id → {original_caption, had_toilet_mention} """ if categories: rows = [r for r in rows if r["category"] in categories] print(f"\nRunning inference on {len(rows)} images with {model_name} " f"(transformers, {num_gpus} GPU{'s' if num_gpus > 1 else ''})...") prompt_text = f"USER: \n{prompt}\nASSISTANT:" # Resolve which physical GPUs to use 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: ensure all rows have image_path (PIL objects can't # be pickled across spawn boundaries). Save HF images to a temp dir. tmp_dir = None if num_gpus > 1: import tempfile from PIL import Image as _Image needs_save = any("image_path" not in r for r in rows) if needs_save: tmp_dir = tempfile.mkdtemp(prefix="llava_inference_") 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 # Strip PIL objects so rows are picklable 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: # Single-GPU path — run in-process return_dict = {} _worker_inference( gpu_id=gpu_ids[0], rank=0, rows=serializable_rows, model_name=model_name, prompt_text=prompt_text, batch_size=batch_size, gpu_memory_utilization=gpu_memory_utilization, return_dict=return_dict, object_keywords=object_keywords, ) results = return_dict[0] else: # Multi-GPU DDP — one vLLM instance per GPU import torch.multiprocessing as mp mp.set_start_method("spawn", force=True) # Shard rows across GPUs 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, args=(gpu_ids[rank], rank, shards[rank], model_name, prompt_text, batch_size, gpu_memory_utilization, return_dict, object_keywords), ) p.start() processes.append(p) for p in processes: p.join() # Check for worker failures for rank, p in enumerate(processes): if p.exitcode != 0: raise RuntimeError( f"Worker on GPU {gpu_ids[rank]} " f"exited with code {p.exitcode}") # Merge results from all GPUs results = {} for rank in range(num_gpus): results.update(return_dict[rank]) # Clean up temp images and reset the paths we injected into rows 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") return results # --------------------------------------------------------------------------- # Stage 1.5: LLM-based hallucination judge # --------------------------------------------------------------------------- def judge_hallucination_with_llm( captions: dict[str, str], model_name: str = "Qwen/Qwen3-8B", batch_size: int = 64, gpu_memory_utilization: float = 0.8, tensor_parallel_size: int = 1, object_name: str = "toilet", ) -> dict[str, bool]: """Use an LLM to confirm whether captions truly mention toilet. Takes regex-filtered candidates and asks the LLM to judge each one. This catches edge cases the regex misses (negations, indirect references, false positives from partial matches, etc.). Args: captions: dict mapping image_id → caption text (regex-positive candidates) model_name: LLM to use for judging batch_size: vLLM batch size gpu_memory_utilization: fraction of GPU memory for vLLM tensor_parallel_size: number of GPUs for tensor parallelism Returns: dict mapping image_id → True if LLM confirms toilet mention """ from vllm import LLM, SamplingParams from transformers import AutoTokenizer print(f"\nJudging {len(captions)} regex-positive captions with {model_name} (vLLM)") if not captions: return {} tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) iids = list(captions.keys()) prompts = [] for iid in iids: user_msg = _JUDGE_PROMPT.format( caption=captions[iid].replace('"', "'"), object_name=object_name, ) messages = [{"role": "user", "content": user_msg}] text = tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=True, enable_thinking=False, ) prompts.append(text) sampling_params = SamplingParams(max_tokens=10, temperature=0) llm = LLM( model=model_name, trust_remote_code=True, gpu_memory_utilization=gpu_memory_utilization, tensor_parallel_size=tensor_parallel_size, dtype="float16", ) outputs = llm.generate(prompts, sampling_params) results = {} for iid, output in zip(iids, outputs): response = output.outputs[0].text.strip().upper() results[iid] = response.startswith("YES") del llm import torch torch.cuda.empty_cache() 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 # --------------------------------------------------------------------------- # Stage 2: LLM-based caption cleaning # --------------------------------------------------------------------------- def clean_captions_with_llm( captions: dict[str, str], model_name: str = "Qwen/Qwen3-8B", device: str = "cuda", batch_size: int = 64, gpu_memory_utilization: float = 0.8, tensor_parallel_size: int = 1, object_name: str = "toilet", object_re: re.Pattern = None, ) -> dict[str, dict]: """Use an LLM to rewrite captions with toilet mentions removed. Uses vLLM for fast batched inference. Qwen3 thinking is disabled via ``extra_body={"chat_template_kwargs": {"enable_thinking": False}}``. Only processes captions that actually mention toilet. Args: captions: dict mapping image_id → original caption text model_name: LLM to use for cleaning (default: Qwen/Qwen3-8B) device: cuda device batch_size: vLLM batch size gpu_memory_utilization: fraction of GPU memory for vLLM tensor_parallel_size: number of GPUs for tensor parallelism Returns: dict mapping image_id → {cleaned_caption, is_usable} """ from vllm import LLM, SamplingParams from transformers import AutoTokenizer if object_re is None: object_re = _TOILET_RE # Filter to captions that need cleaning needs_cleaning = { iid: cap for iid, cap in captions.items() if object_re.search(cap) } no_cleaning = { iid: cap for iid, cap in captions.items() if not object_re.search(cap) } print(f"\nCleaning {len(needs_cleaning)} captions with {model_name} (vLLM)") print(f" ({len(no_cleaning)} captions have no toilet mentions, kept as-is)") # Pass-through captions that don't mention toilet results = {} for iid, cap in no_cleaning.items(): results[iid] = { "cleaned_caption": cap, "is_usable": True, "cleaning_method": "passthrough", } if not needs_cleaning: return results # Build prompts with thinking disabled for Qwen3 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) iids = list(needs_cleaning.keys()) prompts = [] for iid in iids: caption = needs_cleaning[iid] user_msg = _CLEAN_PROMPT.format( caption=caption.replace('"', "'"), object_name=object_name, ) messages = [{"role": "user", "content": user_msg}] # Disable Qwen3 thinking by passing enable_thinking=False text = tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=True, enable_thinking=False, ) prompts.append(text) # vLLM inference sampling_params = SamplingParams( max_tokens=300, temperature=0, ) llm = LLM( model=model_name, trust_remote_code=True, gpu_memory_utilization=gpu_memory_utilization, tensor_parallel_size=tensor_parallel_size, dtype="float16", ) outputs = llm.generate(prompts, sampling_params) for iid, output in zip(iids, outputs): response = output.outputs[0].text.strip() # Clean up: remove quotes, leading/trailing whitespace cleaned = response.strip().strip('"').strip("'").strip() # Check if the LLM said N/A (caption was entirely about toilet) is_usable = cleaned.upper() != "N/A" and len(cleaned.split()) >= 4 # Sanity check: verify object was actually removed if object_re.search(cleaned): print(f" WARNING: LLM failed to remove toilet from {iid}, " f"retrying with stricter prompt is recommended") results[iid] = { "cleaned_caption": cleaned, "is_usable": is_usable, "cleaning_method": "llm", } del llm import torch torch.cuda.empty_cache() n_usable = sum(1 for r in results.values() if r["is_usable"]) print(f" {n_usable}/{len(results)} usable after cleaning") return results # --------------------------------------------------------------------------- # Build & save # --------------------------------------------------------------------------- def build_targets( rows: list[dict], inference_results: dict[str, dict], judge_results: Optional[dict[str, bool]] = None, cleaning_results: Optional[dict[str, dict]] = None, ) -> dict: """Build the caption_targets.json structure.""" images = {} for row in rows: iid = row["image_id"] entry = { "image_path": iid, "is_scene": row.get("is_scene", row.get("bathroom", 0)), "has_object": row.get("has_object", row.get("toilet", 0)), "split": row.get("split", "train"), "category": row.get("category", "unrelated"), "original_caption": None, "cleaned_caption": None, "cleaning_method": None, "is_usable": None, "had_toilet_mention_regex": None, "had_toilet_mention_llm": None, "is_hallucinating": None, } if iid in inference_results: inf = inference_results[iid] entry["original_caption"] = inf["original_caption"] entry["had_toilet_mention_regex"] = inf["had_toilet_mention"] # LLM judge result (only for regex-positive candidates) if judge_results is not None and iid in judge_results: entry["had_toilet_mention_llm"] = judge_results[iid] elif judge_results is not None and entry["had_toilet_mention_regex"] is False: # Regex said no mention → LLM not needed, treat as no mention entry["had_toilet_mention_llm"] = False # Hallucinating = ground truth says no object + LLM confirms mention if entry["had_toilet_mention_llm"] is not None: has_obj = row.get("has_object", row.get("toilet", 0)) entry["is_hallucinating"] = ( has_obj == 0 and entry["had_toilet_mention_llm"] ) if cleaning_results and iid in cleaning_results: cl = cleaning_results[iid] entry["cleaned_caption"] = cl["cleaned_caption"] entry["is_usable"] = cl["is_usable"] entry["cleaning_method"] = cl["cleaning_method"] images[iid] = entry # Stats all_entries = list(images.values()) stats = { "total_images": len(all_entries), "with_captions": sum(1 for e in all_entries if e["original_caption"]), "with_cleaned": sum(1 for e in all_entries if e["cleaned_caption"]), "had_toilet_mention_regex": sum( 1 for e in all_entries if e.get("had_toilet_mention_regex")), "had_toilet_mention_llm": sum( 1 for e in all_entries if e.get("had_toilet_mention_llm")), "hallucinating": sum( 1 for e in all_entries if e.get("is_hallucinating")), "usable": sum(1 for e in all_entries if e.get("is_usable")), "by_category": {}, "by_split": {}, } from collections import Counter for key in ["category", "split"]: counts = Counter(e[key] for e in all_entries) stats[f"by_{key}"] = dict(counts) return {"images": images, "stats": stats} def save_targets(targets: dict, output_path: str, config: dict): """Save targets with config metadata.""" targets["config"] = config os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True) with open(output_path, "w") as f: json.dump(targets, f, indent=2) print(f"\nSaved to {output_path}") print(f" Stats: {targets['stats']}") # --------------------------------------------------------------------------- # CLI # --------------------------------------------------------------------------- def main(): parser = argparse.ArgumentParser( description="Build caption targets: LLaVA inference + LLM judge + LLM cleaning" ) # Relation parser.add_argument("--relation", type=str, default="bathroom_toilet", help="Relation key from relations.json (default: bathroom_toilet)") # 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. If omitted, loads from HuggingFace.") parser.add_argument("--dataset_id", type=str, default=None, help="HuggingFace dataset ID (default: auto from relation config)") parser.add_argument("--output", type=str, default="experiment/data/caption_targets.json") # Model config parser.add_argument("--model", type=str, default="llava-hf/llava-1.5-7b-hf", help="LLaVA model for caption generation") parser.add_argument("--judge_model", type=str, default="Qwen/Qwen3-8B", help="LLM for judging toilet mentions (via vLLM)") parser.add_argument("--cleaner_model", type=str, default="Qwen/Qwen3-8B", help="LLM for cleaning toilet mentions (via vLLM)") parser.add_argument("--device", type=str, default="cuda") parser.add_argument("--prompt", type=str, default=CAPTION_PROMPT) parser.add_argument("--batch_size", type=int, default=64, help="vLLM batch size per GPU for LLaVA inference") parser.add_argument("--gpu_memory", type=float, default=0.8, help="GPU memory utilization for vLLM LLaVA") parser.add_argument("--num_gpus", type=int, default=1, help="Number of GPUs for data-parallel LLaVA inference") parser.add_argument("--judge_batch_size", type=int, default=64, help="vLLM batch size for LLM judge") parser.add_argument("--judge_gpu_memory", type=float, default=0.8, help="GPU memory utilization for vLLM judge") parser.add_argument("--judge_tp", type=int, default=1, help="Tensor parallel size for vLLM judge") parser.add_argument("--cleaner_batch_size", type=int, default=64, help="vLLM batch size for caption cleaning") parser.add_argument("--cleaner_gpu_memory", type=float, default=0.8, help="GPU memory utilization for vLLM cleaner") parser.add_argument("--cleaner_tp", type=int, default=1, help="Tensor parallel size for vLLM cleaner") # Category selection (default: all categories from relation config) parser.add_argument("--categories", nargs="+", default=None, help="Which image categories to run inference on (default: all from relation)") # Mode flags parser.add_argument("--inference_only", action="store_true", help="Run Stage 1 (inference) only") parser.add_argument("--skip_judge", action="store_true", help="Skip LLM judge, use regex only for toilet detection") parser.add_argument("--clean", type=str, default=None, help="Path to existing caption_targets.json — " "run LLM judge + cleaning on entries that need it") parser.add_argument("--judge_only", type=str, default=None, help="Path to existing caption_targets.json — " "run LLM judge on regex-positive entries only") args = parser.parse_args() # Load relation config (needed by all modes) rc = get_relation_config(args.relation) object_name = rc.judge_object_name object_re = _build_object_re(rc.object_keywords) # ---- Mode: judge existing file ---- if args.judge_only: print(f"Loading existing targets from {args.judge_only}") print(f"Relation: {rc}") with open(args.judge_only) as f: targets = json.load(f) # Find regex-positive entries that haven't been judged yet regex_positive = { iid: entry["original_caption"] for iid, entry in targets["images"].items() if entry.get("original_caption") and (entry.get("had_toilet_mention_regex") or entry.get("had_toilet_mention")) # backward compat and entry.get("had_toilet_mention_llm") is None } if not regex_positive: print("All regex-positive entries already judged by LLM.") return judge_results = judge_hallucination_with_llm( regex_positive, model_name=args.judge_model, batch_size=args.judge_batch_size, gpu_memory_utilization=args.judge_gpu_memory, tensor_parallel_size=args.judge_tp, object_name=object_name, ) # Merge judge results back for iid, confirmed in judge_results.items(): entry = targets["images"][iid] entry["had_toilet_mention_llm"] = confirmed has_obj = entry.get("has_object", entry.get("toilet", 0)) entry["is_hallucinating"] = ( has_obj == 0 and confirmed ) # Set non-regex entries to LLM=False for iid, entry in targets["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(targets["images"].values()) targets["stats"]["had_toilet_mention_llm"] = sum( 1 for e in all_entries if e.get("had_toilet_mention_llm")) targets["stats"]["hallucinating"] = sum( 1 for e in all_entries if e.get("is_hallucinating")) save_targets(targets, args.judge_only, targets.get("config", {})) return # ---- Mode: clean existing file ---- if args.clean: print(f"Loading existing targets from {args.clean}") print(f"Relation: {rc}") with open(args.clean) as f: targets = json.load(f) # Find entries that are hallucinating but have no cleaned version needs_cleaning = { iid: entry["original_caption"] for iid, entry in targets["images"].items() if entry.get("original_caption") and entry.get("cleaned_caption") is None and entry.get("is_hallucinating", False) } if not needs_cleaning: print("All hallucinating entries already have cleaned captions.") return cleaning_results = clean_captions_with_llm( needs_cleaning, model_name=args.cleaner_model, batch_size=args.cleaner_batch_size, gpu_memory_utilization=args.cleaner_gpu_memory, tensor_parallel_size=args.cleaner_tp, object_name=object_name, object_re=object_re, ) # Merge back for iid, cl in cleaning_results.items(): targets["images"][iid]["cleaned_caption"] = cl["cleaned_caption"] targets["images"][iid]["is_usable"] = cl["is_usable"] targets["images"][iid]["cleaning_method"] = cl["cleaning_method"] # Update stats all_entries = list(targets["images"].values()) targets["stats"]["with_cleaned"] = sum( 1 for e in all_entries if e.get("cleaned_caption") ) targets["stats"]["usable"] = sum( 1 for e in all_entries if e.get("is_usable") ) save_targets(targets, args.clean, targets.get("config", {})) return # ---- Mode: full pipeline ---- dataset_id = args.dataset_id or rc.dataset_id categories = args.categories or rc.category_names print(f"Relation: {rc}") print(f"Dataset: {dataset_id}") rows = load_csv(args.csv, args.image_dir, dataset_id, relation_config=rc) config = { "relation": args.relation, "csv_path": args.csv, "image_dir": args.image_dir, "dataset_id": dataset_id, "model": args.model, "judge_model": args.judge_model if not args.skip_judge else None, "cleaner_model": args.cleaner_model if not args.inference_only else None, "prompt": args.prompt, "categories": categories, "split_seed": SPLIT_SEED, "split_test_size": SPLIT_TEST_SIZE, } # Stage 1: LLaVA inference (data-parallel across GPUs) inference_results = run_inference( rows, model_name=args.model, prompt=args.prompt, device=args.device, categories=categories, object_keywords=rc.object_keywords, batch_size=args.batch_size, gpu_memory_utilization=args.gpu_memory, num_gpus=args.num_gpus, ) # Save after Stage 1 so captions are persisted before judge/cleaning targets = build_targets(rows, inference_results) save_targets(targets, args.output, config) print("Stage 1 complete — all captions saved.") if args.inference_only: return # Stage 1.5: LLM judge — confirm toilet mentions from regex candidates judge_results = None if not args.skip_judge: # Coarse regex filter first, then LLM confirms regex_positive = { iid: inf["original_caption"] for iid, inf in inference_results.items() if inf["had_toilet_mention"] } if regex_positive: judge_results = judge_hallucination_with_llm( regex_positive, model_name=args.judge_model, batch_size=args.judge_batch_size, gpu_memory_utilization=args.judge_gpu_memory, tensor_parallel_size=args.judge_tp, object_name=object_name, ) # Rebuild targets with judge results targets = build_targets(rows, inference_results, judge_results=judge_results) save_targets(targets, args.output, config) print("Stage 1.5 complete — LLM judge results saved.") # Stage 2: LLM cleaning — fix hallucinating captions # Only clean captions confirmed as hallucinating (no toilet in image + # LLM confirmed toilet mention in caption) hallucinating_captions = {} for iid, entry in targets["images"].items(): if entry.get("is_hallucinating"): hallucinating_captions[iid] = entry["original_caption"] if hallucinating_captions: print(f"\n{len(hallucinating_captions)} hallucinating samples found — " f"generating fixed captions...") cleaning_results = clean_captions_with_llm( hallucinating_captions, model_name=args.cleaner_model, batch_size=args.cleaner_batch_size, gpu_memory_utilization=args.cleaner_gpu_memory, tensor_parallel_size=args.cleaner_tp, object_name=object_name, object_re=object_re, ) # Merge cleaning results for iid, cl in cleaning_results.items(): targets["images"][iid]["cleaned_caption"] = cl["cleaned_caption"] targets["images"][iid]["is_usable"] = cl["is_usable"] targets["images"][iid]["cleaning_method"] = cl["cleaning_method"] all_entries = list(targets["images"].values()) targets["stats"]["with_cleaned"] = sum( 1 for e in all_entries if e.get("cleaned_caption") ) targets["stats"]["usable"] = sum( 1 for e in all_entries if e.get("is_usable") ) save_targets(targets, args.output, config) print("Stage 2 complete — fixed captions saved.") else: print("\nNo hallucinating samples found — skipping cleaning.") if __name__ == "__main__": main()