File size: 3,995 Bytes
0d06934 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | """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
# HAAS_EXPORT_DIR: dir containing the haas_onnx package (scripts/ of Sanami/haas-ja-onnx)
# HAAS_NEMO_DIR: the extracted Haas-ja-v1.0.nemo (model_config.yaml, model_weights.ckpt)
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 # noqa: F401
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) # [U=1, B=1, H]
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) # joint_net = activation, dropout, linear
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)])
|