path: root/worldalign/vg_prepare.py

from __future__ import annotations

import argparse
import json
from pathlib import Path
import random
from typing import Any, Iterable

from datasets import Dataset, load_dataset

from .common import write_json


DATASET = "ranjaykrishna/visual_genome"


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser()
    p.add_argument("--output-dir", default="artifacts/vg")
    p.add_argument("--cache-dir", default="/tmp/yurenh2-worldalign-vg-hf")
    p.add_argument("--nodes", type=int, default=100_000)
    p.add_argument("--views", type=int, default=32)
    p.add_argument("--min-regions", type=int, default=32)
    p.add_argument("--seed", type=int, default=20260728)
    p.add_argument(
        "--with-extra-tiers",
        action="store_true",
        help="Also load visible relation/attribute and QA annotations.",
    )
    return p.parse_args()


def load_config(name: str, cache_dir: str) -> Dataset:
    return load_dataset(
        DATASET,
        name,
        split="train",
        trust_remote_code=True,
        with_image=False,
        cache_dir=cache_dir,
    )


def by_image_id(dataset: Dataset, field: str) -> dict[int, list[dict[str, Any]]]:
    return {
        int(image_id): values
        for image_id, values in zip(dataset["image_id"], dataset[field])
    }


def first_name(obj: dict[str, Any]) -> str:
    names = obj.get("names") or []
    return str(names[0]).strip() if names else ""


def visible_relation_texts(
    image_id: int,
    relationships: dict[int, list[dict[str, Any]]],
    attributes: dict[int, list[dict[str, Any]]],
) -> list[str]:
    texts: list[str] = []
    for rel in relationships.get(image_id, []):
        subject = first_name(rel["subject"])
        predicate = str(rel.get("predicate") or "").strip()
        object_ = first_name(rel["object"])
        if subject and predicate and object_:
            texts.append(f"{subject} {predicate} {object_}.")
    for obj in attributes.get(image_id, []):
        name = first_name(obj)
        for attribute in obj.get("attributes") or []:
            attribute = str(attribute).strip()
            if name and attribute:
                texts.append(f"{name} is {attribute}.")
    return texts


def qa_texts(
    image_id: int, questions: dict[int, list[dict[str, Any]]]
) -> list[str]:
    texts: list[str] = []
    for item in questions.get(image_id, []):
        question = str(item.get("question") or "").strip()
        answer = str(item.get("answer") or "").strip()
        if question and answer:
            texts.append(f"Question: {question} Answer: {answer}")
    return texts


def fixed_mixture(
    groups: list[list[str]], budget: int, rng: random.Random
) -> list[str]:
    """Draw a fixed-size, approximately balanced mixture without count leakage."""
    groups = [list(dict.fromkeys(group)) for group in groups if group]
    for group in groups:
        rng.shuffle(group)
    selected: list[str] = []
    while len(selected) < budget and groups:
        next_groups: list[list[str]] = []
        for group in groups:
            if group and len(selected) < budget:
                selected.append(group.pop())
            if group:
                next_groups.append(group)
        groups = next_groups
    if not selected:
        raise ValueError("Cannot construct a text bundle from empty groups")
    seed_values = selected.copy()
    while len(selected) < budget:
        selected.append(rng.choice(seed_values))
    rng.shuffle(selected)
    return selected


def write_jsonl(path: Path, records: Iterable[dict[str, Any]]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", encoding="utf-8") as handle:
        for record in records:
            handle.write(json.dumps(record, ensure_ascii=False) + "\n")


def main() -> None:
    args = parse_args()
    if args.views < 1:
        raise ValueError("--views must be positive")
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    rng = random.Random(args.seed)

    region_ds = load_config("region_descriptions_v1.2.0", args.cache_dir)
    relation_lookup: dict[int, list[dict[str, Any]]] = {}
    attribute_lookup: dict[int, list[dict[str, Any]]] = {}
    question_lookup: dict[int, list[dict[str, Any]]] = {}
    if args.with_extra_tiers:
        relation_lookup = by_image_id(
            load_config("relationships_v1.2.0", args.cache_dir),
            "relationships",
        )
        attribute_lookup = by_image_id(
            load_config("attributes_v1.2.0", args.cache_dir),
            "attributes",
        )
        question_lookup = by_image_id(
            load_config("question_answers_v1.2.0", args.cache_dir),
            "qas",
        )

    eligible = [
        idx
        for idx, regions in enumerate(region_ds["regions"])
        if len(regions) >= args.min_regions
    ]
    rng.shuffle(eligible)
    selected = eligible[: min(args.nodes, len(eligible))]

    # The two orders and all opaque identifiers are independent. Source IDs
    # appear only in the vision preprocessing file and private truth file.
    vision_order = selected.copy()
    text_order = selected.copy()
    rng.shuffle(vision_order)
    rng.shuffle(text_order)
    vision_opaque = {idx: f"v{rank:07d}" for rank, idx in enumerate(vision_order)}
    text_opaque = {idx: f"t{rank:07d}" for rank, idx in enumerate(text_order)}

    vision_records: list[dict[str, Any]] = []
    text_records: list[dict[str, Any]] = []
    truth: list[dict[str, Any]] = []

    for idx in selected:
        row = region_ds[int(idx)]
        regions = list(row["regions"])
        # Fix the view count and use a deterministic per-node sample. This
        # removes region-count and ordering side channels.
        node_rng = random.Random(args.seed ^ int(row["image_id"]))
        node_rng.shuffle(regions)
        regions = regions[: args.views]

        visual_regions = [
            {
                "x": int(region["x"]),
                "y": int(region["y"]),
                "width": int(region["width"]),
                "height": int(region["height"]),
            }
            for region in regions
        ]
        phrases = [str(region["phrase"]).strip() for region in regions]
        node_rng.shuffle(phrases)
        image_id = int(row["image_id"])
        v_id = vision_opaque[idx]
        t_id = text_opaque[idx]
        vision_records.append(
            {
                "node_id": v_id,
                "source_image_id": image_id,
                "url": row["url"],
                "width": int(row["width"]),
                "height": int(row["height"]),
                "regions": visual_regions,
            }
        )
        text_record: dict[str, Any] = {
            "node_id": t_id,
            "region_closed": phrases,
        }
        if args.with_extra_tiers:
            relation_text = visible_relation_texts(
                image_id, relation_lookup, attribute_lookup
            )
            qa_text = qa_texts(image_id, question_lookup)
            text_record["visible_relations"] = fixed_mixture(
                [phrases, relation_text], args.views, node_rng
            )
            text_record["qa_expanded"] = fixed_mixture(
                [phrases, relation_text, qa_text], args.views, node_rng
            )
        text_records.append(text_record)
        truth.append(
            {
                "vision_node_id": v_id,
                "text_node_id": t_id,
                "source_image_id": image_id,
            }
        )

    rng.shuffle(vision_records)
    rng.shuffle(text_records)
    rng.shuffle(truth)
    write_jsonl(output_dir / "vision_nodes.private.jsonl", vision_records)
    write_jsonl(output_dir / "text_nodes.jsonl", text_records)
    write_jsonl(output_dir / "ground_truth.private.jsonl", truth)

    nested_sizes = [
        size for size in (5_000, 20_000, 50_000, 100_000) if size <= len(selected)
    ]
    if len(selected) not in nested_sizes:
        nested_sizes.append(len(selected))
    manifest = {
        "dataset": DATASET,
        "seed": args.seed,
        "selected_nodes": len(selected),
        "eligible_nodes": len(eligible),
        "views_per_node": args.views,
        "minimum_available_regions": args.min_regions,
        "extra_closure_tiers": args.with_extra_tiers,
        "nested_scaling_sizes": sorted(set(nested_sizes)),
        "training_files": {
            "vision": "vision_nodes.private.jsonl",
            "text": "text_nodes.jsonl",
        },
        "evaluation_only": "ground_truth.private.jsonl",
        "leakage_note": (
            "The source image ID and URL are required only by the vision "
            "feature preprocessor. Alignment training must consume extracted "
            "features with source metadata removed."
        ),
    }
    write_json(output_dir / "manifest.json", manifest)
    print(
        f"Wrote {len(selected):,} hidden-pair nodes to {output_dir}; "
        f"{len(eligible):,} scenes met the region threshold."
    )


if __name__ == "__main__":
    main()