diff --git a/funasr/auto/auto_model.py b/funasr/auto/auto_model.py index 55f9cf9dd..03d394873 100644 --- a/funasr/auto/auto_model.py +++ b/funasr/auto/auto_model.py @@ -1238,7 +1238,11 @@ def inference_with_vad(self, input, input_len=None, **cfg): ) result["text"] = punctuated_surface or punc_res[0]["text"] - timestamp_text = punc_input_text + # Timestamp units need VAD boundary spaces; punctuation surface text does not. + timestamp_text = ( + re.sub(r"<\|[^|]*\|>", "", raw_text).strip() + if raw_text is not None else None + ) sentence_timestamps = result.get("timestamp", []) punc_alignment_failed = False if punc_res is not None: diff --git a/tests/test_punc_model_none.py b/tests/test_punc_model_none.py index 4e14eb705..42dfe34e6 100644 --- a/tests/test_punc_model_none.py +++ b/tests/test_punc_model_none.py @@ -700,6 +700,115 @@ def test_sentence_timestamp_handles_unsized_punc_array(self, mock_prep, mock_loa self.assertEqual(results[0]["sentence_info"][0]["text"], ["你", "好"]) + + def _check_vad_chunk_alignment( + self, chunks, expected_texts, punc_input, *, speaker=False, + raw_text=False, english=False, timestamps=None, + ): + """Keep model responses fixed; exercise real VAD merging and speaker assignment.""" + am = self._make_auto_model( + punc_model=MagicMock(), + spk_model=MagicMock() if speaker else None, + spk_mode="punc_segment", + ) + if speaker: + am.cb_model = MagicMock(return_value=np.arange(len(chunks))) + if timestamps is None: + timestamps = [ + [[i * 200, (i + 1) * 200] for i in range(len(text.split()))] + for text in chunks + ] + segments = [] + offset = 0 + for stamps in timestamps: + segments.append([offset, offset + stamps[-1][1]]) + offset = segments[-1][1] + order = iter(sorted(range(len(chunks)), key=lambda i: segments[i][1] - segments[i][0])) + punc_array = [] + for i, stamps in enumerate(timestamps): + punc_array.extend([1] * (len(stamps) - 1) + [3 if i == len(chunks) - 1 else 2]) + top_text = " ".join(expected_texts) + + def infer(data, input_len=None, model=None, kwargs=None, **cfg): + if model is am.vad_model: + return [{"key": "chunk_alignment", "value": [s[:] for s in segments]}] + if model is am.model: + self.assertEqual(len(data), 1) + i = next(order) + return [{"text": chunks[i], "timestamp": [t[:] for t in timestamps[i]]}] + if model is am.spk_model: + return [{"spk_embedding": torch.ones(1, 4)} for _ in data] + if model is am.punc_model: + self.assertEqual(data, punc_input) + return [{"text": top_text, "punc_array": punc_array}] + self.fail("Unexpected model") + + am.inference = infer + result = am.inference_with_vad( + np.zeros(offset * 16, dtype=np.float32), + sentence_timestamp=True, return_raw_text=raw_text, en_post_proc=english, + )[0] + self.assertEqual(result["text"], top_text) + sentences = result["sentence_info"] + self.assertEqual([s["text"] for s in sentences], expected_texts) + expected_all = [] + for i, sentence in enumerate(sentences): + expected = [[a + segments[i][0], b + segments[i][0]] for a, b in timestamps[i]] + expected_all.extend(expected) + self.assertEqual(sentence["timestamp"], expected) + self.assertEqual((sentence["start"], sentence["end"]), (expected[0][0], expected[-1][1])) + if speaker: + self.assertEqual(sentence["spk"], i) + self.assertEqual(result["timestamp"], expected_all) + if raw_text: + self.assertEqual(result["raw_text"], " ".join(chunks)) + return result + + def test_chinese_vad_boundaries_preserve_timestamp_units(self): + chunks = ["今 天 天 气 很 好", "我 们 去 公 园 吧"] + for speaker in (False, True): + for raw_text in (False, True): + with self.subTest(speaker=speaker, raw_text=raw_text): + self._check_vad_chunk_alignment( + chunks, ["今天天气很好,", "我们去公园吧。"], + "今 天 天 气 很 好我 们 去 公 园 吧", + speaker=speaker, raw_text=raw_text, + ) + + def test_chinese_vad_boundaries_restore_sorted_segments(self): + for speaker in (False, True): + with self.subTest(speaker=speaker): + self._check_vad_chunk_alignment( + ["今 天 天 气 很 好", "我 们 去 公 园 吧", "一 起 出 发"], + ["今天天气很好,", "我们去公园吧,", "一起出发。"], + "今 天 天 气 很 好我 们 去 公 园 吧一 起 出 发", + speaker=speaker, + ) + + def test_chinese_vad_boundaries_remove_rich_tags_only_for_alignment(self): + self._check_vad_chunk_alignment( + ["<|zh|><|NEUTRAL|>今 天 天 气 很 好", "<|zh|>我 们 去 公 园 吧"], + ["今天天气很好,", "我们去公园吧。"], + "今 天 天 气 很 好我 们 去 公 园 吧", + speaker=True, raw_text=True, + ) + + def test_mixed_vad_boundaries_preserve_timestamp_units(self): + self._check_vad_chunk_alignment( + ["你 好 hello", "world 我 来"], + ["你好 hello,", " world 我来。"], + "你 好 hello world 我 来", + raw_text=True, + ) + + def test_english_vad_boundaries_preserve_timestamp_units(self): + self._check_vad_chunk_alignment( + ["hello world", "from funasr"], + ["hello world,", "from funasr."], + "hello world from funasr", + english=True, raw_text=True, + ) + class TestCTTransformerPunctuation(unittest.TestCase): """Keep CT-Transformer text and punctuation-array endings consistent."""