File size: 3,285 Bytes
bea429c | 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 | """Check native scorer versus fixed K16 Core ML over real upstream demo requests."""
from __future__ import annotations
import argparse
import json
import statistics
import time
from pathlib import Path
import numpy as np
from assets import ROOT, check_base_access
from native_reference import choose, encode_options, pad_batch
from scorer import score_coreml
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--package", type=Path, default=ROOT / "build" / "system_one_gemma_fp16_L256_K16.mlpackage")
parser.add_argument("--repeats", type=int, default=10)
args = parser.parse_args()
check_base_access()
import coremltools as ct
import torch
from export_model import load_trained_scorer
tokenizer, native = load_trained_scorer()
coreml = ct.models.MLModel(str(args.package), compute_units=ct.ComputeUnit.ALL)
demos = json.loads((ROOT / "build" / "source" / "demos.json").read_text())
rows = []
for demo in demos:
sequences = encode_options(tokenizer, demo["state"], demo["question"], demo["options"])
ids, masks = pad_batch(sequences, tokenizer.pad_token_id)
with torch.inference_mode():
native_logits = (
native(
input_ids=torch.tensor(ids, dtype=torch.long),
attention_mask=torch.tensor(masks, dtype=torch.long),
)
.logits.reshape(-1)
.float()
.tolist()
)
native_choice = choose(demo["options"], native_logits)
times = []
converted = None
for _ in range(args.repeats):
start = time.perf_counter()
converted = score_coreml(coreml, tokenizer, demo["state"], demo["question"], demo["options"])
times.append((time.perf_counter() - start) * 1000)
assert converted is not None
error = max(abs(a - b) for a, b in zip(native_choice["probabilities"], converted["probabilities"]))
rows.append(
{
"demo": demo["name"],
"candidates": len(demo["options"]),
"native_selected_index": native_choice["selected_index"],
"coreml_selected_index": converted["selected_index"],
"max_calibrated_probability_error": error,
"median_end_to_end_ms": statistics.median(times),
}
)
report = {
"package": args.package.name,
"source": "pinned upstream demos.json",
"questions": len(rows),
"selected_id_agreements": sum(r["native_selected_index"] == r["coreml_selected_index"] for r in rows),
"max_calibrated_probability_error": float(np.max([r["max_calibrated_probability_error"] for r in rows])),
"rows": rows,
}
report["passed"] = (
report["selected_id_agreements"] == len(rows) and report["max_calibrated_probability_error"] <= 0.02
)
target = ROOT / "reports" / "verification-L256-K16.json"
target.parent.mkdir(parents=True, exist_ok=True)
target.write_text(json.dumps(report, indent=2) + "\n")
print(json.dumps(report, indent=2))
if not report["passed"]:
raise SystemExit(1)
if __name__ == "__main__":
main()
|