diff --git a/.github/workflows/test-text-postprocessing.yml b/.github/workflows/test-text-postprocessing.yml index e76358e8ee..55364def5d 100644 --- a/.github/workflows/test-text-postprocessing.yml +++ b/.github/workflows/test-text-postprocessing.yml @@ -9,6 +9,7 @@ on: - "runtime/python/libtorch/funasr_torch/utils/postprocess_utils.py" - "tests/test_postprocess_hotwords.py" - "tests/test_rich_transcription_postprocess.py" + - "tests/test_sentence_postprocess_timestamps.py" - ".github/workflows/test-text-postprocessing.yml" push: branches: [main] @@ -19,6 +20,7 @@ on: - "runtime/python/libtorch/funasr_torch/utils/postprocess_utils.py" - "tests/test_postprocess_hotwords.py" - "tests/test_rich_transcription_postprocess.py" + - "tests/test_sentence_postprocess_timestamps.py" - ".github/workflows/test-text-postprocessing.yml" permissions: @@ -52,6 +54,7 @@ jobs: python -m pytest -q \ tests/test_postprocess_hotwords.py \ tests/test_rich_transcription_postprocess.py \ + tests/test_sentence_postprocess_timestamps.py \ --junitxml=text-postprocessing-results.xml - uses: actions/upload-artifact@v4 if: always() diff --git a/funasr/utils/postprocess_utils.py b/funasr/utils/postprocess_utils.py index 28c5d4546d..9774300ac9 100644 --- a/funasr/utils/postprocess_utils.py +++ b/funasr/utils/postprocess_utils.py @@ -173,9 +173,10 @@ def sentence_postprocess(words: List[Any], time_stamp: List[List] = None): word_lists = [] word_item = "" ts_lists = [] + kept_indices = [] # wash words lists - for i in words: + for index, i in enumerate(words): word = "" if isinstance(i, str): word = i @@ -186,6 +187,11 @@ def sentence_postprocess(words: List[Any], time_stamp: List[List] = None): continue else: middle_lists.append(word) + kept_indices.append(index) + + # time_stamp has one entry per token, including the ones washed out above + if time_stamp is not None and max(kept_indices, default=-1) < len(time_stamp): + time_stamp = [time_stamp[index] for index in kept_indices] # all chinese characters if isAllChinese(middle_lists): diff --git a/runtime/python/libtorch/funasr_torch/utils/postprocess_utils.py b/runtime/python/libtorch/funasr_torch/utils/postprocess_utils.py index 8dac6b0cad..0390c5da87 100644 --- a/runtime/python/libtorch/funasr_torch/utils/postprocess_utils.py +++ b/runtime/python/libtorch/funasr_torch/utils/postprocess_utils.py @@ -139,9 +139,10 @@ def sentence_postprocess(words: List[Any], time_stamp: List[List] = None): word_lists = [] word_item = "" ts_lists = [] + kept_indices = [] # wash words lists - for i in words: + for index, i in enumerate(words): word = "" if isinstance(i, str): word = i @@ -152,6 +153,11 @@ def sentence_postprocess(words: List[Any], time_stamp: List[List] = None): continue else: middle_lists.append(word) + kept_indices.append(index) + + # time_stamp has one entry per token, including the ones washed out above + if time_stamp is not None and max(kept_indices, default=-1) < len(time_stamp): + time_stamp = [time_stamp[index] for index in kept_indices] # all chinese characters if isAllChinese(middle_lists): diff --git a/runtime/python/onnxruntime/funasr_onnx/utils/postprocess_utils.py b/runtime/python/onnxruntime/funasr_onnx/utils/postprocess_utils.py index ef8b206a9c..7158aba60f 100644 --- a/runtime/python/onnxruntime/funasr_onnx/utils/postprocess_utils.py +++ b/runtime/python/onnxruntime/funasr_onnx/utils/postprocess_utils.py @@ -141,9 +141,10 @@ def sentence_postprocess(words: List[Any], time_stamp: List[List] = None): word_lists = [] word_item = "" ts_lists = [] + kept_indices = [] # wash words lists - for i in words: + for index, i in enumerate(words): word = "" if isinstance(i, str): word = i @@ -154,6 +155,11 @@ def sentence_postprocess(words: List[Any], time_stamp: List[List] = None): continue else: middle_lists.append(word) + kept_indices.append(index) + + # time_stamp has one entry per token, including the ones washed out above + if time_stamp is not None and max(kept_indices, default=-1) < len(time_stamp): + time_stamp = [time_stamp[index] for index in kept_indices] # all chinese characters if isAllChinese(middle_lists): diff --git a/tests/test_sentence_postprocess_timestamps.py b/tests/test_sentence_postprocess_timestamps.py new file mode 100644 index 0000000000..5ad8d775c1 --- /dev/null +++ b/tests/test_sentence_postprocess_timestamps.py @@ -0,0 +1,81 @@ +import importlib.util +import unittest +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[1] +POSTPROCESS_UTILS = ( + "funasr/utils/postprocess_utils.py", + "runtime/python/onnxruntime/funasr_onnx/utils/postprocess_utils.py", + "runtime/python/libtorch/funasr_torch/utils/postprocess_utils.py", +) + + +def _load_sentence_postprocess(relative_path): + spec = importlib.util.spec_from_file_location( + relative_path.replace("/", "_").replace(".py", ""), REPO_ROOT / relative_path + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module.sentence_postprocess + + +class SentencePostprocessTimestampTest(unittest.TestCase): + def assert_timestamps(self, tokens, expected_words, expected_timestamps): + timestamps = [[i * 100, (i + 1) * 100] for i in range(len(tokens))] + for path in POSTPROCESS_UTILS: + with self.subTest(path=path): + _, result_timestamps, words = _load_sentence_postprocess(path)( + tokens, timestamps + ) + self.assertEqual(words, expected_words) + self.assertEqual(result_timestamps, expected_timestamps) + + def test_chinese_words_keep_their_timestamps_after_unk(self): + self.assert_timestamps( + ["你", "", "好", "吗"], + ["你", "好", "吗"], + [[0, 100], [200, 300], [300, 400]], + ) + + def test_english_words_keep_their_timestamps_after_unk(self): + self.assert_timestamps( + ["hello", "", "wor@@", "ld"], + ["hello", "world"], + [[0, 100], [200, 400]], + ) + + def test_mixed_words_keep_their_timestamps_after_unk(self): + self.assert_timestamps( + ["你", "", "hello", "好"], + ["你", "hello", "好"], + [[0, 100], [200, 300], [300, 400]], + ) + + def test_tokens_without_a_timestamp_at_the_end_are_ignored(self): + for path in POSTPROCESS_UTILS: + with self.subTest(path=path): + _, timestamps, words = _load_sentence_postprocess(path)( + ["你", "好", ""], [[0, 100], [100, 200]] + ) + self.assertEqual(words, ["你", "好"]) + self.assertEqual(timestamps, [[0, 100], [100, 200]]) + + def test_too_few_timestamps_do_not_crash_when_a_token_is_filtered(self): + cases = ( + (["你", "", "好", "吗"], ["你", "好", "吗"]), + (["hello", "", "there", "big", "world"], ["hello", "there", "big", "world"]), + (["你", "", "hello", "好", "world"], ["你", "hello", "好", "world"]), + ) + for path in POSTPROCESS_UTILS: + for tokens, expected_words in cases: + with self.subTest(path=path, tokens=tokens): + timestamps = [[i * 100, (i + 1) * 100] for i in range(len(tokens) - 1)] + _, result_timestamps, words = _load_sentence_postprocess(path)( + tokens, timestamps + ) + self.assertEqual(words, expected_words) + self.assertEqual(len(result_timestamps), len(expected_words)) + + +if __name__ == "__main__": + unittest.main()