import sys, time
import numpy as np, coremltools as ct
for path in sys.argv[1:]:
    m = ct.models.MLModel(path, compute_units=ct.ComputeUnit.CPU_AND_GPU)
    spec = {i.name: tuple(i.type.multiArrayType.shape) for i in m.get_spec().description.input}
    x = {}
    for k, shp in spec.items():
        dt = np.int32 if k in ("input_ids", "cand_idx", "query_idx") else np.float16
        x[k] = np.zeros(shp, dtype=dt)
    x["valid"][:] = 1; x["segment"][:] = np.eye(spec["segment"][0])
    if "seg_chunks" in x: x["seg_chunks"][:] = np.eye(64)
    m.predict(x); ts = []
    for _ in range(10):
        t = time.time(); m.predict(x); ts.append((time.time() - t) * 1000)
    print(f"{path.split('/')[-1]:40s} p50 {np.median(ts):6.1f} ms")