fix(paraformer): correct timestamp alpha and peak argument order

funasr/models/paraformer/model.py

@@ -690,9 +690,9 @@ class Paraformer(torch.nn.Module):
                             timestamp_pre_peak_index = timestamp_pre_peak_index[:timestamp_len]
                             timestamp_alphas = timestamp_alphas[:timestamp_len]
                         timestamp_str, timestamp = ts_prediction_lfr6_standard(
-                            timestamp_pre_peak_index,
-                            timestamp_alphas,
-                            copy.copy(token),
+                            us_alphas=timestamp_alphas,
+                            us_peaks=timestamp_pre_peak_index,
+                            char_list=copy.copy(token),
                             vad_offset=kwargs.get("begin_time", 0),
                             upsample_rate=1,
                         )

tests/test_paraformer_timestamp_contract.py

@@ -1,4 +1,4 @@
-"""Regression tests for Paraformer timestamp flag precedence."""
+"""Regression tests for Paraformer timestamp flags and predictor arguments."""
 
 import importlib
 import unittest
@@ -83,6 +83,62 @@ class TestParaformerTimestampContract(unittest.TestCase):
                     )
                     self.assertEqual("timestamp" in results[0], expected_timestamp)
 
+    def test_inference_timestamp_boundaries_with_real_helper(self):
+        from funasr.models.paraformer.cif_predictor import cif
+        from funasr.utils import timestamp_tools
+
+        cases = (
+            # Three fires already provide the two tokens' boundary positions.
+            ([0.0] + [0.25] * 12, [[150, 390], [390, 780]], False),
+            # Two fires require the helper to normalize alphas to three boundaries.
+            (
+                [0.1, 0.1, 0.4, 0.6] + [0.0] * 4 + [0.1] * 10,
+                [[90, 510], [510, 1080]],
+                True,
+            ),
+        )
+        for weights, expected, needs_fallback in cases:
+            for offset in (0, 1000):
+                with self.subTest(needs_fallback=needs_fallback, offset=offset):
+                    model = self._make_paraformer()
+                    alphas = torch.tensor([weights])
+                    width = alphas.shape[1]
+                    hidden = torch.zeros(1, width, 2)
+                    _, peaks = cif(hidden, alphas, threshold=1.0)
+                    model.encode.return_value = (hidden, torch.tensor([width]))
+                    model.calc_predictor.return_value = (
+                        torch.zeros(1, 2, 2),
+                        torch.tensor([2.0]),
+                        alphas,
+                        peaks,
+                    )
+                    model.cal_decoder_with_predictor.return_value = (
+                        torch.tensor([[[0.0, 0.0, 0.0, 4.0]] * 2]),
+                        torch.tensor([2]),
+                    )
+                    with patch.object(
+                        timestamp_tools,
+                        "cif_wo_hidden",
+                        wraps=timestamp_tools.cif_wo_hidden,
+                    ) as fallback:
+                        results, _ = model.inference(
+                            hidden,
+                            data_lengths=torch.tensor([[width]]),
+                            key=["utt"],
+                            tokenizer=_Tokenizer(),
+                            frontend=None,
+                            device="cpu",
+                            data_type="fbank",
+                            pred_timestamp=True,
+                            begin_time=offset,
+                        )
+                    self.assertEqual(results[0]["text"], "你 你")
+                    self.assertEqual(
+                        results[0]["timestamp"],
+                        [[start + offset, end + offset] for start, end in expected],
+                    )
+                    self.assertEqual(fallback.call_count, int(needs_fallback))
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_paraformer_timestamp_padding.py

@@ -42,8 +42,8 @@ class _DummyParaformer(Paraformer):
         return (
             torch.zeros(batch, 1, 2),
             torch.ones(batch),
-            torch.zeros(batch, self.alphas_width),
-            torch.zeros(batch, self.pre_peak_width),
+            torch.arange(batch * self.alphas_width).reshape(batch, -1).float() / 100,
+            1 + torch.arange(batch * self.pre_peak_width).reshape(batch, -1).float() / 100,
         )
 
     def cal_decoder_with_predictor(
@@ -60,8 +60,18 @@ class TestParaformerTimestampPadding(unittest.TestCase):
         model = _DummyParaformer(predictor, encoder_lens, predictor_width)
         seen = []
 
-        def fake_timestamp(arg0, arg1, char_list, **kwargs):
-            seen.append((arg0.shape[-1], arg1.shape[-1]))
+        def fake_timestamp(us_alphas, us_peaks, char_list, **kwargs):
+            sample = len(seen)
+            alpha_len, peak_len = us_alphas.shape[-1], us_peaks.shape[-1]
+            torch.testing.assert_close(
+                us_alphas,
+                (sample * model.alphas_width + torch.arange(alpha_len)).float() / 100,
+            )
+            torch.testing.assert_close(
+                us_peaks,
+                1 + (sample * model.pre_peak_width + torch.arange(peak_len)).float() / 100,
+            )
+            seen.append((alpha_len, peak_len))
             return "", [[0, 100]]
 
         with mock.patch.object(