experiment/data/eval_lora_hallucination.py

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