1 files changed, 150 insertions, 0 deletions
diff --git a/worldalign/synth_decode.py b/worldalign/synth_decode.py
new file mode 100644
index 0000000..183e4a9
--- /dev/null
+++ b/worldalign/synth_decode.py
@@ -0,0 +1,150 @@
+"""Close the loop: recovered pairs to a map to generated descriptions.
+
+Matching is transductive -- it aligns one fixed population. A usable
+system needs a map, so the recovered assignment is treated as pseudo-pair
+supervision for a small vision-to-text-state regressor, and the map is
+then applied to scenes that took no part in the matching. Descriptions
+are read out through the frozen text tower.
+
+Two controls decide whether the output is image-conditioned: the same
+pipeline with the recovered assignment replaced by a random one, and the
+same map applied to a shuffled image. Hidden pairs score; nothing here
+trains on them.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from .common import read_json, seed_everything, write_json
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--data-dir", default="artifacts/synth_v1")
+    parser.add_argument("--states", required=True,
+                        help=".pt with vision_states, text_states, rows for "
+                             "the matched population and the held-out split.")
+    parser.add_argument("--assignment", required=True,
+                        help=".pt with the recovered permutation and truth.")
+    parser.add_argument("--ridge", type=float, default=1.0)
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--output", default="artifacts/synth_v1/decode.json")
+    return parser.parse_args()
+
+
+def fit_map(source: np.ndarray, target: np.ndarray, ridge: float) -> np.ndarray:
+    source = np.concatenate([source, np.ones((len(source), 1))], axis=1)
+    gram = source.T @ source + ridge * np.eye(source.shape[1])
+    return np.linalg.solve(gram, source.T @ target)
+
+
+def apply_map(weights: np.ndarray, source: np.ndarray) -> np.ndarray:
+    source = np.concatenate([source, np.ones((len(source), 1))], axis=1)
+    return source @ weights
+
+
+def nearest_caption(
+    predicted: np.ndarray, bank: np.ndarray
+) -> np.ndarray:
+    predicted = predicted / np.linalg.norm(predicted, axis=1, keepdims=True).clip(1e-9)
+    bank = bank / np.linalg.norm(bank, axis=1, keepdims=True).clip(1e-9)
+    return (predicted @ bank.T).argmax(1)
+
+
+def factor_scores(
+    predicted_rows: list[int], truth_rows: list[int], scenes: list[dict]
+) -> dict:
+    """Do the retrieved descriptions state the right world facts?"""
+    colour_f1, count_ok, group_ok = [], [], []
+    for predicted, truth in zip(predicted_rows, truth_rows):
+        p, t = scenes[predicted], scenes[truth]
+        pc = {g["color"] for g in p["groups"]}
+        tc = {g["color"] for g in t["groups"]}
+        overlap = len(pc & tc)
+        precision = overlap / max(len(pc), 1)
+        recall = overlap / max(len(tc), 1)
+        colour_f1.append(
+            0.0 if precision + recall == 0 else 2 * precision * recall / (precision + recall)
+        )
+        count_ok.append(
+            sorted(g["count"] for g in p["groups"])
+            == sorted(g["count"] for g in t["groups"])
+        )
+        group_ok.append(len(p["groups"]) == len(t["groups"]))
+    return {
+        "colour_set_f1": float(np.mean(colour_f1)),
+        "count_multiset_exact": float(np.mean(count_ok)),
+        "group_count_exact": float(np.mean(group_ok)),
+    }
+
+
+def main() -> None:
+    args = parse_args()
+    seed_everything(args.seed)
+    scenes = read_json(Path(args.data_dir, "scenes.private.json"))["scenes"]
+    states = torch.load(args.states, map_location="cpu", weights_only=False)
+    recovered = torch.load(args.assignment, map_location="cpu", weights_only=False)
+
+    match_vision = states["match_vision"].double().numpy()
+    match_text = states["match_text"].double().numpy()
+    held_vision = states["held_vision"].double().numpy()
+    held_text = states["held_text"].double().numpy()
+    held_rows = states["held_rows"]
+
+    assignment = np.asarray(recovered["assignment"])
+    generator = np.random.default_rng(args.seed)
+    random_assignment = generator.permutation(len(assignment))
+
+    report = {
+        "protocol": (
+            "The recovered assignment supplies pseudo-pairs for a ridge "
+            "map from vision states to text states; the map is applied to "
+            "held-out scenes that took no part in matching. Random "
+            "assignment and shuffled-image controls bound the claim."
+        ),
+        "matched_population": len(assignment),
+        "held_out": len(held_rows),
+        "recovery_accuracy": float(
+            (assignment == np.arange(len(assignment))).mean()
+        ),
+        "conditions": {},
+    }
+
+    for label, permutation in (
+        ("recovered_pairs", assignment),
+        ("random_pairs", random_assignment),
+    ):
+        weights = fit_map(match_vision, match_text[permutation], args.ridge)
+        predicted = apply_map(weights, held_vision)
+        retrieved = nearest_caption(predicted, held_text)
+        exact = float((retrieved == np.arange(len(held_rows))).mean())
+        entry = {
+            "held_out_retrieval_exact": exact,
+            "chance": 1.0 / len(held_rows),
+            **factor_scores(
+                [held_rows[i] for i in retrieved], held_rows, scenes
+            ),
+        }
+        if label == "recovered_pairs":
+            shuffled = generator.permutation(len(held_vision))
+            predicted_shuffled = apply_map(weights, held_vision[shuffled])
+            retrieved_shuffled = nearest_caption(predicted_shuffled, held_text)
+            entry["shuffled_image_control"] = factor_scores(
+                [held_rows[i] for i in retrieved_shuffled], held_rows, scenes
+            )
+        report["conditions"][label] = entry
+        print(json.dumps({label: entry}))
+
+    write_json(args.output, report)
+    print(f"Wrote {args.output}")
+
+
+if __name__ == "__main__":
+    main()