| """Convert the vela3-fast decision backbone encoder to a fixed-length fp16 Core ML package: python convert.py L [name] |
| -> coreml/<name>_L<L>.mlpackage (inputs input_ids / attention_mask int32 [1, L], output hidden fp16 [1, L, 1024]).""" |
| import sys |
| import numpy as np |
| import torch |
| import coremltools as ct |
| import ct_patch |
| from vela3_export import Vela3EncoderExport, load_decision |
|
|
| L = int(sys.argv[1]); name = sys.argv[2] if len(sys.argv) > 2 else "vela3_decision" |
| m = load_decision() |
| g = Vela3EncoderExport(m.encoder, L).eval() |
| x = torch.zeros(1, L, dtype=torch.int32); mk = torch.ones(1, L, dtype=torch.int32) |
| with torch.no_grad(): |
| tr = torch.jit.trace(g, (x, mk)) |
| ml = ct.convert(tr, inputs=[ct.TensorType("input_ids", shape=(1, L), dtype=np.int32), ct.TensorType("attention_mask", shape=(1, L), dtype=np.int32)], |
| outputs=[ct.TensorType("hidden", dtype=np.float16)], minimum_deployment_target=ct.target.macOS15, |
| compute_precision=ct.precision.FLOAT16, convert_to="mlprogram") |
| ml.short_description = f"vela3-fast UltraBERT-380M encoder ({name}), fixed length {L}" |
| ml.save(f"coreml/{name}_L{L}.mlpackage") |
| print("saved", f"coreml/{name}_L{L}.mlpackage") |
|
|
