summaryrefslogtreecommitdiff
path: root/worldalign/vg_prepare.py
diff options
context:
space:
mode:
authorYuren Hao <blackhao0426@gmail.com>2026-08-01 14:10:03 -0500
committerYuren Hao <blackhao0426@gmail.com>2026-08-01 14:10:03 -0500
commita62cf4d2a99b4a7985c61b2a7feb92a82a8218b7 (patch)
treeee2248078db7edf3812a07f195afa3d9bd6f10c6 /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.py265
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()