| """Convert the Haas-ja-v1.0 FastConformer encoder to CoreML (.mlpackage, fp16). |
| |
| Usage: python convert.py <T_mel> [out.mlpackage] |
| Input: audio_signal [1,128,T_mel] f32 (zero past length), length [1] i32 -> encoded [1,1024,T_mel/8] |
| """ |
| import os, sys |
| |
| |
| sys.path.insert(0, os.environ.get("HAAS_EXPORT_DIR", "export")) |
| import numpy as np, torch, coremltools as ct |
| from haas_onnx.load import load_model |
| import cmlpatch |
|
|
| T = int(sys.argv[1]) |
| out = sys.argv[2] if len(sys.argv) > 2 else f"encoder_T{T}.mlpackage" |
|
|
|
|
| class Enc(torch.nn.Module): |
| def __init__(self, enc): |
| super().__init__() |
| self.enc = enc |
|
|
| def forward(self, audio_signal, length=None): |
| if length is None: |
| length = torch.full((1,), audio_signal.shape[-1], dtype=torch.long) |
| y, _ = self.enc(audio_signal=audio_signal, length=length.long()) |
| return y |
|
|
|
|
| model = load_model(os.environ.get("HAAS_NEMO_DIR", "nemo_x")) |
| enc = model.encoder.eval() |
| NL = int(os.environ.get("NL", "0")) |
| if NL: |
| enc.layers = enc.layers[:NL] |
| NOLEN = os.environ.get("NOLEN") == "1" |
| enc.update_max_seq_length(T, device="cpu") if hasattr(enc, "update_max_seq_length") else None |
| w = Enc(enc).eval() |
| x = torch.randn(1, 128, T) |
| L = torch.tensor([T], dtype=torch.int32) |
| args = (x,) if NOLEN else (x, L) |
| with torch.no_grad(): |
| ref = w(*args).numpy() |
| tr = torch.jit.trace(w, args, check_trace=False) |
| ins = [ct.TensorType("audio_signal", shape=(1, 128, T), dtype=np.float32)] |
| if not NOLEN: |
| ins.append(ct.TensorType("length", shape=(1,), dtype=np.int32)) |
| ml = ct.convert( |
| tr, |
| inputs=ins, |
| outputs=[ct.TensorType("encoded", dtype=np.float32)], |
| convert_to="mlprogram", compute_precision=ct.precision.FLOAT16, |
| minimum_deployment_target=ct.target.iOS17, |
| ) |
| ml.save(out) |
| feed = {"audio_signal": x.numpy()} |
| if not NOLEN: |
| feed["length"] = L.numpy() |
| y = ml.predict(feed)["encoded"] |
| print("shape", y.shape, "max abs diff", float(np.abs(y - ref).max()), "ref max", float(np.abs(ref).max()), |
| "cos", float((y * ref).sum() / np.linalg.norm(y) / np.linalg.norm(ref))) |
|
|
