diff options
Diffstat (limited to 'worldalign/natural_pipeline.py')
| -rw-r--r-- | worldalign/natural_pipeline.py | 204 |
1 files changed, 204 insertions, 0 deletions
diff --git a/worldalign/natural_pipeline.py b/worldalign/natural_pipeline.py new file mode 100644 index 0000000..c30af4d --- /dev/null +++ b/worldalign/natural_pipeline.py @@ -0,0 +1,204 @@ +"""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() |