File size: 3,259 Bytes
26afca4 | 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 | #!/usr/bin/env python3
"""Run the native Swift host and compare its EOS-bound output with a BF16 receipt."""
from __future__ import annotations
import argparse
import json
import subprocess
import time
from pathlib import Path
EOS_TOKEN_ID = 2
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--executable", type=Path, required=True)
parser.add_argument("--model-dir", type=Path, required=True)
parser.add_argument("--pixel-values", type=Path, required=True)
parser.add_argument("--expected-receipt", type=Path, required=True)
parser.add_argument("--expected-row", type=int, default=0)
parser.add_argument("--max-new-tokens", type=int, default=64)
parser.add_argument("--output", type=Path, required=True)
args = parser.parse_args()
expected_receipt = json.loads(args.expected_receipt.read_text())
if "rows" in expected_receipt:
row = expected_receipt["rows"][args.expected_row]
expected_tokens = row["reference_tokens"]
expected_text = row["reference_text"]
source = row.get("image")
else:
expected_tokens = expected_receipt.get("native_tokens", expected_receipt["coreml_tokens"])
expected_tokens = expected_tokens[: expected_tokens.index(EOS_TOKEN_ID) + 1]
expected_text = None
source = "synthetic-canary"
command = [
str(args.executable.expanduser().resolve()),
"--model-dir",
str(args.model_dir.expanduser().resolve()),
"--pixel-values-fp32",
str(args.pixel_values.expanduser().resolve()),
"--max-tokens",
str(args.max_new_tokens),
]
started = time.perf_counter()
completed = subprocess.run(command, check=False, capture_output=True, text=True)
seconds = time.perf_counter() - started
stdout_lines = completed.stdout.splitlines()
if not stdout_lines:
raise RuntimeError(f"Swift runtime produced no stdout: {completed.stderr}")
tokens = [int(value) for value in stdout_lines[0].split()]
text = "\n".join(stdout_lines[1:])
receipt = {
"source": source,
"command": [
args.executable.name,
"--model-dir",
".",
"--pixel-values-fp32",
args.pixel_values.name,
"--max-tokens",
str(args.max_new_tokens),
],
"exit_code": completed.returncode,
"seconds": seconds,
"tokens": tokens,
"expected_tokens": expected_tokens,
"token_exact": tokens == expected_tokens,
"stop_reason": "eos" if tokens and tokens[-1] == EOS_TOKEN_ID else "max_new_tokens",
"text": text,
"expected_text": expected_text,
"text_exact": expected_text is None or text == expected_text,
"stderr": completed.stderr,
}
args.output.parent.mkdir(parents=True, exist_ok=True)
args.output.write_text(json.dumps(receipt, ensure_ascii=False, indent=2) + "\n")
print(json.dumps(receipt, ensure_ascii=False, indent=2))
if (
completed.returncode != 0
or not receipt["token_exact"]
or not receipt["text_exact"]
or receipt["stop_reason"] != "eos"
):
raise SystemExit(1)
if __name__ == "__main__":
main()
|