diff --git a/.github/workflows/test-numpy-compatibility.yml b/.github/workflows/test-numpy-compatibility.yml index 054a5f9aa7..288f6926c8 100644 --- a/.github/workflows/test-numpy-compatibility.yml +++ b/.github/workflows/test-numpy-compatibility.yml @@ -21,6 +21,8 @@ on: - "tests/test_punc_model_none.py" - "tests/test_paraformer_empty_token_batch.py" - "tests/test_paraformer_hotword_file.py" + - "funasr/models/ct_transformer_streaming/model.py" + - "tests/test_ct_transformer_streaming_empty_input.py" - "tests/test_pcm_input_format.py" - "tests/test_load_audio_bytes.py" - "funasr/frontends/**" @@ -56,6 +58,8 @@ on: - "tests/test_punc_model_none.py" - "tests/test_paraformer_empty_token_batch.py" - "tests/test_paraformer_hotword_file.py" + - "funasr/models/ct_transformer_streaming/model.py" + - "tests/test_ct_transformer_streaming_empty_input.py" - "tests/test_pcm_input_format.py" - "tests/test_load_audio_bytes.py" - "funasr/frontends/**" @@ -129,6 +133,7 @@ jobs: tests/test_punc_model_none.py \ tests/test_paraformer_empty_token_batch.py \ tests/test_paraformer_hotword_file.py \ + tests/test_ct_transformer_streaming_empty_input.py \ tests/test_cli.py \ --junitxml=numpy-results.xml - uses: actions/upload-artifact@v4 diff --git a/funasr/models/ct_transformer_streaming/model.py b/funasr/models/ct_transformer_streaming/model.py index 06add5f2e4..49a29cd0a0 100644 --- a/funasr/models/ct_transformer_streaming/model.py +++ b/funasr/models/ct_transformer_streaming/model.py @@ -102,6 +102,9 @@ def inference( if cache is None: cache = {} assert len(data_in) == 1 + if not data_in[0] or (isinstance(data_in[0], str) and not data_in[0].strip()): + meta_data = {"batch_data_time": -1} + return [{"key": key[0] if key else "", "text": "", "punc_array": None}], meta_data if len(cache) == 0: cache["pre_text"] = [] diff --git a/tests/test_ct_transformer_streaming_empty_input.py b/tests/test_ct_transformer_streaming_empty_input.py new file mode 100644 index 0000000000..32432170d0 --- /dev/null +++ b/tests/test_ct_transformer_streaming_empty_input.py @@ -0,0 +1,77 @@ +"""Regression tests for empty text in the realtime CT-Transformer punctuation model.""" + +import unittest + +import torch + +from funasr.auto.auto_model import AutoModel +from funasr.models.ct_transformer_streaming.model import CTTransformerStreaming +from funasr.tokenizer.char_tokenizer import CharTokenizer + +PUNC_LIST = ["", "_", ",", "。", "?", "、"] +TOKEN_LIST = ["", "", "", "", "跨", "境", "河", "流", "是"] + + +def _tiny_streaming_punc_model(): + """Build a small randomly initialised CTTransformerStreaming that never adds punctuation.""" + model = CTTransformerStreaming( + encoder="SANMVadEncoder", + encoder_conf={ + "input_size": 8, + "output_size": 8, + "attention_heads": 2, + "linear_units": 16, + "num_blocks": 2, + "dropout_rate": 0.0, + "positional_dropout_rate": 0.0, + "attention_dropout_rate": 0.0, + "input_layer": "pe", + "kernel_size": 3, + "sanm_shfit": 0, + "padding_idx": 0, + }, + vocab_size=len(TOKEN_LIST), + punc_list=PUNC_LIST, + embed_unit=8, + att_unit=8, + sentence_end_id=3, + ) + with torch.no_grad(): + model.decoder.weight.zero_() + model.decoder.bias.zero_() + model.decoder.bias[PUNC_LIST.index("_")] = 1.0 + return model.eval() + + +class TestCTTransformerStreamingEmptyInput(unittest.TestCase): + def setUp(self): + self.model = _tiny_streaming_punc_model() + self.tokenizer = CharTokenizer(token_list=TOKEN_LIST, unk_symbol="") + + def test_empty_text_returns_empty_result(self): + for text in ("", " "): + with self.subTest(text=text): + results, _ = self.model.inference( + [text], key=["utt"], tokenizer=self.tokenizer, cache={}, device="cpu" + ) + self.assertEqual(results, [{"key": "utt", "text": "", "punc_array": None}]) + + def test_empty_segment_keeps_pending_words(self): + am = AutoModel.__new__(AutoModel) + am.model = self.model + am.vad_model = None + am.punc_model = None + am.kwargs = {"device": "cpu", "disable_pbar": True, "tokenizer": self.tokenizer} + + cache = {} + self.assertEqual(am.generate(input="跨境河流", cache=cache)[0]["text"], "跨境河流") + self.assertEqual(cache["pre_text"], ["跨", "境", "河", "流"]) + + self.assertEqual(am.generate(input="", cache=cache)[0]["text"], "") + self.assertEqual(cache["pre_text"], ["跨", "境", "河", "流"]) + + self.assertEqual(am.generate(input="是", cache=cache)[0]["text"], "是") + + +if __name__ == "__main__": + unittest.main()