File size: 4,824 Bytes
613264a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
"""Validate Core AI against full upstream BF16 text and image inference."""

import argparse
import asyncio
import json
import sys
import time
from pathlib import Path

import numpy as np
import torch

from coreai_d3 import D3CoreAI
from validation_cases import CASES


async def main():
    p = argparse.ArgumentParser()
    p.add_argument("name")
    a = p.parse_args()
    root = Path("upstream") / a.name
    bundle = Path("bundles") / (a.name + "-CoreAI")
    sys.path.insert(0, str(root.resolve()))
    from d3_runtime import D3

    torch.set_num_threads(4)
    reference = D3.from_pretrained(root, device="cpu", verify="full")
    engine = await D3CoreAI.load(bundle)
    cases = [(s, q, None, None) for s, q in CASES]
    cases.append(
        (
            "refund " * 350,
            {
                "type": "choice",
                "instructions": "What topic?",
                "criteria": {"refund": None, "shipping": None},
            },
            None,
            None,
        )
    )
    cases.append(
        (
            "A receipt.",
            {
                "type": "choice",
                "instructions": "Which item appears on the receipt?",
                "criteria": {"blender": None, "laptop": None, "bicycle": None},
            },
            [root / "assets/example-receipt.png"],
            None,
        )
    )
    cases.append(
        (
            "A receipt.",
            {"type": "noul", "instructions": "Does this receipt list a blender?"},
            [root / "assets/example-receipt.png"],
            None,
        )
    )
    cases.append(
        (
            "Two receipts.",
            {"type": "noul", "instructions": "Do the receipts list a blender?"},
            [root / "assets/example-receipt.png"] * 2,
            None,
        )
    )
    cases.append(
        (
            "A camera clip.",
            {
                "type": "choice",
                "instructions": "Which way does the square move?",
                "criteria": {
                    "right": "From left to right",
                    "left": "From right to left",
                    "still": "It does not move",
                },
            },
            None,
            [root / "assets/example-video.mp4"],
        )
    )
    records = []
    for state, q, images, videos in cases:
        start = time.monotonic()
        logits, keys, encoded = await engine.logits(state, q, images, videos)
        actual = torch.tensor(logits).softmax(-1).numpy()
        if images or videos:
            # Compare identical resized pixels, prompts, positions and the source BF16 model.
            inputs = {
                k: torch.from_numpy(np.asarray(v))
                for k, v in encoded.items()
                if k
                in [
                    "input_ids",
                    "attention_mask",
                    "pixel_values",
                    "pixel_values_videos",
                    "image_grid_thw",
                    "video_grid_thw",
                    "mm_token_type_ids",
                ]
            }
            with torch.inference_mode():
                h = reference.backbone(**inputs, use_cache=False).last_hidden_state[
                    :, -1
                ]
                scores = (h.float() @ reference.readout.T)[0, : len(keys)]
                expected = scores.softmax(-1).numpy()
        else:
            seq = encoded["input_ids"][0].tolist()
            expected = np.array(reference.probabilities([seq], [len(keys)])[0])
        r = {
            "type": q["type"],
            "modality": "video"
            if videos
            else "images"
            if images and len(images) > 1
            else "image"
            if images
            else "text",
            "tokens": len(encoded["input_ids"][0]),
            "options": len(keys),
            "coreai_choice": int(actual.argmax()),
            "upstream_choice": int(expected.argmax()),
            "max_probability_error": float(abs(expected - actual).max()),
            "total_seconds": time.monotonic() - start,
        }
        print(json.dumps(r), flush=True)
        records.append(r)
        if not np.isfinite(actual).all():
            raise AssertionError("Nonfinite Core AI probabilities")
        if r["coreai_choice"] != r["upstream_choice"]:
            raise AssertionError(r)
    (bundle / "validation.json").write_text(json.dumps(records, indent=2) + "\n")
    try:
        await engine.logits("refund " * 1000, CASES[0][1])
    except ValueError:
        pass
    else:
        raise AssertionError("Oversized request was accepted")
    print(
        "PASS: selected answers agree; probability drift is recorded, not treated as strict parity",
        flush=True,
    )


if __name__ == "__main__":
    asyncio.run(main())