1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
|
"""Ground-truth factor probes for synthetic-world representations.
Linear probes from a representation to the discrete scene factors measure
which world variables survive the encoder and its pooling. Probes are
fitted per modality on that modality's own training split and evaluated
on held-out scenes; scene truth is generator metadata, so this is a
within-modality diagnostic, not a cross-modal signal.
"""
from __future__ import annotations
import argparse
import json
from pathlib import Path
import numpy as np
import torch
from .common import read_json, write_json
from .synth_world import COLORS, SHAPES
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument("--data-dir", default="artifacts/synth_v0")
parser.add_argument("--features", required=True)
parser.add_argument(
"--side", choices=["vision", "text"], required=True,
help="Selects the training split whose scenes fit the probes.",
)
parser.add_argument("--ridge", type=float, default=1.0)
parser.add_argument("--output", required=True)
return parser.parse_args()
def scene_targets(scene: dict) -> dict[str, np.ndarray | float]:
color_presence = np.zeros(len(COLORS))
shape_presence = np.zeros(len(SHAPES))
for group in scene["groups"]:
color_presence[list(COLORS).index(group["color"])] = 1.0
shape_presence[SHAPES.index(group["shape"])] = 1.0
return {
"color_presence": color_presence,
"shape_presence": shape_presence,
"group_count": float(len(scene["groups"])),
"object_total": float(sum(g["count"] for g in scene["groups"])),
"relation_count": float(len(scene["relations"])),
}
def ridge_fit(
x: np.ndarray, y: np.ndarray, ridge: float
) -> tuple[np.ndarray, np.ndarray]:
x = np.concatenate([x, np.ones((len(x), 1))], axis=1)
gram = x.T @ x + ridge * np.eye(x.shape[1])
weights = np.linalg.solve(gram, x.T @ y)
return weights, x @ weights
def evaluate(
features_train: np.ndarray,
features_test: np.ndarray,
train_targets: np.ndarray,
test_targets: np.ndarray,
ridge: float,
binary: bool,
) -> dict:
weights, _ = ridge_fit(features_train, train_targets, ridge)
prediction = (
np.concatenate([features_test, np.ones((len(features_test), 1))], axis=1)
@ weights
)
if binary:
accuracy = float(((prediction > 0.5) == (test_targets > 0.5)).mean())
balanced = []
for column in range(test_targets.shape[1]):
truth = test_targets[:, column] > 0.5
if truth.any() and (~truth).any():
hit = (prediction[:, column] > 0.5) == truth
balanced.append(
(hit[truth].mean() + hit[~truth].mean()) / 2.0
)
return {
"accuracy": accuracy,
"balanced_accuracy": float(np.mean(balanced)),
}
residual = prediction[:, 0] - test_targets[:, 0]
variance = test_targets[:, 0].var()
return {
"r2": float(1.0 - residual.var() / max(variance, 1e-9)),
"mae": float(np.abs(residual).mean()),
}
def main() -> None:
args = parse_args()
manifest = read_json(Path(args.data_dir, "manifest.json"))
scenes = read_json(Path(args.data_dir, "scenes.private.json"))["scenes"]
state = torch.load(args.features, map_location="cpu", weights_only=False)
lookup = {int(row): i for i, row in enumerate(state["rows"])}
features = state["features"].float().numpy()
split_key = "vision_only_train" if args.side == "vision" else "text_only_train"
train_rows = [row for row in manifest[split_key] if row in lookup][:8000]
test_rows = [row for row in manifest["test"] if row in lookup]
x_train = features[[lookup[row] for row in train_rows]]
x_test = features[[lookup[row] for row in test_rows]]
report = {"features": args.features, "side": args.side, "targets": {}}
for name in ("color_presence", "shape_presence"):
y_train = np.stack([scene_targets(scenes[row])[name] for row in train_rows])
y_test = np.stack([scene_targets(scenes[row])[name] for row in test_rows])
report["targets"][name] = evaluate(
x_train, x_test, y_train, y_test, args.ridge, binary=True
)
for name in ("group_count", "object_total", "relation_count"):
y_train = np.array(
[[scene_targets(scenes[row])[name]] for row in train_rows]
)
y_test = np.array([[scene_targets(scenes[row])[name]] for row in test_rows])
report["targets"][name] = evaluate(
x_train, x_test, y_train, y_test, args.ridge, binary=False
)
write_json(args.output, report)
print(json.dumps(report, indent=2))
if __name__ == "__main__":
main()
|
