From fe89f75d7113c2eb8e29f623adfe7764023274d7 Mon Sep 17 00:00:00 2001 From: Tanishq Date: Tue, 3 Nov 2026 21:56:30 +0530 Subject: [PATCH] fix(runtime): decode every batch in ContextualParaformer and SeacoParaformer ContextualParaformer.__call__ in funasr_onnx and funasr_torch computes the hotword embedding once per call, but the batch loop overwrote it with its batched copy, so from the second batch on the copy had one more dimension. In funasr_onnx, ONNX Runtime rejected the 4-D bias_embed ("Invalid rank for input: bias_embed Got: 4 Expected: 3"), the ONNXRuntimeError branch logged "input wav is silence or noise" and the batch added no result: with the default batch_size=1, model([a, b, c], hotwords) returned only the result for a. In funasr_torch, repeat() raised RuntimeError on the second batch. SeacoParaformer inherits __call__ in both packages. Build the batched copy in a new variable, so every batch gets a [batch, num_hotwords, 512] bias_embed. Extend tests/test_runtime_eval_native.py with tiny embedder and backbone graphs, exported to ONNX and traced to TorchScript, and check that both classes in both packages return one result per file when a call has more files than batch_size. --- .../libtorch/funasr_torch/paraformer_bin.py | 8 +- .../onnxruntime/funasr_onnx/paraformer_bin.py | 6 +- tests/test_runtime_eval_native.py | 112 ++++++++++++++++++ 3 files changed, 120 insertions(+), 6 deletions(-) diff --git a/runtime/python/libtorch/funasr_torch/paraformer_bin.py b/runtime/python/libtorch/funasr_torch/paraformer_bin.py index fe405753cd..dc9d2a5966 100644 --- a/runtime/python/libtorch/funasr_torch/paraformer_bin.py +++ b/runtime/python/libtorch/funasr_torch/paraformer_bin.py @@ -342,14 +342,16 @@ def __call__( for beg_idx in range(0, waveform_nums, self.batch_size): end_idx = min(waveform_nums, beg_idx + self.batch_size) feats, feats_len = self.extract_feat(waveform_list[beg_idx:end_idx]) - bias_embed = torch.unsqueeze(bias_embed, 0).repeat(feats.shape[0], 1, 1) + batch_bias_embed = torch.unsqueeze(bias_embed, 0).repeat(feats.shape[0], 1, 1) try: with torch.no_grad(): if int(self.device_id) == -1: - outputs = self.bb_infer(feats, feats_len, bias_embed) + outputs = self.bb_infer(feats, feats_len, batch_bias_embed) am_scores, valid_token_lens = outputs[0], outputs[1] else: - outputs = self.bb_infer(feats.cuda(), feats_len.cuda(), bias_embed.cuda()) + outputs = self.bb_infer( + feats.cuda(), feats_len.cuda(), batch_bias_embed.cuda() + ) am_scores, valid_token_lens = outputs[0].cpu(), outputs[1].cpu() except: # logging.warning(traceback.format_exc()) diff --git a/runtime/python/onnxruntime/funasr_onnx/paraformer_bin.py b/runtime/python/onnxruntime/funasr_onnx/paraformer_bin.py index 94538033c9..d822e73183 100644 --- a/runtime/python/onnxruntime/funasr_onnx/paraformer_bin.py +++ b/runtime/python/onnxruntime/funasr_onnx/paraformer_bin.py @@ -374,10 +374,10 @@ def __call__( for beg_idx in range(0, waveform_nums, self.batch_size): end_idx = min(waveform_nums, beg_idx + self.batch_size) feats, feats_len = self.extract_feat(waveform_list[beg_idx:end_idx]) - bias_embed = np.expand_dims(bias_embed, axis=0) - bias_embed = np.repeat(bias_embed, feats.shape[0], axis=0) + batch_bias_embed = np.expand_dims(bias_embed, axis=0) + batch_bias_embed = np.repeat(batch_bias_embed, feats.shape[0], axis=0) try: - outputs = self.bb_infer(feats, feats_len, bias_embed) + outputs = self.bb_infer(feats, feats_len, batch_bias_embed) am_scores, valid_token_lens = outputs[0], outputs[1] if len(outputs) == 4: diff --git a/tests/test_runtime_eval_native.py b/tests/test_runtime_eval_native.py index c3b60a8ccb..62747fb2c1 100644 --- a/tests/test_runtime_eval_native.py +++ b/tests/test_runtime_eval_native.py @@ -161,3 +161,115 @@ def test_real_cli_outputs_and_audio_duration(native_fixture, tmp_path, backend, } ) ) + + +@pytest.fixture(scope="module") +def hotword_fixture(tmp_path_factory): + torch = pytest.importorskip("torch") + np = pytest.importorskip("numpy") + onnx = pytest.importorskip("onnx") + pytest.importorskip("onnxruntime") + pytest.importorskip("kaldi_native_fbank") + directory = tmp_path_factory.mktemp("runtime-hotword") + tokens = ["", "", "", "你", "好", "魔", "搭", ""] + + class EmbedderFixture(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("weight", torch.ones(len(tokens), 512)) + + def forward(self, hotword): + # [num_hotwords, 10] -> [10, num_hotwords, 512], like the exported model_eb + return torch.nn.functional.embedding(hotword.long(), self.weight).transpose(0, 1) + + class BackboneFixture(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("bias", torch.tensor([0.0, 0.0, 0.0, 5.0, 0.0, 0.0, 0.0, 0.0])) + + def forward(self, speech, speech_lengths, bias_embed): + # bias_embed [batch, num_hotwords, 512] is added as zeros; every frame predicts token 3 + hotword_bias = bias_embed.sum(dim=(1, 2), keepdim=True) * 0.0 + logits = speech[..., :1] * 0.0 + hotword_bias + self.bias + return logits, (speech_lengths // 4).to(torch.int32) + + embedder, backbone = EmbedderFixture().eval(), BackboneFixture().eval() + hotword = torch.zeros(3, 10, dtype=torch.int32) + speech = torch.zeros(2, 12, 4) + lengths = torch.tensor([12, 8], dtype=torch.int32) + bias_embed = torch.zeros(2, 3, 512) + torch.jit.trace(embedder, (hotword,)).save(str(directory / "model_eb.torchscript")) + torch.jit.trace(backbone, (speech, lengths, bias_embed)).save( + str(directory / "model_bb.torchscript") + ) + torch.onnx.export( + embedder, + (hotword,), + str(directory / "model_eb.onnx"), + dynamo=False, + input_names=["hotword"], + output_names=["hw_embed"], + dynamic_axes={"hotword": {0: "num_hotwords"}, "hw_embed": {1: "num_hotwords"}}, + opset_version=17, + ) + torch.onnx.export( + backbone, + (speech, lengths, bias_embed), + str(directory / "model.onnx"), + dynamo=False, + input_names=["speech", "speech_lengths", "bias_embed"], + output_names=["logits", "token_num"], + dynamic_axes={ + "speech": {0: "batch_size", 1: "feats_length"}, + "speech_lengths": {0: "batch_size"}, + "bias_embed": {0: "batch_size", 1: "num_hotwords"}, + "logits": {0: "batch_size", 1: "logits_length"}, + "token_num": {0: "batch_size"}, + }, + opset_version=17, + ) + for name in ("model_eb.onnx", "model.onnx"): + onnx.checker.check_model(str(directory / name)) + (directory / "tokens.json").write_text(json.dumps(tokens)) + (directory / "config.yaml").write_text( + json.dumps( + { + "frontend_conf": {"n_mels": 4, "lfr_m": 1, "lfr_n": 1, "dither": 0.0}, + "model_conf": {"predictor_bias": 1}, + } + ) + ) + (directory / "am.mvn").write_text( + " 4 4\n 1 [ 0 0 0 0 ]\n" + " 4 4\n 1 [ 1 1 1 1 ]\n" + ) + paths = [] + for index, count in enumerate([4000, 8000, 12000]): + path = directory / f"hotword {index}.wav" + signal = (np.sin(np.arange(count) * 2 * np.pi * 440 / 16000) * 3000).astype("