diff options
| author | Yuren Hao <blackhao0426@gmail.com> | 2026-08-01 14:10:03 -0500 |
|---|---|---|
| committer | Yuren Hao <blackhao0426@gmail.com> | 2026-08-01 14:10:03 -0500 |
| commit | a62cf4d2a99b4a7985c61b2a7feb92a82a8218b7 (patch) | |
| tree | ee2248078db7edf3812a07f195afa3d9bd6f10c6 /worldalign/vg_prepare.py | |
World Alignment: unpaired cross-modal correspondence by relational identifiability
Method: scene states are sets of part states; relation fields are built
within each modality and are invariant to how each side labels its own
features; the cross-modal bridge is a coupling searched under an energy
that is a closed-form functional of one matrix; solving is spectral
initialisation followed by exact local refinement.
Evidence: in a procedurally generated closed world, blind recovery of a
hidden image-caption correspondence reaches 95.3% at 256 scenes against
0.39% chance, and the recovered pairs transfer to 200 held-out scenes at
93.0% exact retrieval with random-pair and shuffled-image controls at or
near chance. Cross-modal value correspondence is derived from disjoint
corpora rather than declared. On Visual Genome the field correlation
reaches 0.656 against the 0.9 that polynomial recovery needs, with the
deficit attributed away from segmentation and discretisation.
Protocol: no image-text pair enters any objective, optimiser,
initialisation, or model selection; hidden pairs score orderings only.
Co-Authored-By: Claude <noreply@anthropic.com>
Diffstat (limited to 'worldalign/vg_prepare.py')
| -rw-r--r-- | worldalign/vg_prepare.py | 265 |
1 files changed, 265 insertions, 0 deletions
diff --git a/worldalign/vg_prepare.py b/worldalign/vg_prepare.py new file mode 100644 index 0000000..7d553a8 --- /dev/null +++ b/worldalign/vg_prepare.py @@ -0,0 +1,265 @@ +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() |
