fix(auto-model): preserve explicit timestamp and speaker text

Signed-off-by: John6666 <[email protected]>

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()

tests/test_punc_model_none.py

@@ -443,6 +443,57 @@ class TestPuncModelNone(unittest.TestCase):
         )
         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")