experiment/data/build_caption_targets.py

"""
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_id>": {
          "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: <image>\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()