| """Convert the TDT prediction net (KDA) and joint to CoreML for on-device decoding (fp32, meant for the CPU). |
| |
| decoder.mlpackage token [1] i32, rec [2,1,20480], conv [2,1,20480] -> g [1,640] (joint.pred applied), rec_out, conv_out |
| joint_enc.mlpackage encoded [1,1024,188] -> f [1,188,640] (joint.enc), once per 15 s window |
| joint.mlpackage f [1,640], g [1,640] -> token_id [1] i32 (argmax of 3073), duration_id [1] i32 (argmax of 5) |
| The decoder only runs after a non-blank token; the joint runs every step. |
| """ |
| import os, sys |
| |
| |
| sys.path.insert(0, os.environ.get("HAAS_EXPORT_DIR", "export")) |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
| import numpy as np, torch, coremltools as ct |
| from haas_onnx.load import load_model |
| import cmlpatch |
|
|
| OUT = os.path.dirname(os.path.abspath(__file__)) |
| N_DUR = 5 |
|
|
|
|
| class Dec(torch.nn.Module): |
| def __init__(self, model): |
| super().__init__() |
| self.embed = model.decoder.prediction["embed"] |
| self.rnn = model.decoder.prediction["dec_rnn"] |
| self.pred = model.joint.pred |
|
|
| def forward(self, token, rec, conv): |
| y = self.embed(token.long()).unsqueeze(0) |
| g, (rec, conv) = self.rnn(y, (rec, conv)) |
| return self.pred(g[0]), rec, conv |
|
|
|
|
| class JointEnc(torch.nn.Module): |
| def __init__(self, model): |
| super().__init__() |
| self.enc = model.joint.enc |
|
|
| def forward(self, encoded): |
| return self.enc(encoded.transpose(1, 2)) |
|
|
|
|
| class Joint(torch.nn.Module): |
| def __init__(self, model): |
| super().__init__() |
| self.net = model.joint.joint_net |
|
|
| def forward(self, f, g): |
| logits = self.net(f + g) |
| tok = logits[:, :-N_DUR].argmax(-1).to(torch.int32) |
| dur = logits[:, -N_DUR:].argmax(-1).to(torch.int32) |
| return tok, dur |
|
|
|
|
| def conv(module, args, ins, outs, name): |
| with torch.no_grad(): |
| tr = torch.jit.trace(module.eval(), args, check_trace=False) |
| ml = ct.convert(tr, inputs=ins, outputs=[ct.TensorType(o) for o in outs], convert_to="mlprogram", |
| compute_precision=ct.precision.FLOAT32, minimum_deployment_target=ct.target.iOS17) |
| ml.save(f"{OUT}/{name}.mlpackage") |
| return ml |
|
|
|
|
| if __name__ == "__main__": |
| model = load_model(os.environ.get("HAAS_NEMO_DIR", "nemo_x")) |
| print(model.joint.joint_net) |
| W = model.decoder.prediction["dec_rnn"].width |
| tok = torch.tensor([42], dtype=torch.int32) |
| rec, cv = torch.randn(2, 1, W) * 0.1, torch.randn(2, 1, W) * 0.1 |
| d = Dec(model) |
| m = conv(d, (tok, rec, cv), [ct.TensorType("token", shape=(1,), dtype=np.int32), |
| ct.TensorType("rec", shape=(2, 1, W), dtype=np.float32), |
| ct.TensorType("conv", shape=(2, 1, W), dtype=np.float32)], |
| ["g", "rec_out", "conv_out"], "decoder") |
| with torch.no_grad(): |
| r = d(tok, rec, cv) |
| y = m.predict({"token": tok.numpy(), "rec": rec.numpy(), "conv": cv.numpy()}) |
| print("decoder maxdiff", [float(np.abs(y[k] - v.numpy()).max()) for k, v in zip(["g", "rec_out", "conv_out"], r)]) |
| e = torch.randn(1, 1024, 188) |
| m = conv(JointEnc(model), (e,), [ct.TensorType("encoded", shape=(1, 1024, 188), dtype=np.float32)], ["f"], "joint_enc") |
| print("joint_enc maxdiff", float(np.abs(m.predict({"encoded": e.numpy()})["f"] - JointEnc(model)(e).detach().numpy()).max())) |
| f, g = torch.randn(1, 640), torch.randn(1, 640) |
| m = conv(Joint(model), (f, g), [ct.TensorType("f", shape=(1, 640), dtype=np.float32), |
| ct.TensorType("g", shape=(1, 640), dtype=np.float32)], ["token_id", "duration_id"], "joint") |
| print("joint", m.predict({"f": f.numpy(), "g": g.numpy()}), [t.item() for t in Joint(model)(f, g)]) |
|
|
