diff --git a/funasr/auto/auto_model.py b/funasr/auto/auto_model.py index 3dbb7a4f2..4f41dcc21 100644 --- a/funasr/auto/auto_model.py +++ b/funasr/auto/auto_model.py @@ -1164,6 +1164,7 @@ def inference_with_vad(self, input, input_len=None, **cfg): { "start": vadsegment[0], "end": vadsegment[1], + "text": rest["text"], "sentence": rest["text"], "timestamp": ts, } diff --git a/funasr/models/paraformer/model.py b/funasr/models/paraformer/model.py index 3f5b01517..cb0e7a77d 100644 --- a/funasr/models/paraformer/model.py +++ b/funasr/models/paraformer/model.py @@ -555,7 +555,11 @@ def inference( is_use_lm = ( kwargs.get("lm_weight", 0.0) > 0.00001 and kwargs.get("lm_file", None) is not None ) - pred_timestamp = kwargs.get("pred_timestamp", False) + pred_timestamp = ( + kwargs["pred_timestamp"] + if "pred_timestamp" in kwargs + else kwargs.get("output_timestamp", False) + ) if self.beam_search is None and (is_use_lm or is_use_ctc): logging.info("enable beam_search") self.init_beam_search(**kwargs) diff --git a/tests/test_paraformer_timestamp_contract.py b/tests/test_paraformer_timestamp_contract.py new file mode 100644 index 000000000..de2f04544 --- /dev/null +++ b/tests/test_paraformer_timestamp_contract.py @@ -0,0 +1,88 @@ +"""Regression tests for Paraformer timestamp flag precedence.""" + +import importlib +import unittest +from unittest.mock import MagicMock, patch + +import torch + + +class _Tokenizer: + def ids2tokens(self, token_ids): + return ["你" for _ in token_ids] + + def tokens2text(self, tokens): + return "".join(tokens) + + +class TestParaformerTimestampContract(unittest.TestCase): + def _make_paraformer(self): + from funasr.models.paraformer.model import Paraformer + + model = Paraformer.__new__(Paraformer) + torch.nn.Module.__init__(model) + model.beam_search = None + model.sos = 1 + model.eos = 2 + model.blank_id = 0 + model.encode = MagicMock( + return_value=(torch.zeros((1, 2, 2)), torch.tensor([2])) + ) + model.calc_predictor = MagicMock( + return_value=( + torch.zeros((1, 1, 2)), + torch.tensor([1.0]), + torch.ones((1, 2)), + torch.ones((1, 2)), + ) + ) + model.cal_decoder_with_predictor = MagicMock( + return_value=( + torch.tensor([[[0.0, 0.0, 0.0, 4.0]]]), + torch.tensor([1]), + ) + ) + return model + + @staticmethod + def _sentence_postprocess(tokens, timestamp=None): + text = "".join(tokens) + if timestamp is None: + return text, None + return text, [[0, 100]], None + + def test_pred_timestamp_precedence_and_output_timestamp_fallback(self): + paraformer_module = importlib.import_module("funasr.models.paraformer.model") + cases = ( + ({"output_timestamp": False}, False), + ({"output_timestamp": True}, True), + ({"pred_timestamp": True, "output_timestamp": False}, True), + ({"pred_timestamp": False, "output_timestamp": True}, False), + ) + + with patch.object( + paraformer_module, + "ts_prediction_lfr6_standard", + return_value=("", [[0, 100]]), + ), patch.object( + paraformer_module.postprocess_utils, + "sentence_postprocess", + side_effect=self._sentence_postprocess, + ): + for timestamp_kwargs, expected_timestamp in cases: + with self.subTest(timestamp_kwargs=timestamp_kwargs): + results, _ = self._make_paraformer().inference( + torch.zeros((1, 2, 2)), + data_lengths=torch.tensor([[2]]), + key=["utt"], + tokenizer=_Tokenizer(), + frontend=None, + device="cpu", + data_type="fbank", + **timestamp_kwargs, + ) + self.assertEqual("timestamp" in results[0], expected_timestamp) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_punc_model_none.py b/tests/test_punc_model_none.py index db0823382..b8dc22256 100644 --- a/tests/test_punc_model_none.py +++ b/tests/test_punc_model_none.py @@ -443,6 +443,57 @@ def mock_inference(data, *args, **kwargs): ) mock_distribute_spk.assert_called_once() + @patch("funasr.auto.auto_model.distribute_spk") + @patch("funasr.auto.auto_model.postprocess") + @patch("funasr.auto.auto_model.sv_chunk") + @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") + def test_no_timestamp_speaker_fallback_keeps_sentence_and_adds_text( + self, + mock_prep, + mock_load, + mock_slice, + mock_sv_chunk, + mock_postprocess, + _mock_distribute_spk, + ): + """Speaker VAD fallback exposes documented text without dropping sentence.""" + am = self._make_auto_model( + punc_model=None, spk_model=MagicMock(), spk_mode="vad_segment" + ) + am.cb_model = MagicMock(return_value=np.array([0])) + inference_calls = [] + results_seq = [ + [{"key": "test_utt", "value": [[0, 1000]]}], + [{"text": "hello"}], + [{"spk_embedding": torch.tensor([[1.0, 0.0]])}], + ] + + def mock_inference(data, *args, **kwargs): + inference_calls.append(kwargs) + return results_seq.pop(0) + + am.inference = MagicMock(side_effect=mock_inference) + mock_prep.return_value = ( + ["test_utt"], + [np.zeros(16000, dtype=np.float32)], + ) + mock_load.return_value = np.zeros(16000, dtype=np.float32) + mock_slice.return_value = ([np.zeros(16000, dtype=np.float32)], [16000]) + mock_sv_chunk.return_value = [ + [0.0, 1.0, np.zeros(16000, dtype=np.float32)] + ] + mock_postprocess.return_value = [{"start": 0.0, "end": 1.0, "spk": 0}] + + results = am.inference_with_vad("dummy_input") + + self.assertIs(inference_calls[1]["output_timestamp"], True) + self.assertIs(inference_calls[1]["return_time_stamps"], True) + sentence = results[0]["sentence_info"][0] + self.assertEqual(sentence["text"], "hello") + self.assertEqual(sentence["sentence"], "hello") + @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")