From ecaf729fe45346c3c1ad4746a9dae4a7e2475aa7 Mon Sep 17 00:00:00 2001 From: phoenixray2000 Date: Sun, 25 Oct 2026 03:55:12 +0800 Subject: [PATCH] fix(auto): preserve timestamp token boundaries across VAD segments --- funasr/auto/auto_model.py | 2 +- tests/test_punc_model_none.py | 60 +++++++++++++++++++++++++++++++++++ 2 files changed, 61 insertions(+), 1 deletion(-) diff --git a/funasr/auto/auto_model.py b/funasr/auto/auto_model.py index 55f9cf9dd8..3376f1a2e5 100644 --- a/funasr/auto/auto_model.py +++ b/funasr/auto/auto_model.py @@ -1238,7 +1238,7 @@ 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_text = raw_text 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 4e14eb705d..e0a106e7e0 100644 --- a/tests/test_punc_model_none.py +++ b/tests/test_punc_model_none.py @@ -275,6 +275,66 @@ def mock_inference(data, *args, **kwargs): ], ) + @patch("funasr.auto.auto_model.slice_padding_audio_samples") + @patch("funasr.auto.auto_model.load_audio_text_image_video") + @patch("funasr.auto.auto_model.prepare_data_iterator") + @patch("funasr.auto.auto_model.sv_chunk") + @patch("funasr.auto.auto_model.postprocess") + def test_sentence_timestamps_preserve_token_boundaries_across_vad_segments( + self, mock_postprocess, mock_sv_chunk, mock_prep, mock_load, mock_slice + ): + mock_prep.return_value = (["test_utt"], [np.zeros(40000, dtype=np.float32)]) + mock_load.return_value = np.zeros(40000, dtype=np.float32) + mock_slice.return_value = ([np.zeros(16000, dtype=np.float32)], [16000]) + mock_sv_chunk.side_effect = lambda segments: segments + mock_postprocess.return_value = [[0.0, 1.0, 0], [1.5, 2.5, 1]] + + for with_speakers in (False, True): + with self.subTest(with_speakers=with_speakers): + am = self._make_auto_model( + punc_model=MagicMock(), + spk_model=MagicMock() if with_speakers else None, + spk_mode="punc_segment" if with_speakers else None, + ) + results_seq = [ + [{"key": "test_utt", "value": [[0, 1000], [1500, 2500]]}], + [{"text": "客 户", "timestamp": [[0, 500], [500, 1000]]}], + ] + if with_speakers: + am.cb_model = MagicMock(return_value=np.array([0, 1])) + results_seq.append([{"spk_embedding": torch.tensor([[1.0, 0.0]])}]) + results_seq.append([{"text": "需 求", "timestamp": [[0, 500], [500, 1000]]}]) + if with_speakers: + results_seq.append([{"spk_embedding": torch.tensor([[0.0, 1.0]])}]) + results_seq.append([{"text": "客户,需求。", "punc_array": [1, 2, 1, 3]}]) + am.inference = MagicMock(side_effect=lambda *args, **kwargs: results_seq.pop(0)) + + result = am.inference_with_vad( + "dummy_input", sentence_timestamp=True, return_raw_text=True + )[0] + + expected = [ + { + "text": "客户,", + "raw_text": "客 户", + "start": 0, + "end": 1000, + "timestamp": [[0, 500], [500, 1000]], + }, + { + "text": "需求。", + "raw_text": "需 求", + "start": 1500, + "end": 2500, + "timestamp": [[1500, 2000], [2000, 2500]], + }, + ] + if with_speakers: + for index, sentence in enumerate(expected): + sentence["spk"] = index + self.assertEqual(result["text"], "客户,需求。") + self.assertEqual(result["sentence_info"], expected) + @patch("funasr.auto.auto_model.slice_padding_audio_samples") @patch("funasr.auto.auto_model.load_audio_text_image_video") @patch("funasr.auto.auto_model.prepare_data_iterator")