diff --git a/.github/workflows/test-funasr-onnx-package.yml b/.github/workflows/test-funasr-onnx-package.yml index 796ede71ae..58dfd29a8c 100644 --- a/.github/workflows/test-funasr-onnx-package.yml +++ b/.github/workflows/test-funasr-onnx-package.yml @@ -9,6 +9,7 @@ on: - "runtime/python/libtorch/funasr_torch/utils/postprocess_utils.py" - "tests/test_funasr_onnx_release.py" - "tests/test_funasr_onnx_installed.py" + - "tests/test_funasr_onnx_paraformer_online_overlap.py" - "tests/test_rich_transcription_postprocess.py" - "tests/test_rich_transcription_event_segments.py" push: @@ -21,6 +22,7 @@ on: - "runtime/python/libtorch/funasr_torch/utils/postprocess_utils.py" - "tests/test_funasr_onnx_release.py" - "tests/test_funasr_onnx_installed.py" + - "tests/test_funasr_onnx_paraformer_online_overlap.py" - "tests/test_rich_transcription_postprocess.py" - "tests/test_rich_transcription_event_segments.py" workflow_dispatch: @@ -95,6 +97,9 @@ jobs: - name: Verify installed wrapper error paths run: python tests/test_funasr_onnx_installed.py + - name: Verify streaming Paraformer overlap cache + run: python tests/test_funasr_onnx_paraformer_online_overlap.py + - name: Verify source distribution error paths run: | python -m pip install --force-reinstall --no-deps dist/funasr-onnx/*.tar.gz diff --git a/runtime/python/onnxruntime/funasr_onnx/paraformer_online_bin.py b/runtime/python/onnxruntime/funasr_onnx/paraformer_online_bin.py index 81b991d4b5..33bea03f23 100644 --- a/runtime/python/onnxruntime/funasr_onnx/paraformer_online_bin.py +++ b/runtime/python/onnxruntime/funasr_onnx/paraformer_online_bin.py @@ -119,14 +119,19 @@ def add_overlap_chunk(self, feats: np.ndarray, cache: dict = {}): # process last chunk overlap_feats = np.concatenate((cache["feats"], feats), axis=1) if cache["is_final"]: - cache["feats"] = overlap_feats[:, -self.chunk_size[0] :, :] + cache["feats"] = self._tail_frames(overlap_feats, self.chunk_size[0]) if not cache["last_chunk"]: padding_length = sum(self.chunk_size) - overlap_feats.shape[1] overlap_feats = np.pad(overlap_feats, ((0, 0), (0, padding_length), (0, 0))) else: - cache["feats"] = overlap_feats[:, -(self.chunk_size[0] + self.chunk_size[2]) :, :] + cache["feats"] = self._tail_frames(overlap_feats, self.chunk_size[0] + self.chunk_size[2]) return overlap_feats + @staticmethod + def _tail_frames(feats: np.ndarray, num_frames: int) -> np.ndarray: + # feats[:, -0:] would return every frame, so compute the start index explicitly. + return feats[:, max(feats.shape[1] - num_frames, 0) :, :] + def __call__(self, audio_in: np.ndarray, **kwargs): waveforms = np.expand_dims(audio_in, axis=0) param_dict = kwargs.get("param_dict", dict()) diff --git a/tests/test_funasr_onnx_paraformer_online_overlap.py b/tests/test_funasr_onnx_paraformer_online_overlap.py new file mode 100644 index 0000000000..ac9afab37c --- /dev/null +++ b/tests/test_funasr_onnx_paraformer_online_overlap.py @@ -0,0 +1,62 @@ +"""Run against an installed funasr-onnx package.""" + +import importlib.util +import unittest + +import numpy as np + +if importlib.util.find_spec("funasr_onnx") is None: + raise unittest.SkipTest("Install the standalone funasr-onnx distribution first.") + +from funasr_onnx.paraformer_online_bin import Paraformer + + +def make_model(chunk_size): + model = Paraformer.__new__(Paraformer) + model.chunk_size = chunk_size + return model + + +def make_cache(chunk_size, is_final, last_chunk=False, dims=2): + left = np.full((1, chunk_size[0] + chunk_size[2], dims), -1.0, dtype=np.float32) + return {"feats": left, "is_final": is_final, "last_chunk": last_chunk} + + +def frames(n, dims=2): + return np.arange(n * dims, dtype=np.float32).reshape(1, n, dims) + + +class AddOverlapChunkTest(unittest.TestCase): + def test_zero_left_context_keeps_no_frames_on_final_chunk(self): + chunk_size = [0, 10, 5] + model = make_model(chunk_size) + cache = make_cache(chunk_size, is_final=True) + + model.add_overlap_chunk(frames(10), cache) + + self.assertEqual(cache["feats"].shape[1], 0) + + def test_zero_left_and_right_context_keeps_no_frames(self): + chunk_size = [0, 10, 0] + model = make_model(chunk_size) + cache = make_cache(chunk_size, is_final=False) + + model.add_overlap_chunk(frames(10), cache) + + self.assertEqual(cache["feats"].shape[1], 0) + + def test_nonzero_context_keeps_tail_frames(self): + chunk_size = [5, 10, 5] + model = make_model(chunk_size) + + cache = make_cache(chunk_size, is_final=False) + overlap = model.add_overlap_chunk(frames(10), cache) + np.testing.assert_array_equal(cache["feats"], overlap[:, -10:, :]) + + cache = make_cache(chunk_size, is_final=True) + overlap = model.add_overlap_chunk(frames(10), cache) + np.testing.assert_array_equal(cache["feats"], overlap[:, 20 - 5 : 20, :]) + + +if __name__ == "__main__": + unittest.main()