From b2fb22d1addcfa0284b3ca0c63bf2e02d02a6ad4 Mon Sep 17 00:00:00 2001 From: John6666 <186692226+John6666cat@users.noreply.github.com> Date: Sat, 29 Aug 2026 17:09:19 +0900 Subject: [PATCH] fix(paraformer): trim timestamp predictor padding per sample Signed-off-by: John6666 <186692226+John6666cat@users.noreply.github.com> --- funasr/models/paraformer/model.py | 18 ++- tests/test_paraformer_timestamp_padding.py | 129 +++++++++++++++++++++ 2 files changed, 145 insertions(+), 2 deletions(-) create mode 100644 tests/test_paraformer_timestamp_padding.py diff --git a/funasr/models/paraformer/model.py b/funasr/models/paraformer/model.py index 3f5b01517..382880faf 100644 --- a/funasr/models/paraformer/model.py +++ b/funasr/models/paraformer/model.py @@ -671,9 +671,23 @@ def inference( text_postprocessed = tokenizer.tokens2text(token) if pred_timestamp: + timestamp_pre_peak_index = pre_peak_index[i] + timestamp_alphas = alphas[i] + predictor = getattr(self, "predictor", None) + if getattr(predictor, "tail_mask", None) is True: + timestamp_len = int(encoder_out_lens[i].item()) + if float(getattr(predictor, "tail_threshold", 0.0)) > 0.0: + timestamp_len += 1 + timestamp_len = min( + timestamp_len, + timestamp_pre_peak_index.shape[-1], + timestamp_alphas.shape[-1], + ) + timestamp_pre_peak_index = timestamp_pre_peak_index[:timestamp_len] + timestamp_alphas = timestamp_alphas[:timestamp_len] timestamp_str, timestamp = ts_prediction_lfr6_standard( - pre_peak_index[i], - alphas[i], + timestamp_pre_peak_index, + timestamp_alphas, copy.copy(token), vad_offset=kwargs.get("begin_time", 0), upsample_rate=1, diff --git a/tests/test_paraformer_timestamp_padding.py b/tests/test_paraformer_timestamp_padding.py new file mode 100644 index 000000000..0169ff28a --- /dev/null +++ b/tests/test_paraformer_timestamp_padding.py @@ -0,0 +1,129 @@ +import unittest +from types import SimpleNamespace +from unittest import mock + +import torch + +import funasr.models.paraformer.model as paraformer_model +from funasr.models.paraformer.model import Paraformer + + +class _DummyTokenizer: + def ids2tokens(self, ids): + return ["x" for _ in ids] + + def tokens2text(self, tokens): + return "".join(tokens) + + +class _DummyParaformer(Paraformer): + def __init__(self, predictor, encoder_lens, predictor_width): + torch.nn.Module.__init__(self) + if predictor is not None: + self.predictor = predictor + self.encoder_lens = torch.tensor(encoder_lens, dtype=torch.long) + if isinstance(predictor_width, tuple): + self.alphas_width, self.pre_peak_width = predictor_width + else: + self.alphas_width = self.pre_peak_width = predictor_width + self.beam_search = None + self.ctc = None + self.sos = 1 + self.eos = 2 + self.blank_id = 0 + + def encode(self, speech, speech_lengths): + batch = len(self.encoder_lens) + max_len = int(self.encoder_lens.max().item()) + return torch.zeros(batch, max_len, 2), self.encoder_lens + + def calc_predictor(self, encoder_out, encoder_out_lens): + batch = len(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), + ) + + def cal_decoder_with_predictor( + self, encoder_out, encoder_out_lens, pre_acoustic_embeds, pre_token_length + ): + batch = len(encoder_out_lens) + decoder_out = torch.zeros(batch, 1, 4) + decoder_out[:, :, 3] = 1.0 + return decoder_out, pre_token_length + + +class TestParaformerTimestampPadding(unittest.TestCase): + 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])) + return "", [[0, 100]] + + with mock.patch.object( + paraformer_model, "ts_prediction_lfr6_standard", side_effect=fake_timestamp + ), mock.patch.object( + paraformer_model.postprocess_utils, + "sentence_postprocess", + side_effect=lambda token, timestamp: ("".join(token), timestamp, None), + ): + model.inference( + torch.zeros(len(encoder_lens), max(encoder_lens), 2), + data_lengths=torch.tensor(encoder_lens, dtype=torch.long), + key=[f"utt-{i}" for i in range(len(encoder_lens))], + tokenizer=_DummyTokenizer(), + frontend=None, + data_type="fbank", + device="cpu", + pred_timestamp=True, + ) + return seen + + def test_tail_mask_true_uses_per_sample_encoder_extent_plus_tail(self): + predictor = SimpleNamespace(tail_mask=True, tail_threshold=0.45) + self.assertEqual( + self._timestamp_helper_shapes(predictor, [4, 8], 9), + [(5, 5), (9, 9)], + ) + + def test_tail_mask_false_preserves_full_predictor_extent(self): + predictor = SimpleNamespace(tail_mask=False, tail_threshold=0.45) + self.assertEqual( + self._timestamp_helper_shapes(predictor, [4, 8], 9), + [(9, 9), (9, 9)], + ) + + def test_predictor_without_tail_mask_preserves_full_extent(self): + predictor = SimpleNamespace(tail_threshold=0.45) + self.assertEqual( + self._timestamp_helper_shapes(predictor, [4, 8], 9), + [(9, 9), (9, 9)], + ) + + def test_missing_predictor_attribute_preserves_full_extent(self): + self.assertEqual( + self._timestamp_helper_shapes(None, [4, 8], 9), + [(9, 9), (9, 9)], + ) + + def test_zero_tail_threshold_does_not_add_tail_frame(self): + predictor = SimpleNamespace(tail_mask=True, tail_threshold=0.0) + self.assertEqual( + self._timestamp_helper_shapes(predictor, [4, 8], 8), + [(4, 4), (8, 8)], + ) + + def test_trim_length_is_clamped_to_both_timestamp_tensor_widths(self): + predictor = SimpleNamespace(tail_mask=True, tail_threshold=0.45) + self.assertEqual( + self._timestamp_helper_shapes(predictor, [8], (7, 6)), + [(6, 6)], + ) + + +if __name__ == "__main__": + unittest.main()