File size: 2,903 Bytes
3bfcc7d 970988e 3bfcc7d 970988e 3bfcc7d 970988e 3bfcc7d 970988e 3bfcc7d | 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 | """Selected fixed-bucket end-to-end Core ML extraction latency on this Mac."""
import argparse
import json
import platform
import statistics
import time
from pathlib import Path
import coremltools as ct
import psutil
from gliner2 import Schema
from extraction_runtime import CoreMLBoundaryExtractor
UNITS = {
"cpu_only": ct.ComputeUnit.CPU_ONLY,
"cpu_and_gpu": ct.ComputeUnit.CPU_AND_GPU,
"cpu_and_neural_engine": ct.ComputeUnit.CPU_AND_NE,
"all": ct.ComputeUnit.ALL,
}
def percentile(values, fraction):
ordered = sorted(values)
return ordered[min(round(fraction * (len(ordered) - 1)), len(ordered) - 1)]
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--model-dir", required=True)
parser.add_argument("--precision", choices=["fp16", "fp32"], default="fp32")
parser.add_argument("--feature-package")
parser.add_argument("--units", choices=list(UNITS), default="all")
parser.add_argument("--warmup", type=int, default=20)
parser.add_argument("--iterations", type=int, default=200)
args = parser.parse_args()
text = "Alice founded Acme in Toronto in 2020."
schema = Schema().entities(["person", "organization", "location"])
started = time.perf_counter()
runtime = CoreMLBoundaryExtractor(
args.model_dir,
precision=args.precision,
compute_units=UNITS[args.units],
feature_package=args.feature_package,
)
load_ms = (time.perf_counter() - started) * 1000
for _ in range(args.warmup):
runtime.extract(text, schema)
process = psutil.Process()
latencies = []
peak_rss = process.memory_info().rss
for _ in range(args.iterations):
start = time.perf_counter()
runtime.extract(text, schema)
latencies.append((time.perf_counter() - start) * 1000)
peak_rss = max(peak_rss, process.memory_info().rss)
report = {
"purpose": "selected end-to-end entity extraction latency, no benchmark scoring",
"fixture": text,
"shape": "L128/W64/Q8/C192",
"precision": args.precision,
"compute_units": args.units,
"feature_package": args.feature_package,
"warmup": args.warmup,
"iterations": args.iterations,
"load_ms": load_ms,
"p50_ms": statistics.median(latencies),
"p95_ms": percentile(latencies, 0.95),
"mean_ms": statistics.mean(latencies),
"peak_process_rss_bytes": peak_rss,
"macos": platform.mac_ver()[0],
"machine": platform.machine(),
"coremltools": ct.__version__,
}
folder = Path(args.model_dir)
variant = "w8-embedding" if args.feature_package else "baseline"
path = folder / f"benchmark-{args.precision}-{args.units}-{variant}.json"
path.write_text(json.dumps(report, indent=2) + "\n")
print(json.dumps(report, indent=2))
if __name__ == "__main__":
main()
|