File size: 2,865 Bytes
0deb31c | 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 | """Compare the exported Jeff FP16 Core ML package with its trained native model."""
from __future__ import annotations
import argparse
import json
import time
from pathlib import Path
import coremltools as ct
import numpy as np
import torch
from huggingface_hub import snapshot_download
from jeff.backends.torch_backend import TorchBackend
from export import FIXTURES, REVISION, SOURCE, make_batch, model_inputs
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--precision", choices=("fp16", "fp32"), default="fp16")
args = parser.parse_args()
torch.set_num_threads(2)
checkpoint = snapshot_download(SOURCE, revision=REVISION, local_files_only=True)
backend = TorchBackend(checkpoint, device="cpu", dtype="float32", attn_kernel="eager", batch_size=1)
package = Path(f"build/JeffDecision-L128-{args.precision.upper()}.mlpackage")
model = ct.models.MLModel(str(package), compute_units=ct.ComputeUnit.CPU_ONLY)
output_name = model.get_spec().description.output[0].name
results = []
for name, text, group in FIXTURES:
batch = make_batch(backend, text, group)
native = backend.model.model(**batch, include_media=False).cat_logits.detach().float().numpy()[0]
tensor_inputs = model_inputs(batch, backend.model.config)
inputs = {
key: tensor.numpy()
for key, tensor in zip(
("input_ids", "attention_mask", "parent_position", "category_positions"),
tensor_inputs,
)
}
start = time.perf_counter()
converted = model.predict(inputs)[output_name][0, :len(group.labels)].astype(np.float32)
elapsed_ms = (time.perf_counter() - start) * 1000
error = float(np.max(np.abs(native - converted))) if np.isfinite(converted).all() else float("inf")
agreement = int(np.argmax(native) == np.argmax(converted))
results.append({
"name": name,
"labels": list(group.labels),
"native_logits": native.tolist(),
"coreml_logits": converted.tolist(),
"max_logit_error": error,
"top_label_agreement": bool(agreement),
"coreml_wall_ms": elapsed_ms,
})
print(
f"{name}: logits={converted.tolist()}, error={error:.6f}, "
f"top_label={bool(agreement)}, wall_ms={elapsed_ms:.1f}",
flush=True,
)
Path(f"build/coreml-parity-{args.precision}.json").write_text(json.dumps(results, indent=2) + "\n")
if not all(row["top_label_agreement"] for row in results):
raise AssertionError("Core ML changed the chosen label on a parity fixture")
if max(row["max_logit_error"] for row in results) > 0.25:
raise AssertionError("Core ML logit error exceeds the 0.25 tolerance")
if __name__ == "__main__":
main()
|