fix(runtime): trim batch padding from CIF peaks before funasr_onnx and funasr_torch Paraformer timestamps

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

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

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(
+        "<AddShift> 28 28\n<LearnRateCoef> 1 [ " + "0 " * 28 + "]\n"
+        "<Rescale> 28 28\n<LearnRateCoef> 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("<i2").tobytes())
+        paths.append(str(path))
+    return directory, paths
+
+
+def timestamp_outputs(feats, feats_len):
+    """BiCifParaformer-style outputs for 你 好, padded to the longest row like the exported graph.
+
+    Each row's three CIF peaks lie inside its own 3 * feats_len upsampled frames; the
+    frames after them are zero padding.
+    """
+    import numpy as np
+
+    batch = len(feats_len)
+    am_scores = np.zeros((batch, 3, 8), dtype=np.float32)
+    am_scores[:, [0, 1, 2], [3, 4, 2]] = 1.0  # 你 好 </s>
+    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")