File size: 4,044 Bytes
616b1e2 78cd6a5 616b1e2 78cd6a5 616b1e2 78cd6a5 616b1e2 | 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 | """Compare a converted RLCD package with its pinned native full-value scorer."""
from __future__ import annotations
import argparse
import json
import sys
import time
from pathlib import Path
import coremltools as ct
import numpy as np
import torch
from decision import SOURCE_REVISION, CandidateScorer, load_native, source_path
from fixtures import CONTEXT, SCHEMA
from preprocessing import Shape, batch_arrays, prepare_candidates, select_values
from release_integrity import package_file_hashes
def score_case(case_id, context, schema, shape, tokenizer, scorer, coreml):
candidates = prepare_candidates(tokenizer, context, schema, shape)
expected, actual, call_ms = [], [], []
for offset in range(0, len(candidates), shape.candidates):
group = candidates[offset : offset + shape.candidates]
arrays = batch_arrays(tokenizer, group, shape)
with torch.inference_mode():
tensors = [torch.from_numpy(array) for array in arrays.values()]
expected.extend(scorer(*tensors).float().numpy().tolist()[: len(group)])
start = time.perf_counter()
result = coreml.predict(arrays)
call_ms.append((time.perf_counter() - start) * 1000)
actual.extend(np.asarray(result["scores"]).reshape(-1).tolist()[: len(group)])
errors = [abs(left - right) for left, right in zip(expected, actual)]
return {
"case_id": case_id,
"candidates": len(candidates),
"max_log_likelihood_error": max(errors),
"same_selected_values": select_values(candidates, expected) == select_values(candidates, actual),
"model_call_ms": call_ms,
"expected": expected,
"actual": actual,
}
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("package", type=Path)
parser.add_argument("--length", type=int, default=256)
parser.add_argument("--batch", type=int, default=8)
parser.add_argument("--max-value-tokens", type=int, default=16)
parser.add_argument("--compute-units", choices=["cpu", "cpu-ane", "all"], default="cpu-ane")
parser.add_argument("--upstream-cases", type=int, default=0)
args = parser.parse_args()
torch.set_num_threads(2)
shape = Shape(args.length, args.batch, args.max_value_tokens)
tokenizer, native = load_native("cpu")
scorer = CandidateScorer(native, shape.length).eval()
units = {
"cpu": ct.ComputeUnit.CPU_ONLY,
"cpu-ane": ct.ComputeUnit.CPU_AND_NE,
"all": ct.ComputeUnit.ALL,
}[args.compute_units]
coreml = ct.models.MLModel(str(args.package), compute_units=units)
fixtures = [("coreml-fixture", SCHEMA, CONTEXT)]
if args.upstream_cases:
sys.path.insert(0, str(source_path()))
from rlcd.tasks import CASES # noqa: E402
fixtures = [(case_id, schema, context) for case_id, schema, context, _ in CASES[: args.upstream_cases]]
cases = [
score_case(case_id, context, schema, shape, tokenizer, scorer, coreml) for case_id, schema, context in fixtures
]
report = {
"package": args.package.name,
"source_revision": SOURCE_REVISION,
"package_files_sha256": package_file_hashes(args.package),
"compute_units": args.compute_units,
"fixture_set": "upstream-rlcd-tasks" if args.upstream_cases else "coreml-fixture",
"max_log_likelihood_error": max(case["max_log_likelihood_error"] for case in cases),
"same_selected_values": all(case["same_selected_values"] for case in cases),
"cases": cases,
}
output = Path(f"build/coreml-parity-{args.package.stem}-{args.compute_units}-cases{args.upstream_cases}.json")
output.parent.mkdir(parents=True, exist_ok=True)
output.write_text(json.dumps(report, indent=2) + "\n")
print(json.dumps(report, indent=2), flush=True)
if not report["same_selected_values"] or report["max_log_likelihood_error"] > 0.5:
raise SystemExit("Core ML scorer diverged from native")
if __name__ == "__main__":
main()
|