mirror of
https://github.com/modelscope/FunASR.git
synced 2026-10-12 02:19:18 +00:00
fix: align SenseVoice sentence timestamps (#3407)
This commit is contained in:
zhifu gao
+255
-16
@@ -8,6 +8,7 @@ import time
|
||||
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 @@ class AutoModel:
|
||||
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)
|
||||
punc_res = self.inference(
|
||||
result["text"], model=self.punc_model, kwargs=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(
|
||||
punc_input_text,
|
||||
model=self.punc_model,
|
||||
kwargs=self.punc_kwargs,
|
||||
**cfg,
|
||||
)
|
||||
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 @@ class AutoModel:
|
||||
"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 @@ class AutoModel:
|
||||
"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
|
||||
|
||||
@@ -416,11 +416,15 @@ class CTTransformer(torch.nn.Module):
|
||||
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 @@ class CTTransformer(torch.nn.Module):
|
||||
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 @@ class CTTransformer(torch.nn.Module):
|
||||
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
|
||||
|
||||
@@ -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 @@ class TestPuncModelNone(unittest.TestCase):
|
||||
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 = ["<unk>", "_", "，", "。", "？", "、"]
|
||||
am.punc_kwargs = {}
|
||||
am.spk_model = spk_model
|
||||
am.cb_model = None
|
||||
@@ -114,6 +118,321 @@ class TestPuncModelNone(unittest.TestCase):
|
||||
# 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 protected]"
|
||||
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 protected].",
|
||||
)
|
||||
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 protected]."],
|
||||
)
|
||||
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 = ["<unk>", "_", "，", "。", "？", "、"]
|
||||
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()
|
||||
|
||||
Reference in New Issue
Block a user