| """Manual reduced-output ANE experiment; not the deployment interface. |
| |
| Usage: python scripts/deployment/ane_output_probe.py NEW_OUTPUT.mlpackage |
| Conversion, timing and accuracy checks start only when this command is run. |
| """ |
| import argparse |
| import sys |
| import time |
| from pathlib import Path |
|
|
| ROOT = Path(__file__).resolve().parents[2] |
| sys.path.insert(0, str(ROOT / "src")) |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("output", type=Path, help="New scratch .mlpackage path") |
| args = parser.parse_args() |
| out = args.output |
| if out.exists(): |
| parser.error("Output exists; choose a new experiment path.") |
| if out.suffix != ".mlpackage": |
| parser.error("Output must end with .mlpackage.") |
| if sys.platform != "darwin": |
| parser.error("Run this experiment manually on Mac.") |
| import numpy as np, torch, coremltools as ct |
| from torch import nn |
| from vimeml.deployment.bundle import BundleLM |
| from vimeml.deployment.graph import ConversionGraph |
| lm = BundleLM(ROOT/"artifacts/deployment/tiny-ja-v1-inference-v1") |
| K = 16 |
| class Head(nn.Module): |
| def __init__(s, model): super().__init__(); s.g = ConversionGraph(model) |
| def forward(s, input_ids, target_ids): |
| logp = torch.log_softmax(s.g(input_ids), dim=-1) |
| target = torch.gather(logp, 2, target_ids.long().unsqueeze(-1)).squeeze(-1) |
| top, idx = torch.topk(logp, K, dim=-1) |
| return target, top, idx.int() |
| h = Head(lm.model).eval() |
| ex = (torch.zeros(1,16,dtype=torch.int32), torch.zeros(1,16,dtype=torch.int32)) |
| tr = torch.jit.trace(h, ex) |
| shapes = ct.EnumeratedShapes(shapes=[(1,L) for L in (16,32,64,128)], default=(1,16)) |
| m = ct.convert(tr, convert_to="mlprogram", minimum_deployment_target=ct.target.iOS18, compute_precision=ct.precision.FLOAT16, |
| inputs=[ct.TensorType(name="input_ids", shape=shapes, dtype=np.int32), ct.TensorType(name="target_ids", shape=shapes, dtype=np.int32)], |
| outputs=[ct.TensorType(name="target_logprobs", dtype=np.float32), ct.TensorType(name="top_logprobs", dtype=np.float32), ct.TensorType(name="top_ids", dtype=np.int32)]) |
| m.save(str(out)) |
| for units in ("CPU_ONLY","CPU_AND_NE"): |
| mm = ct.models.MLModel(str(out), compute_units=getattr(ct.ComputeUnit, units)); row=[units] |
| for L in (16,64,128): |
| x={"input_ids":np.full((1,L),5,np.int32),"target_ids":np.full((1,L),6,np.int32)}; mm.predict(x); ts=[] |
| for _ in range(20): t=time.perf_counter(); mm.predict(x); ts.append((time.perf_counter()-t)*1000) |
| row.append(f"T{L}={np.median(ts):.2f}ms") |
| print(" ".join(row)) |
| |
| import json |
| items=json.load(open(ROOT/"artifacts/benchmarks/ajimee-jwtd-v2-v1/evaluation_items.json")) |
| mm = ct.models.MLModel(str(out), compute_units=ct.ComputeUnit.CPU_AND_NE) |
| diffs=[]; topagree=[] |
| for it in items[:100]: |
| ids=[2]+lm.processor.encode((it.get("context_text") or "")+it["expected_output"][0])[:127] |
| n=len(ids)-1; L=next(b for b in (16,32,64,128) if b>=n) |
| inp=np.zeros((1,L),np.int32); inp[0,:n]=ids[:-1]; tgt=np.zeros((1,L),np.int32); tgt[0,:n]=ids[1:] |
| o=mm.predict({"input_ids":inp,"target_ids":tgt}) |
| with torch.no_grad(): ref=lm.model(torch.tensor([ids[:-1]])).float().log_softmax(-1)[0] |
| r=ref.gather(1,torch.tensor(ids[1:])[:,None]).squeeze(1).numpy() |
| diffs.append(abs(o["target_logprobs"][0,:n]-r).max()) |
| topagree.append(np.mean(o["top_ids"][0,:n,0]==ref.argmax(-1).numpy())) |
| print(f"ANE target logprob max|diff| mean={np.mean(diffs):.4f} max={np.max(diffs):.4f}; top1 agree={np.mean(topagree):.4f}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
