"""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  # noqa: F401
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")