diff --git a/funasr/auto/auto_model.py b/funasr/auto/auto_model.py index a92309198..6546f7c80 100644 --- a/funasr/auto/auto_model.py +++ b/funasr/auto/auto_model.py @@ -8,6 +8,7 @@ import copy import torch import random +import re import string import logging import os.path @@ -52,6 +53,182 @@ def _resolve_ncpu(config, fallback=4): return max(value, 1) +def _join_vad_texts(texts): + """Remove rich tags and join VAD text without adding spaces between Chinese chunks.""" + cleaned = [re.sub(r"<\|[^|]*\|>", "", text).strip() for text in texts] + cleaned = [text for text in cleaned if text] + if not cleaned: + return "" + joined = cleaned[0] + for text in cleaned[1:]: + separator = "" + if not ("\u3400" <= joined[-1] <= "\u9fff" and "\u3400" <= text[0] <= "\u9fff"): + separator = " " + joined += separator + text + return joined + + +def _get_punc_tokens(text, punc_array, punc_model): + """Return the surface tokens represented by a CT-Transformer punctuation array.""" + try: + from funasr.models.ct_transformer.utils import split_words + + tokens = split_words( + text, + jieba_usr_dict=getattr(punc_model, "jieba_usr_dict", None), + ) + except Exception: + return None + + expanded_tokens = [] + for token in tokens: + if token and "\u0e00" <= token[0] <= "\u9fa5" and len(token) > 1: + expanded_tokens.extend(token) + else: + expanded_tokens.append(token) + try: + punc_length = len(punc_array) + except TypeError: + return None + if len(expanded_tokens) != punc_length: + return None + return expanded_tokens + + +def _punctuate_surface_text(text, punc_array, punc_model): + """Insert predicted punctuation without changing the ASR surface text.""" + tokens = _get_punc_tokens(text, punc_array, punc_model) + if tokens is None: + return None + + spans = _surface_token_spans(text, tokens) + if spans is None: + return None + + parts = [] + cursor = 0 + for token, punc_id, (_, end) in zip(tokens, punc_array, spans): + parts.append(text[cursor:end]) + parts.append(_punc_symbol(punc_id, token, punc_model)) + cursor = end + parts.append(text[cursor:]) + return "".join(parts) + + +def _punc_symbol(punc_id, token, punc_model): + """Return the punctuation character represented by a model punctuation ID.""" + punc_list = getattr(punc_model, "punc_list", None) + fallback_punc = {1: "", 2: ",", 3: "。", 4: "?", 5: "、"} + punc_id = int(punc_id) + try: + punctuation = punc_list[punc_id] + except (IndexError, TypeError): + punctuation = fallback_punc.get(punc_id, "") + if punctuation == "_": + punctuation = "" + if punctuation and token[0].isascii(): + punctuation = {",": ",", "。": ".", "?": "?", "、": ","}.get( + punctuation, punctuation + ) + return punctuation + + +def _surface_token_spans(text, tokens): + """Map punctuation tokens back to exact spans in the original surface text.""" + spans = [] + cursor = 0 + for token in tokens: + while cursor < len(text) and text[cursor].isspace(): + cursor += 1 + surface_token = text[cursor : cursor + len(token)] + if surface_token.casefold() != token.casefold(): + return None + spans.append((cursor, cursor + len(token))) + cursor += len(token) + if text[cursor:].strip(): + return None + return spans + + +def _merge_timestamp_units(text, words, timestamps, punc_array, punc_model): + """Merge timestamp/BPE units to the punctuation model's surface tokens.""" + expanded_tokens = _get_punc_tokens(text, punc_array, punc_model) + if expanded_tokens is None: + return None + + def normalize(value): + return "".join(value.split()).casefold() + + merged_timestamps = [] + word_index = 0 + for token in expanded_tokens: + token_text = normalize(token) + if not token_text: + return None + start_index = word_index + word_text = "" + while word_index < len(words) and len(word_text) < len(token_text): + word_text += normalize(words[word_index]) + word_index += 1 + if word_text != token_text or start_index == word_index: + return None + if not all( + isinstance(item, (list, tuple)) and len(item) >= 2 + for item in timestamps[start_index:word_index] + ): + return None + merged_timestamps.append( + [timestamps[start_index][0], timestamps[word_index - 1][1]] + ) + + if word_index != len(words): + return None + return " ".join(expanded_tokens), merged_timestamps + + +def _timestamp_sentences_from_surface( + text, timestamps, punc_array, punc_model, return_raw_text=False +): + """Build sentence timestamps while preserving the exact ASR surface text.""" + tokens = _get_punc_tokens(text, punc_array, punc_model) + if tokens is None or len(tokens) != len(timestamps): + return None + spans = _surface_token_spans(text, tokens) + if spans is None: + return None + + sentences = [] + sentence_start = 0 + for index, (token, punc_id) in enumerate(zip(tokens, punc_array)): + punctuation = _punc_symbol(punc_id, token, punc_model) + if not punctuation: + continue + raw_sentence = text[spans[sentence_start][0] : spans[index][1]].strip() + sentence = { + "text": raw_sentence + punctuation, + "start": timestamps[sentence_start][0], + "end": timestamps[index][1], + "timestamp": timestamps[sentence_start : index + 1], + } + if return_raw_text: + sentence["raw_text"] = raw_sentence + sentences.append(sentence) + sentence_start = index + 1 + + if sentence_start < len(tokens): + raw_sentence = text[spans[sentence_start][0] : spans[-1][1]].strip() + sentence = { + "text": raw_sentence, + "start": timestamps[sentence_start][0], + "end": timestamps[-1][1], + "timestamp": timestamps[sentence_start:], + } + if return_raw_text: + sentence["raw_text"] = raw_sentence + sentences.append(sentence) + return sentences + + def _get_import_errors(): """Internal: get import errors.""" try: @@ -800,18 +977,76 @@ def inference_with_vad(self, input, input_len=None, **cfg): if not len(result["text"].strip()): continue return_raw_text = kwargs.get("return_raw_text", False) + aligned_words = result.get("words") + aligned_timestamps = result.get("timestamp") + aligned_word_text = None + if ( + isinstance(aligned_words, list) + and aligned_words + and all(isinstance(word, str) and word.strip() for word in aligned_words) + and isinstance(aligned_timestamps, list) + and len(aligned_words) == len(aligned_timestamps) + ): + aligned_word_text = " ".join(aligned_words) + # step.3 compute punc model raw_text = None + punc_input_text = None punc_res = None + punc_array = None if self.punc_model is not None and "timestamps" not in result: deep_update(self.punc_kwargs, cfg) + raw_text = copy.copy(result["text"]) + punc_input_text = _join_vad_texts( + item.get("text", "") for item in restored_data + ) punc_res = self.inference( - result["text"], model=self.punc_model, kwargs=self.punc_kwargs, **cfg + punc_input_text, + model=self.punc_model, + kwargs=self.punc_kwargs, + **cfg, ) - raw_text = copy.copy(result["text"]) if return_raw_text: result["raw_text"] = raw_text - result["text"] = punc_res[0]["text"] + punc_array = punc_res[0].get("punc_array") + punctuated_surface = None + if aligned_word_text is not None: + punctuated_surface = _punctuate_surface_text( + punc_input_text, punc_array, self.punc_model + ) + result["text"] = punctuated_surface or punc_res[0]["text"] + + timestamp_text = punc_input_text + sentence_timestamps = result.get("timestamp", []) + if punc_res is not None: + try: + punc_length = len(punc_array) + except TypeError: + punc_length = -1 + punc_array = None + if aligned_word_text is not None and punc_length == len(aligned_words): + timestamp_text = aligned_word_text + elif aligned_word_text is not None and punc_length > 0: + merged_units = _merge_timestamp_units( + punc_input_text, + aligned_words, + aligned_timestamps, + punc_array, + self.punc_model, + ) + if merged_units is not None: + timestamp_text, sentence_timestamps = merged_units + if punc_array is not None and punc_length != len(sentence_timestamps): + punc_array = None + surface_sentence_list = None + if aligned_word_text is not None and punc_array is not None: + surface_sentence_list = _timestamp_sentences_from_surface( + punc_input_text, + sentence_timestamps, + punc_array, + self.punc_model, + return_raw_text=return_raw_text, + ) # speaker embedding cluster after resorted if self.spk_model is not None and kwargs.get("return_spk_res", True): @@ -874,18 +1109,20 @@ def inference_with_vad(self, input, input_len=None, **cfg): "Missing punc_model, which is required for punc_segment speaker diarization." ) sentence_list = [] + elif surface_sentence_list is not None: + sentence_list = surface_sentence_list elif kwargs.get("en_post_proc", False): sentence_list = timestamp_sentence_en( - punc_res[0]["punc_array"], - result["timestamp"], - raw_text, + punc_array, + sentence_timestamps, + timestamp_text, return_raw_text=return_raw_text, ) else: sentence_list = timestamp_sentence( - punc_res[0]["punc_array"], - result["timestamp"], - raw_text, + punc_array, + sentence_timestamps, + timestamp_text, return_raw_text=return_raw_text, ) distribute_spk(sentence_list, sv_output) @@ -898,19 +1135,21 @@ def inference_with_vad(self, input, input_len=None, **cfg): "punc_model is required for sentence_timestamp, skipping sentence segmentation." ) sentence_list = [] + elif surface_sentence_list is not None: + sentence_list = surface_sentence_list else: if kwargs.get("en_post_proc", False): sentence_list = timestamp_sentence_en( - punc_res[0]["punc_array"], - result["timestamp"], - raw_text, + punc_array, + sentence_timestamps, + timestamp_text, return_raw_text=return_raw_text, ) else: sentence_list = timestamp_sentence( - punc_res[0]["punc_array"], - result["timestamp"], - raw_text, + punc_array, + sentence_timestamps, + timestamp_text, return_raw_text=return_raw_text, ) result["sentence_info"] = sentence_list diff --git a/funasr/models/ct_transformer/model.py b/funasr/models/ct_transformer/model.py index ab70258a0..1ee0bf566 100644 --- a/funasr/models/ct_transformer/model.py +++ b/funasr/models/ct_transformer/model.py @@ -416,11 +416,15 @@ def inference( new_mini_sentence_punc_out = new_mini_sentence_punc[:-1] + [ self.sentence_end_id ] + if len(punctuations): + punctuations[-1] = self.sentence_end_id elif new_mini_sentence[-1] == ",": new_mini_sentence_out = new_mini_sentence[:-1] + "." new_mini_sentence_punc_out = new_mini_sentence_punc[:-1] + [ self.sentence_end_id ] + if len(punctuations): + punctuations[-1] = self.sentence_end_id elif ( new_mini_sentence[-1] != "。" and new_mini_sentence[-1] != "?" @@ -431,7 +435,7 @@ def inference( self.sentence_end_id ] if len(punctuations): - punctuations[-1] = 2 + punctuations[-1] = self.sentence_end_id elif ( new_mini_sentence[-1] != "." and new_mini_sentence[-1] != "?" @@ -442,7 +446,7 @@ def inference( self.sentence_end_id ] if len(punctuations): - punctuations[-1] = 2 + punctuations[-1] = self.sentence_end_id # keep a punctuations array for punc segment if punc_array is None: punc_array = punctuations diff --git a/tests/test_punc_model_none.py b/tests/test_punc_model_none.py index 3ab90a08c..26e354c72 100644 --- a/tests/test_punc_model_none.py +++ b/tests/test_punc_model_none.py @@ -1,8 +1,9 @@ -"""Tests for issue #2839: punc_model=None should not cause UnboundLocalError.""" +"""Regression tests for AutoModel VAD punctuation and sentence timestamps.""" import unittest from unittest.mock import MagicMock, patch import numpy as np +import torch class TestPuncModelNone(unittest.TestCase): @@ -16,6 +17,9 @@ def _make_auto_model(self, punc_model=None, spk_model=None, spk_mode=None): am.model = MagicMock() am.vad_model = MagicMock() am.punc_model = punc_model + if punc_model is not None: + punc_model.jieba_usr_dict = None + punc_model.punc_list = ["", "_", ",", "。", "?", "、"] am.punc_kwargs = {} am.spk_model = spk_model am.cb_model = None @@ -114,6 +118,321 @@ def mock_inference(data, input_len=None, model=None, kwargs=None, **cfg): # Text should be updated with punctuated version self.assertEqual(results[0]["text"], "Hello, world.") + @patch("funasr.auto.auto_model.slice_padding_audio_samples") + @patch("funasr.auto.auto_model.load_audio_text_image_video") + @patch("funasr.auto.auto_model.prepare_data_iterator") + def test_sentence_timestamp_uses_asr_words_for_unspaced_text( + self, mock_prep, mock_load, mock_slice + ): + """Sentence timestamps must align SenseVoice words with token timestamps.""" + am = self._make_auto_model(punc_model=MagicMock()) + vad_result = [{"key": "test_utt", "value": [[0, 2000]]}] + asr_result = [ + { + "text": "<|zh|><|NEUTRAL|><|Speech|><|woitn|>你好世界", + "timestamp": [[0, 500], [500, 1000], [1000, 1500], [1500, 2000]], + "words": ["你", "好", "世", "界"], + } + ] + punc_result = [{"text": "你好,世界。", "punc_array": [1, 2, 1, 3]}] + results_seq = [vad_result, asr_result, punc_result] + + def mock_inference(data, *args, **kwargs): + if len(results_seq) == 1: + self.assertEqual(data, "你好世界") + return results_seq.pop(0) + + am.inference = MagicMock(side_effect=mock_inference) + mock_prep.return_value = (["test_utt"], [np.zeros(2000, dtype=np.float32)]) + mock_load.return_value = np.zeros(32000, dtype=np.float32) + mock_slice.return_value = ([np.zeros(32000, dtype=np.float32)], [32000]) + + results = am.inference_with_vad( + "dummy_input", sentence_timestamp=True, return_raw_text=True + ) + + self.assertEqual( + results[0]["raw_text"], + "<|zh|><|NEUTRAL|><|Speech|><|woitn|>你好世界", + ) + self.assertEqual( + results[0]["sentence_info"], + [ + { + "text": "你好,", + "start": 0, + "end": 1000, + "timestamp": [[0, 500], [500, 1000]], + "raw_text": "你好", + }, + { + "text": "世界。", + "start": 1000, + "end": 2000, + "timestamp": [[1000, 1500], [1500, 2000]], + "raw_text": "世界", + }, + ], + ) + + @patch("funasr.auto.auto_model.slice_padding_audio_samples") + @patch("funasr.auto.auto_model.load_audio_text_image_video") + @patch("funasr.auto.auto_model.prepare_data_iterator") + def test_sentence_timestamp_aligns_words_across_vad_segments( + self, mock_prep, mock_load, mock_slice + ): + """Every SenseVoice VAD prefix must be excluded from punctuation alignment.""" + am = self._make_auto_model(punc_model=MagicMock()) + tag = "<|zh|><|NEUTRAL|><|Speech|><|woitn|>" + results_seq = [ + [{"key": "test_utt", "value": [[0, 1000], [1000, 2000]]}], + [{"text": f"{tag}你好", "timestamp": [[0, 500], [500, 1000]], "words": ["你", "好"]}], + [{"text": f"{tag}世界", "timestamp": [[0, 500], [500, 1000]], "words": ["世", "界"]}], + [{"text": "你好,世界。", "punc_array": [1, 2, 1, 3]}], + ] + + def mock_inference(data, *args, **kwargs): + if len(results_seq) == 1: + self.assertEqual(data, "你好世界") + return results_seq.pop(0) + + am.inference = MagicMock(side_effect=mock_inference) + mock_prep.return_value = (["test_utt"], [np.zeros(32000, dtype=np.float32)]) + mock_load.return_value = np.zeros(32000, dtype=np.float32) + mock_slice.return_value = ([np.zeros(16000, dtype=np.float32)], [16000]) + + results = am.inference_with_vad("dummy_input", sentence_timestamp=True) + + self.assertEqual( + [item["text"] for item in results[0]["sentence_info"]], ["你好,", "世界。"] + ) + self.assertEqual( + [item["timestamp"] for item in results[0]["sentence_info"]], + [ + [[0, 500], [500, 1000]], + [[1000, 1500], [1500, 2000]], + ], + ) + + @patch("funasr.auto.auto_model.distribute_spk") + @patch("funasr.auto.auto_model.postprocess") + @patch("funasr.auto.auto_model.sv_chunk") + @patch("funasr.auto.auto_model.slice_padding_audio_samples") + @patch("funasr.auto.auto_model.load_audio_text_image_video") + @patch("funasr.auto.auto_model.prepare_data_iterator") + def test_speaker_punc_segment_uses_aligned_words( + self, + mock_prep, + mock_load, + mock_slice, + mock_sv_chunk, + mock_postprocess, + mock_distribute_spk, + ): + """Speaker punctuation segmentation must use the same SenseVoice alignment.""" + am = self._make_auto_model( + punc_model=MagicMock(), spk_model=MagicMock(), spk_mode="punc_segment" + ) + am.cb_model = MagicMock(return_value=np.array([0])) + tag = "<|zh|><|NEUTRAL|><|Speech|><|woitn|>" + results_seq = [ + [{"key": "test_utt", "value": [[0, 2000]]}], + [ + { + "text": f"{tag}你好世界", + "timestamp": [[0, 500], [500, 1000], [1000, 1500], [1500, 2000]], + "words": ["你", "好", "世", "界"], + } + ], + [{"spk_embedding": torch.tensor([[1.0, 0.0]])}], + [{"text": "你好,世界。", "punc_array": [1, 2, 1, 3]}], + ] + + def mock_inference(data, *args, **kwargs): + if len(results_seq) == 1: + self.assertEqual(data, "你好世界") + return results_seq.pop(0) + + am.inference = MagicMock(side_effect=mock_inference) + mock_prep.return_value = (["test_utt"], [np.zeros(32000, dtype=np.float32)]) + mock_load.return_value = np.zeros(32000, dtype=np.float32) + mock_slice.return_value = ([np.zeros(32000, dtype=np.float32)], [32000]) + mock_sv_chunk.return_value = [[0.0, 2.0, np.zeros(32000, dtype=np.float32)]] + mock_postprocess.return_value = [{"start": 0.0, "end": 2.0, "spk": 0}] + + results = am.inference_with_vad("dummy_input") + + self.assertEqual( + [item["text"] for item in results[0]["sentence_info"]], ["你好,", "世界。"] + ) + mock_distribute_spk.assert_called_once() + + @patch("funasr.auto.auto_model.slice_padding_audio_samples") + @patch("funasr.auto.auto_model.load_audio_text_image_video") + @patch("funasr.auto.auto_model.prepare_data_iterator") + def test_punctuation_preserves_english_surface_text(self, mock_prep, mock_load, mock_slice): + """Timestamp/BPE units must not add spaces to contractions, URLs, or emails.""" + am = self._make_auto_model(punc_model=MagicMock()) + tag = "<|en|><|NEUTRAL|><|Speech|><|woitn|>" + surface_text = "don't stop https://nature.com email@example.com" + words = [ + "don", + "'", + "t", + "stop", + "https", + ":", + "/", + "/", + "nature", + ".", + "com", + "email", + "@", + "example", + ".", + "com", + ] + mock_prep.return_value = (["test_utt"], [np.zeros(25600, dtype=np.float32)]) + mock_load.return_value = np.zeros(25600, dtype=np.float32) + mock_slice.return_value = ([np.zeros(25600, dtype=np.float32)], [25600]) + + for en_post_proc in (False, True): + with self.subTest(en_post_proc=en_post_proc): + results_seq = [ + [{"key": "test_utt", "value": [[0, 1600]]}], + [ + { + "text": tag + surface_text, + "timestamp": [[i * 100, (i + 1) * 100] for i in range(len(words))], + "words": words, + } + ], + [ + { + "text": " Don ' t stop. Https : / / nature .com. Email @ example .com.", + "punc_array": [1, 3, 1, 3], + } + ], + ] + + def mock_inference(data, *args, **kwargs): + if len(results_seq) == 1: + self.assertEqual(data, surface_text) + return results_seq.pop(0) + + am.inference = MagicMock(side_effect=mock_inference) + results = am.inference_with_vad( + "dummy_input", + sentence_timestamp=True, + en_post_proc=en_post_proc, + ) + + self.assertEqual( + results[0]["text"], + "don't stop. https://nature.com email@example.com.", + ) + self.assertEqual( + [(item["start"], item["end"]) for item in results[0]["sentence_info"]], + [(0, 400), (400, 1600)], + ) + self.assertEqual( + [item["text"] for item in results[0]["sentence_info"]], + ["don't stop.", "https://nature.com email@example.com."], + ) + self.assertEqual(results[0]["sentence_info"][-1]["timestamp"][-1], [1100, 1600]) + + @patch("funasr.auto.auto_model.slice_padding_audio_samples") + @patch("funasr.auto.auto_model.load_audio_text_image_video") + @patch("funasr.auto.auto_model.prepare_data_iterator") + def test_sentence_timestamp_ignores_empty_asr_words(self, mock_prep, mock_load, mock_slice): + """Malformed word metadata must fall back to the existing aligned text.""" + am = self._make_auto_model(punc_model=MagicMock()) + results_seq = [ + [{"key": "test_utt", "value": [[0, 1000]]}], + [ + { + "text": "你 好", + "timestamp": [[0, 500], [500, 1000]], + "words": ["你", ""], + } + ], + [{"text": "你好。", "punc_array": [1, 3]}], + ] + am.inference = MagicMock(side_effect=lambda *args, **kwargs: results_seq.pop(0)) + mock_prep.return_value = (["test_utt"], [np.zeros(1000, dtype=np.float32)]) + mock_load.return_value = np.zeros(16000, dtype=np.float32) + mock_slice.return_value = ([np.zeros(16000, dtype=np.float32)], [16000]) + + results = am.inference_with_vad("dummy_input", sentence_timestamp=True) + + self.assertEqual(results[0]["sentence_info"][0]["text"], "你好。") + + @patch("funasr.auto.auto_model.slice_padding_audio_samples") + @patch("funasr.auto.auto_model.load_audio_text_image_video") + @patch("funasr.auto.auto_model.prepare_data_iterator") + def test_sentence_timestamp_handles_unsized_punc_array(self, mock_prep, mock_load, mock_slice): + """Malformed punctuation metadata must use the legacy no-punctuation fallback.""" + am = self._make_auto_model(punc_model=MagicMock()) + results_seq = [ + [{"key": "test_utt", "value": [[0, 1000]]}], + [ + { + "text": "你 好", + "timestamp": [[0, 500], [500, 1000]], + "words": ["你", "好"], + } + ], + [{"text": "你好。", "punc_array": 3}], + ] + am.inference = MagicMock(side_effect=lambda *args, **kwargs: results_seq.pop(0)) + mock_prep.return_value = (["test_utt"], [np.zeros(1000, dtype=np.float32)]) + mock_load.return_value = np.zeros(16000, dtype=np.float32) + mock_slice.return_value = ([np.zeros(16000, dtype=np.float32)], [16000]) + + results = am.inference_with_vad("dummy_input", sentence_timestamp=True) + + self.assertEqual(results[0]["sentence_info"][0]["text"], ["你", "好"]) + + +class TestCTTransformerPunctuation(unittest.TestCase): + """Keep CT-Transformer text and punctuation-array endings consistent.""" + + @patch("funasr.models.ct_transformer.model.load_audio_text_image_video") + def test_forced_period_uses_sentence_end_id(self, mock_load): + from funasr.models.ct_transformer.model import CTTransformer + + model = CTTransformer.__new__(CTTransformer) + torch.nn.Module.__init__(model) + model.jieba_usr_dict = None + model.punc_list = ["", "_", ",", "。", "?", "、"] + model.sentence_end_id = 3 + + tokenizer = MagicMock() + tokenizer.encode.side_effect = lambda tokens: np.arange(len(tokens), dtype=np.int64) + + for punc_id in (1, 2, 5): + for text in ("hello world", "你好"): + with self.subTest(text=text, punc_id=punc_id): + + def punc_forward(text, text_lengths, **kwargs): + logits = torch.zeros(1, text.shape[1], len(model.punc_list)) + logits[:, :, punc_id] = 1 + return logits, None + + model.punc_forward = punc_forward + mock_load.return_value = [text] + results, _ = model.inference( + data_in=[text], + key=["test_utt"], + tokenizer=tokenizer, + device="cpu", + split_size=20, + ) + + self.assertTrue(results[0]["text"].endswith((".", "。"))) + self.assertEqual(int(results[0]["punc_array"][-1]), model.sentence_end_id) + if __name__ == "__main__": unittest.main()