From 9396b7d86ee1cb70517950822fbea4d6f194abca Mon Sep 17 00:00:00 2001 From: Tanishq Date: Wed, 11 Nov 2026 06:49:26 +0530 Subject: [PATCH] fix(runtime): trim batch padding from CIF peaks before Paraformer timestamps With a timestamp export (4 outputs) and batch_size > 1, the Paraformer wrappers in funasr_onnx (Paraformer, ContextualParaformer and the SeacoParaformer that inherits it) and funasr_torch (Paraformer) passed each row of us_peaks to time_stamp_lfr6_onnx with the batch padding still attached. time_stamp_lfr6_onnx takes the frame count from the array length, so the last token of every file shorter than the longest in its batch ended past the end of its own audio. With two 440 Hz tones of 1 s and 3 s and BiCif-shaped outputs whose peaks lie inside each row's own frames, batch_size=1 gives the 1 s file [[70, 470], [470, 945]], while batch_size=2 gives [[70, 470], [470, 1935]], 935 ms past its 1000 ms end. Pass each row's peaks cut to 3 * feats_len, as BiCifParaformer and SeacoParaformer in funasr/models and the C++ runtime (paraformer-torch.cpp) already do, so batch_size=2 gives the same timestamps as batch_size=1. Extend tests/test_runtime_eval_native.py: for the four wrappers, with the acoustic model session replaced by these outputs, a batched call must return the same results as one call per file. --- .../libtorch/funasr_torch/paraformer_bin.py | 4 +- .../onnxruntime/funasr_onnx/paraformer_bin.py | 8 +- tests/test_runtime_eval_native.py | 94 +++++++++++++++++++ 3 files changed, 100 insertions(+), 6 deletions(-) diff --git a/runtime/python/libtorch/funasr_torch/paraformer_bin.py b/runtime/python/libtorch/funasr_torch/paraformer_bin.py index dc9d2a596..93159e85c 100644 --- a/runtime/python/libtorch/funasr_torch/paraformer_bin.py +++ b/runtime/python/libtorch/funasr_torch/paraformer_bin.py @@ -120,10 +120,10 @@ def __call__(self, wav_content: Union[str, np.ndarray, List[str]], **kwargs) -> pred = sentence_postprocess(pred) asr_res.append({"preds": pred}) else: - for pred, us_peaks_ in zip(preds, us_peaks): + for pred, us_peaks_, feat_len in zip(preds, us_peaks, feats_len): raw_tokens = pred timestamp, timestamp_raw = time_stamp_lfr6_onnx( - us_peaks_, copy.copy(raw_tokens) + us_peaks_[: feat_len * 3], copy.copy(raw_tokens) ) text_proc, timestamp_proc, _ = sentence_postprocess( raw_tokens, timestamp_raw diff --git a/runtime/python/onnxruntime/funasr_onnx/paraformer_bin.py b/runtime/python/onnxruntime/funasr_onnx/paraformer_bin.py index 517252cd8..5a2da2189 100644 --- a/runtime/python/onnxruntime/funasr_onnx/paraformer_bin.py +++ b/runtime/python/onnxruntime/funasr_onnx/paraformer_bin.py @@ -131,10 +131,10 @@ def __call__(self, wav_content: Union[str, np.ndarray, List[str]], **kwargs) -> pred = sentence_postprocess(pred) asr_res.append({"preds": pred}) else: - for pred, us_peaks_ in zip(preds, us_peaks): + for pred, us_peaks_, feat_len in zip(preds, us_peaks, feats_len): raw_tokens = pred timestamp, timestamp_raw = time_stamp_lfr6_onnx( - us_peaks_, copy.copy(raw_tokens) + us_peaks_[: feat_len * 3], copy.copy(raw_tokens) ) text_proc, timestamp_proc, _ = sentence_postprocess( raw_tokens, timestamp_raw @@ -384,10 +384,10 @@ def __call__( pred = sentence_postprocess(pred) asr_res.append({"preds": pred}) else: - for pred, us_peaks_ in zip(preds, us_peaks): + for pred, us_peaks_, feat_len in zip(preds, us_peaks, feats_len): raw_tokens = pred timestamp, timestamp_raw = time_stamp_lfr6_onnx( - us_peaks_, copy.copy(raw_tokens) + us_peaks_[: feat_len * 3], copy.copy(raw_tokens) ) text_proc, timestamp_proc, _ = sentence_postprocess( raw_tokens, timestamp_raw diff --git a/tests/test_runtime_eval_native.py b/tests/test_runtime_eval_native.py index e0e615480..ee8d9933b 100644 --- a/tests/test_runtime_eval_native.py +++ b/tests/test_runtime_eval_native.py @@ -8,6 +8,7 @@ import json import os from pathlib import Path +import shutil import subprocess import sys import wave @@ -275,6 +276,99 @@ def test_hotword_wrappers_decode_every_batch( assert model(paths, hotwords) == alone +@pytest.fixture(scope="module") +def timestamp_fixture(hotword_fixture, tmp_path_factory): + np = pytest.importorskip("numpy") + source, _ = hotword_fixture + directory = tmp_path_factory.mktemp("runtime-timestamp") + # The graphs only have to load: the timestamp test replaces the acoustic model's session. + for name, target in [ + ("tokens.json", "tokens.json"), + ("model_eb.onnx", "model_eb.onnx"), + ("model.onnx", "model.onnx"), + ("model_bb.torchscript", "model.torchscript"), + ]: + shutil.copy(source / name, directory / target) + (directory / "config.yaml").write_text( + json.dumps( + { + "frontend_conf": {"n_mels": 4, "lfr_m": 7, "lfr_n": 6, "dither": 0.0}, + "model_conf": {"predictor_bias": 1}, + } + ) + ) + (directory / "am.mvn").write_text( + " 28 28\n 1 [ " + "0 " * 28 + "]\n" + " 28 28\n 1 [ " + "1 " * 28 + "]\n" + ) + paths = [] + for seconds in (1, 3): + path = directory / f"tone {seconds}s.wav" + tone = np.sin(np.arange(16000 * seconds) * 2 * np.pi * 440 / 16000) * 3000 + with wave.open(str(path), "wb") as output: + output.setnchannels(1) + output.setsampwidth(2) + output.setframerate(16000) + output.writeframes(tone.astype(" + us_alphas = np.zeros((batch, 3 * feats.shape[1]), dtype=np.float32) + us_peaks = np.zeros_like(us_alphas) + for row, feat_len in enumerate(np.asarray(feats_len).tolist()): + frames = 3 * feat_len + us_alphas[row, :frames] = 3.0 / frames + us_peaks[row, np.linspace(5, frames - 6, 3).round().astype(int)] = 1.0 + return [am_scores, np.full(batch, 3, dtype=np.int32), us_alphas, us_peaks] + + +@pytest.mark.parametrize( + "backend, wrapper", + [ + ("onnx", "Paraformer"), + ("onnx", "ContextualParaformer"), + ("onnx", "SeacoParaformer"), + ("torch", "Paraformer"), + ], +) +def test_batched_timestamps_ignore_padding(timestamp_fixture, monkeypatch, backend, wrapper): + torch = pytest.importorskip("torch") + directory, paths = timestamp_fixture + location = "libtorch" if backend == "torch" else "onnxruntime" + monkeypatch.syspath_prepend(str(ROOT / "runtime/python" / location)) + module = importlib.import_module("funasr_" + backend + ".paraformer_bin") + model = getattr(module, wrapper)(str(directory), batch_size=2, intra_op_num_threads=1) + if backend == "torch": + model.ort_infer = lambda feats, feats_len: tuple( + torch.from_numpy(output) for output in timestamp_outputs(feats, feats_len) + ) + elif wrapper == "Paraformer": + model.ort_infer = lambda inputs: timestamp_outputs(*inputs) + else: + model.ort_infer_bb = lambda inputs: timestamp_outputs(*inputs[:2]) + hotwords = () if wrapper == "Paraformer" else ("魔搭",) + alone = [model(path, *hotwords)[0] for path in paths] + # The 1 s tone ends at 1000 ms, the 3 s tone at 3000 ms. + assert [row["timestamp"] for row in alone] == [ + [[70, 470], [470, 945]], + [[70, 670], [1450, 2925]], + ] + # Batched with the 3 s tone, the 1 s tone must not take on the padded frames. + assert model(paths, *hotwords) == alone + + @pytest.mark.parametrize("extension", ["flac", "ogg", "mp3"]) def test_onnx_wrappers_read_non_wav_paths(native_fixture, tmp_path, monkeypatch, extension): np = pytest.importorskip("numpy")