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")