"""Natural-data field pipeline: continuous states and content projection. Produces the field correlation reported for Visual Genome. Two choices carry most of it, both measured. States stay continuous -- quantising segments and phrases into class codes costs more than half the signal, because relation fields need scene similarity and nothing else. And each side is projected onto the directions that separate scenes rather than parts, fitted per modality on scenes outside the evaluated set. Vision states come from `natural_objects`; language states are positive-PMI context vectors of the corpus, averaged per phrase. Hidden pairs are read only to report the correlation. """ from __future__ import annotations import argparse import json import re from collections import Counter from pathlib import Path import numpy as np import torch import torch.nn.functional as F from scipy.linalg import eigh from sklearn.decomposition import TruncatedSVD from .common import read_json, seed_everything, write_json from .natural_families import load_phrases, statistics from .synth_cc_battery import moment_field def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--vg-dir", default="artifacts/vg_5k") parser.add_argument("--objects", default="artifacts/vg_5k/natural_objects.pt") parser.add_argument("--samples", type=int, default=256) parser.add_argument("--fit-scenes", type=int, default=1200) parser.add_argument("--word-vectors", type=int, default=48) parser.add_argument("--min-count", type=int, default=60) parser.add_argument("--max-vocabulary", type=int, default=600) parser.add_argument("--context-window", type=int, default=2) parser.add_argument("--keep", type=int, default=8) parser.add_argument("--shrinkage", type=float, default=0.05) parser.add_argument("--views", type=int, default=1) parser.add_argument("--seed", type=int, default=0) parser.add_argument("--output", default="artifacts/vg_5k/natural_pipeline.pt") return parser.parse_args() def word_vectors( phrases: list[list[str]], args: argparse.Namespace ) -> dict[str, np.ndarray]: """Positive-PMI context vectors, reduced by truncated SVD.""" counts = Counter(token for tokens in phrases for token in set(tokens)) vocabulary = [ word for word, count in counts.most_common(args.max_vocabulary) if count >= args.min_count ] _, _, context = statistics(phrases, set(vocabulary), args.context_window) features = sorted({key for word in vocabulary for key in context[word]}) index = {key: position for position, key in enumerate(features)} matrix = np.zeros((len(vocabulary), len(features))) for row, word in enumerate(vocabulary): for key, value in context[word].items(): matrix[row, index[key]] = value total = matrix.sum() expected = matrix.sum(1, keepdims=True) * matrix.sum(0, keepdims=True) / total pmi = np.log(np.maximum(matrix, 1e-9) / np.maximum(expected, 1e-9)) pmi[matrix == 0] = 0.0 embedding = TruncatedSVD(args.word_vectors, random_state=args.seed).fit_transform( np.maximum(pmi, 0.0) ) embedding /= np.linalg.norm(embedding, axis=1, keepdims=True).clip(1e-9) return {word: embedding[row] for row, word in enumerate(vocabulary)} def content_directions( sets: list[np.ndarray], shrinkage: float ) -> tuple[np.ndarray, np.ndarray]: """Directions separating scenes more than they separate parts.""" means = np.stack([item.mean(0) for item in sets]) centre = means.mean(0) within = np.concatenate([item - item.mean(0, keepdims=True) for item in sets]) scatter_within = within.T @ within / max(len(within) - 1, 1) centred = means - centre scatter_between = centred.T @ centred / max(len(centred) - 1, 1) trace = np.trace(scatter_within) / len(scatter_within) values, vectors = eigh( scatter_between, scatter_within + shrinkage * trace * np.eye(len(scatter_within)) ) return centre, vectors[:, np.argsort(values)[::-1]] def main() -> None: args = parse_args() seed_everything(args.seed) vg_dir = Path(args.vg_dir) truth = [ json.loads(line) for line in (vg_dir / "ground_truth.private.jsonl") .read_text(encoding="utf-8") .splitlines() if line.strip() ] records = { json.loads(line)["node_id"]: json.loads(line) for line in (vg_dir / "text_nodes.jsonl").read_text(encoding="utf-8").splitlines() if line.strip() } state = torch.load(args.objects, map_location="cpu", weights_only=False) per_view = state.get("view_segments") or [state["segments"]] views = min(args.views, len(per_view)) index = {node: position for position, node in enumerate(state["node_ids"])} vectors = word_vectors(load_phrases(vg_dir / "text_nodes.jsonl", "region_closed"), args) def language_state(node: str) -> np.ndarray | None: rows = [] for phrase in records[node]["region_closed"]: hits = [vectors[token] for token in re.findall(r"[a-z]+", phrase.lower()) if token in vectors] if hits: rows.append(np.mean(hits, axis=0)) return np.stack(rows) if rows else None def vision_state(node: str, view: int) -> np.ndarray | None: segments = per_view[view][index[node]] return ( np.stack([item["feature"] for item in segments]).astype(np.float64) if segments else None ) pairs = [ pair for pair in truth if pair["vision_node_id"] in index and pair["text_node_id"] in records ] fit = pairs[args.samples : args.samples + args.fit_scenes] language_fit = [language_state(p["text_node_id"]) for p in fit] language_fit = [item for item in language_fit if item is not None and len(item) > 1] vision_fit = [vision_state(p["vision_node_id"], 0) for p in fit] vision_fit = [item for item in vision_fit if item is not None and len(item) > 1] language_centre, language_basis = content_directions(language_fit, args.shrinkage) vision_centre, vision_basis = content_directions(vision_fit, args.shrinkage) evaluated = [ pair for pair in pairs[: args.samples] if language_state(pair["text_node_id"]) is not None ] keep = args.keep def project(raw: np.ndarray, centre: np.ndarray, basis: np.ndarray) -> torch.Tensor: width = min(keep, basis.shape[1]) return F.normalize( torch.tensor((raw - centre) @ basis[:, :width], dtype=torch.float32), dim=-1 ) text_sets = [ project(language_state(p["text_node_id"]), language_centre, language_basis) for p in evaluated ] view_fields = [] for view in range(views): sets = [ project(vision_state(p["vision_node_id"], view), vision_centre, vision_basis) for p in evaluated ] view_fields.append(moment_field(sets)) visual_field = torch.stack(view_fields).mean(0) text_field = moment_field(text_sets) mask = ~np.eye(len(evaluated), dtype=bool) correlation = float( np.corrcoef( visual_field.double().numpy()[mask], text_field.double().numpy()[mask] )[0, 1] ) torch.save( {"visual_field": visual_field, "text_field": text_field, "nodes": [p["vision_node_id"] for p in evaluated]}, args.output, ) summary = { "protocol": ( "Continuous states, content projection fitted per modality on " "scenes outside the evaluated set. Hidden pairs are read only " "for the correlation." ), "samples": len(evaluated), "views": views, "kept_directions": keep, "field_correlation_at_truth": correlation, "recovery_threshold": 0.9, } print(json.dumps(summary)) write_json(str(args.output).replace(".pt", ".json"), summary) if __name__ == "__main__": main()