diff --git a/runtime/python/libtorch/funasr_torch/paraformer_bin.py b/runtime/python/libtorch/funasr_torch/paraformer_bin.py index fe405753c..dc9d2a596 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 94538033c..d822e7318 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 c3b60a8cc..62747fb2c 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("