diff --git a/funasr/models/contextual_paraformer/model.py b/funasr/models/contextual_paraformer/model.py index a8f23e9260..9531afcd30 100644 --- a/funasr/models/contextual_paraformer/model.py +++ b/funasr/models/contextual_paraformer/model.py @@ -457,7 +457,11 @@ def inference( ) pre_token_length = pre_token_length.round().long() if torch.max(pre_token_length) < 1: - return [] + # no tokens for any utterance: keep one empty result per key + b = encoder_out.size(0) + if tokenizer is None: + return [{"key": key[i], "token_int": []} for i in range(b)], meta_data + return [{"key": key[i], "text": ""} for i in range(b)], meta_data decoder_outs = self.cal_decoder_with_predictor( encoder_out, diff --git a/funasr/models/paraformer/model.py b/funasr/models/paraformer/model.py index 86808be0dc..c20343aeae 100644 --- a/funasr/models/paraformer/model.py +++ b/funasr/models/paraformer/model.py @@ -616,19 +616,24 @@ def inference( ) pre_token_length = pre_token_length.round().long() + b = encoder_out.size(0) + if isinstance(key[0], (list, tuple)): + key = key[0] + if len(key) < b: + key = key * b if torch.max(pre_token_length) < 1: - return [] + # no tokens for any utterance: keep one empty result per key + if tokenizer is None: + return [{"key": key[i], "token_int": []} for i in range(b)], meta_data + if pred_timestamp: + return [{"key": key[i], "text": "", "timestamp": []} for i in range(b)], meta_data + return [{"key": key[i], "text": ""} for i in range(b)], meta_data decoder_outs = self.cal_decoder_with_predictor( encoder_out, encoder_out_lens, pre_acoustic_embeds, pre_token_length ) decoder_out, ys_pad_lens = decoder_outs[0], decoder_outs[1] results = [] - b, n, d = decoder_out.size() - if isinstance(key[0], (list, tuple)): - key = key[0] - if len(key) < b: - key = key * b for i in range(b): x = encoder_out[i, : encoder_out_lens[i], :] am_scores = decoder_out[i, : pre_token_length[i], :] diff --git a/funasr/models/seaco_paraformer/model.py b/funasr/models/seaco_paraformer/model.py index 6a57875308..c1c3ba55a2 100644 --- a/funasr/models/seaco_paraformer/model.py +++ b/funasr/models/seaco_paraformer/model.py @@ -483,7 +483,13 @@ def inference( pre_acoustic_embeds, pre_token_length = predictor_outs[0], predictor_outs[1] pre_token_length = pre_token_length.round().long() if torch.max(pre_token_length) < 1: - return ([],) + # no tokens for any utterance: keep one empty result per key + b = encoder_out.size(0) + if tokenizer is None: + return [{"key": key[i], "token_int": []} for i in range(b)], meta_data + if self.predictor_name == "CifPredictorV3": + return [{"key": key[i], "text": "", "timestamp": []} for i in range(b)], meta_data + return [{"key": key[i], "text": ""} for i in range(b)], meta_data decoder_out = self._seaco_decode_with_ASF( encoder_out, diff --git a/tests/test_paraformer_empty_token_batch.py b/tests/test_paraformer_empty_token_batch.py new file mode 100644 index 0000000000..7771224e36 --- /dev/null +++ b/tests/test_paraformer_empty_token_batch.py @@ -0,0 +1,146 @@ +import unittest + +import numpy as np +import torch + +from funasr.auto.auto_model import AutoModel +from funasr.frontends.wav_frontend import WavFrontend +from funasr.models.contextual_paraformer.model import ContextualParaformer +from funasr.models.paraformer.model import Paraformer +from funasr.models.seaco_paraformer.model import SeacoParaformer + + +class _DummyTokenizer: + def ids2tokens(self, ids): + return ["x" for _ in ids] + + def tokens2text(self, tokens): + return "".join(tokens) + + +class _ZeroTokenMixin: + """Encoder/predictor stand-ins where every utterance predicts zero tokens.""" + + def _setup(self): + torch.nn.Module.__init__(self) + self.dummy = torch.nn.Parameter(torch.zeros(1)) + self.beam_search = None + self.ctc = None + self.sos = 1 + self.eos = 2 + self.blank_id = 0 + + def encode(self, speech, speech_lengths): + return torch.zeros(speech.size(0), 4, 2), torch.full((speech.size(0),), 4) + + def calc_predictor(self, encoder_out, encoder_out_lens): + batch = encoder_out.size(0) + return ( + torch.zeros(batch, 1, 2), + torch.full((batch,), 0.4), + torch.zeros(batch, 4), + torch.zeros(batch, 4), + ) + + def generate_hotwords_list(self, *args, **kwargs): + return None + + def cal_decoder_with_predictor(self, *args, **kwargs): + raise AssertionError("decoder must not run when no token is predicted") + + _seaco_decode_with_ASF = cal_decoder_with_predictor + + +class _ZeroTokenParaformer(_ZeroTokenMixin, Paraformer): + def __init__(self): + self._setup() + + +class _ZeroTokenContextualParaformer(_ZeroTokenMixin, ContextualParaformer): + def __init__(self): + self._setup() + + +class _ZeroTokenSeacoParaformer(_ZeroTokenMixin, SeacoParaformer): + def __init__(self, predictor_name="CifPredictorV3"): + self._setup() + self.predictor_name = predictor_name + + +def _frontend(): + return WavFrontend(cmvn_file=None, dither=0.0) + + +def _inference(model, keys, **kwargs): + return model.inference( + [np.zeros(1600, dtype=np.float32) for _ in keys], + key=list(keys), + frontend=_frontend(), + device="cpu", + **kwargs, + ) + + +class TestParaformerEmptyTokenBatch(unittest.TestCase): + def test_paraformer_returns_one_keyed_result_per_utterance(self): + results, meta_data = _inference( + _ZeroTokenParaformer(), ["a", "b"], tokenizer=_DummyTokenizer() + ) + self.assertEqual(results, [{"key": "a", "text": ""}, {"key": "b", "text": ""}]) + self.assertIn("batch_data_time", meta_data) + + def test_paraformer_keeps_timestamp_field(self): + results, _ = _inference( + _ZeroTokenParaformer(), ["a", "b"], tokenizer=_DummyTokenizer(), output_timestamp=True + ) + self.assertEqual( + results, + [{"key": "a", "text": "", "timestamp": []}, {"key": "b", "text": "", "timestamp": []}], + ) + + def test_paraformer_without_tokenizer_returns_empty_token_ids(self): + results, _ = _inference(_ZeroTokenParaformer(), ["a", "b"]) + self.assertEqual(results, [{"key": "a", "token_int": []}, {"key": "b", "token_int": []}]) + + def test_contextual_paraformer_returns_one_keyed_result_per_utterance(self): + results, _ = _inference( + _ZeroTokenContextualParaformer(), ["a", "b"], tokenizer=_DummyTokenizer() + ) + self.assertEqual(results, [{"key": "a", "text": ""}, {"key": "b", "text": ""}]) + + def test_seaco_paraformer_returns_one_keyed_result_per_utterance(self): + results, _ = _inference( + _ZeroTokenSeacoParaformer(), ["a", "b"], tokenizer=_DummyTokenizer() + ) + self.assertEqual( + results, + [{"key": "a", "text": "", "timestamp": []}, {"key": "b", "text": "", "timestamp": []}], + ) + + results, _ = _inference( + _ZeroTokenSeacoParaformer("CifPredictorV2"), ["a"], tokenizer=_DummyTokenizer() + ) + self.assertEqual(results, [{"key": "a", "text": ""}]) + + def test_auto_model_keeps_every_utterance_of_the_batch(self): + am = AutoModel.__new__(AutoModel) + am.model = _ZeroTokenParaformer() + am.kwargs = { + "batch_size": 2, + "disable_pbar": True, + "device": "cpu", + "frontend": _frontend(), + "tokenizer": _DummyTokenizer(), + } + + res = AutoModel.inference( + am, + [np.zeros(1600, dtype=np.float32), np.zeros(1600, dtype=np.float32)], + key=["a", "b"], + ) + + self.assertEqual(res, [{"key": "a", "text": ""}, {"key": "b", "text": ""}]) + + +if __name__ == "__main__": + unittest.main()