From cbce0a2ca1bf8813fa41e3edde7e05fc2b7bbc90 Mon Sep 17 00:00:00 2001
From: Xu Xiang <10363155+xx205@users.noreply.github.com>
Date: Thu, 17 Sep 2026 13:41:40 +0800
Subject: [PATCH] fix(paraformer): correct timestamp alpha and peak argument
 order

---
 funasr/models/paraformer/model.py           |  6 +--
 tests/test_paraformer_timestamp_contract.py | 58 ++++++++++++++++++++-
 tests/test_paraformer_timestamp_padding.py  | 18 +++++--
 3 files changed, 74 insertions(+), 8 deletions(-)

diff --git a/funasr/models/paraformer/model.py b/funasr/models/paraformer/model.py
index 31f0b46d5..86808be0d 100644
--- a/funasr/models/paraformer/model.py
+++ b/funasr/models/paraformer/model.py
@@ -690,9 +690,9 @@ def inference(
                             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,
                         )
diff --git a/tests/test_paraformer_timestamp_contract.py b/tests/test_paraformer_timestamp_contract.py
index de2f04544..8a333d98a 100644
--- a/tests/test_paraformer_timestamp_contract.py
+++ b/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 @@ def test_pred_timestamp_precedence_and_output_timestamp_fallback(self):
                     )
                     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()
diff --git a/tests/test_paraformer_timestamp_padding.py b/tests/test_paraformer_timestamp_padding.py
index 0169ff28a..a81f382b2 100644
--- a/tests/test_paraformer_timestamp_padding.py
+++ b/tests/test_paraformer_timestamp_padding.py
@@ -42,8 +42,8 @@ def calc_predictor(self, encoder_out, encoder_out_lens):
         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 @@ def _timestamp_helper_shapes(self, predictor, encoder_lens, predictor_width):
         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(