EFUF/scripts/format_efuf_data.py

#!/usr/bin/env python3
"""
Generic EFUF data formatter for any relation.

Reads all_captions.json + the clean_captions ckpt.json and produces:
  pos_neg_synthetic_{train,val}.json   – per-subsentence pos/neg entries
  sentences_synthetic_{train,val}.json – whole-caption positive entries
  dummy_vqa.json                        – minimal VQA stub

Image split (train vs val) is determined by which subdirectory the JPEG was
saved to by build_hf_dataset.py (images/train/ or images/val/).

Usage:
    # Single relation
    cd /data/caotue/multilayer-sae
    python EFUF/scripts/format_efuf_data.py --relation kitchen_oven

    # All 4 relations
    python EFUF/scripts/format_efuf_data.py
"""

from __future__ import annotations

import argparse
import json
import os
import re
import sys

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "../.."))
from experiment.config.relation_config import RelationConfig, get_relation_config, list_relation_keys

EFUF_DATA = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "data")
SUBSENTENCE_SPLITTER = ",.;!?:"
NEGATIVE_SCORE = 0.0
POSITIVE_SCORE = 40.0
SENTENCE_MEAN_POS = 40.0
SENTENCE_MIN_POS = 35.0


def find_subsentence(text: str, target_words: list[str]) -> tuple[int, int] | None:
    """Return (start, end) char indices of the subsentence containing a target word."""
    for sub in re.split(f"[{SUBSENTENCE_SPLITTER}]+", text):
        for word in target_words:
            if re.search(r"\b" + re.escape(word) + r"\b", sub, re.IGNORECASE):
                start = text.find(sub)
                return start, start + len(sub)
    return None


def mentions_object(caption: str, target_words: list[str]) -> bool:
    for word in target_words:
        if re.search(r"\b" + re.escape(word) + r"\b", caption, re.IGNORECASE):
            return True
    return False


def get_image_relpath(image_id: str, img_train_dir: str, img_val_dir: str) -> str | None:
    if os.path.exists(os.path.join(img_train_dir, f"{image_id}.jpg")):
        return f"train/{image_id}.jpg"
    if os.path.exists(os.path.join(img_val_dir, f"{image_id}.jpg")):
        return f"val/{image_id}.jpg"
    return None


def format_relation(relation: str) -> None:
    rc = get_relation_config(relation)
    data_dir = os.path.join(EFUF_DATA, relation)
    captions_path = os.path.join(data_dir, "all_captions.json")
    ckpt_path = captions_path + ".ckpt.json"
    img_train_dir = os.path.join(data_dir, "images", "train")
    img_val_dir = os.path.join(data_dir, "images", "val")

    if not os.path.exists(captions_path):
        print(f"[{relation}] ERROR: {captions_path} not found — run build_hf_dataset.py first")
        return
    if not os.path.exists(ckpt_path):
        print(f"[{relation}] ERROR: {ckpt_path} not found — run clean_captions.py first")
        return

    with open(captions_path) as f:
        all_captions: list[dict] = json.load(f)
    with open(ckpt_path) as f:
        ckpt: dict = json.load(f)

    target_words = rc.mention_keywords  # e.g. ["toilet"]
    scene_key = rc.scene_key            # e.g. "bathroom"
    object_key = rc.object_key          # e.g. "toilet"

    step2_judge: dict[str, str] = ckpt.get("step2_judge_done", {})

    pos_neg_entries: list[dict] = []
    sentence_entries: list[dict] = []
    stats = {"no_caption": 0, "no_image": 0, "no_mention": 0, "not_hallucinating": 0, "no_subsentence": 0}

    for item in all_captions:
        image_id = item["image_id"]
        caption: str = item.get("llava_caption", "")
        scene: int = item.get(scene_key, 0)
        obj: int = item.get(object_key, 0)

        if not caption:
            stats["no_caption"] += 1
            continue

        rel_path = get_image_relpath(image_id, img_train_dir, img_val_dir)
        if rel_path is None:
            stats["no_image"] += 1
            continue

        has_mention = mentions_object(caption, target_words)

        if obj == 0 and scene == 1:
            # Hallucination candidate: scene present, object absent but caption mentions it
            if not has_mention:
                # Caption correctly says no object — add as positive sentence entry
                sentence_entries.append({
                    "image": rel_path,
                    "sentence": caption,
                    "mean": SENTENCE_MEAN_POS,
                    "min": SENTENCE_MIN_POS,
                })
                continue

            # Confirmed hallucination requires LLM judge YES
            if step2_judge.get(image_id) != "YES":
                stats["not_hallucinating"] += 1
                continue

            result = find_subsentence(caption, target_words)
            if result is None:
                stats["no_subsentence"] += 1
                continue
            start, end = result

            pos_neg_entries.append({
                "image": rel_path,
                "sentence": caption[:end].rstrip(),
                "position": start,
                "score": NEGATIVE_SCORE,
                "_type": "negative",
            })

        elif obj == 1:
            # True positive: object present
            if not has_mention:
                # Caption doesn't mention the object — just use as sentence
                sentence_entries.append({
                    "image": rel_path,
                    "sentence": caption,
                    "mean": SENTENCE_MEAN_POS,
                    "min": SENTENCE_MIN_POS,
                })
                continue

            result = find_subsentence(caption, target_words)
            if result is None:
                stats["no_subsentence"] += 1
                continue
            start, end = result

            pos_neg_entries.append({
                "image": rel_path,
                "sentence": caption[:end].rstrip(),
                "position": start,
                "score": POSITIVE_SCORE,
                "_type": "positive",
            })
            sentence_entries.append({
                "image": rel_path,
                "sentence": caption,
                "mean": SENTENCE_MEAN_POS,
                "min": SENTENCE_MIN_POS,
            })

        else:
            # scene=0: non-scene images — use as positive sentence entries if they have good captions
            if scene == 0 and obj == 0:
                sentence_entries.append({
                    "image": rel_path,
                    "sentence": caption,
                    "mean": SENTENCE_MEAN_POS,
                    "min": SENTENCE_MIN_POS,
                })

    neg_count = sum(1 for e in pos_neg_entries if e["_type"] == "negative")
    pos_count = sum(1 for e in pos_neg_entries if e["_type"] == "positive")
    print(f"[{relation}] pos_neg: {len(pos_neg_entries)} ({pos_count} pos, {neg_count} neg)")
    print(f"[{relation}] sentences: {len(sentence_entries)}")
    print(f"[{relation}] skipped: {stats}")

    # Strip internal _type key
    final_pos_neg = [{k: v for k, v in e.items() if k != "_type"} for e in pos_neg_entries]

    def split_by_folder(entries: list[dict]) -> tuple[list[dict], list[dict]]:
        train = [e for e in entries if e["image"].startswith("train/")]
        val = [e for e in entries if e["image"].startswith("val/")]
        return train, val

    train_pn, val_pn = split_by_folder(final_pos_neg)
    train_sent, val_sent = split_by_folder(sentence_entries)

    print(f"[{relation}]   train pos_neg: {len(train_pn)}, val pos_neg: {len(val_pn)}")
    print(f"[{relation}]   train sent: {len(train_sent)}, val sent: {len(val_sent)}")

    for name, data in [
        ("pos_neg_synthetic_train", train_pn),
        ("pos_neg_synthetic_val", val_pn),
        ("sentences_synthetic_train", train_sent),
        ("sentences_synthetic_val", val_sent),
        ("pos_neg_synthetic", final_pos_neg),
        ("sentences_synthetic", sentence_entries),
    ]:
        path = os.path.join(data_dir, f"{name}.json")
        with open(path, "w") as f:
            json.dump(data, f, indent=2, ensure_ascii=False)

    # Minimal dummy VQA — GoldData loads this pre-processed format directly,
    # so image must be an absolute path (no image_dir_path prepending happens).
    first_train = next((e for e in pos_neg_entries + sentence_entries if e["image"].startswith("train/")), None)
    first_rel = first_train["image"] if first_train else (final_pos_neg + sentence_entries)[0]["image"]
    first_abs = os.path.join(img_train_dir if first_rel.startswith("train/") else img_val_dir,
                             os.path.basename(first_rel))
    dummy_vqa = [{
        "input": "Describe this image.",
        "output": all_captions[0].get("llava_caption", "An image.")[:100],
        "image": first_abs,
    }]
    vqa_path = os.path.join(data_dir, "dummy_vqa.json")
    with open(vqa_path, "w") as f:
        json.dump(dummy_vqa, f, indent=2)

    print(f"[{relation}] Saved all files to {data_dir}")


def main() -> None:
    ap = argparse.ArgumentParser(description="Format EFUF training data for any relation")
    ap.add_argument("--relation", default=None, help="Single relation key; omit for all 4")
    args = ap.parse_args()

    relations = [args.relation] if args.relation else list_relation_keys()
    for rel in relations:
        format_relation(rel)

    print("\nDone.")


if __name__ == "__main__":
    main()