fix: align SenseVoice sentence timestamps

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

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