mirror of
https://github.com/modelscope/FunASR.git
synced 2026-10-12 02:19:18 +00:00
744 lines
31 KiB
Python
744 lines
31 KiB
Python
"""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):
|
||
"""Test that inference_with_vad works when punc_model is None."""
|
||
|
||
def _make_auto_model(self, punc_model=None, spk_model=None, spk_mode=None):
|
||
"""Create a minimal AutoModel instance with mocked dependencies."""
|
||
from funasr.auto.auto_model import AutoModel
|
||
|
||
am = AutoModel.__new__(AutoModel)
|
||
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.spk_kwargs = {}
|
||
am.cb_model = None
|
||
am.spk_mode = spk_mode
|
||
am.vad_kwargs = {}
|
||
am.kwargs = {
|
||
"batch_size_s": 300,
|
||
"batch_size_threshold_s": 60,
|
||
"device": "cpu",
|
||
"disable_pbar": True,
|
||
"frontend": MagicMock(fs=16000),
|
||
"fs": 16000,
|
||
}
|
||
am._reset_runtime_configs = MagicMock()
|
||
return am
|
||
|
||
def _setup_mocks(self, am, mock_slice, mock_load, mock_prep):
|
||
"""Configure standard mocks for a single-segment VAD + ASR flow."""
|
||
# VAD returns one segment [0, 16000ms]
|
||
vad_result = [{"key": "test_utt", "value": [[0, 16000]]}]
|
||
# ASR returns text with timestamps
|
||
asr_result = [{"text": "hello world", "timestamp": [[0, 500], [500, 1000]]}]
|
||
|
||
call_count = [0]
|
||
results_seq = [vad_result, asr_result]
|
||
|
||
def mock_inference(data, input_len=None, model=None, kwargs=None, **cfg):
|
||
idx = call_count[0]
|
||
call_count[0] += 1
|
||
if idx < len(results_seq):
|
||
return results_seq[idx]
|
||
return [{"text": ""}]
|
||
|
||
am.inference = MagicMock(side_effect=mock_inference)
|
||
mock_prep.return_value = (["test_utt"], [np.zeros(16000, dtype=np.float32)])
|
||
mock_load.return_value = np.zeros(16000, dtype=np.float32)
|
||
mock_slice.return_value = ([np.zeros(16000, dtype=np.float32)], [16000])
|
||
|
||
@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_punc_model_none_basic(self, mock_prep, mock_load, mock_slice):
|
||
"""Basic inference with punc_model=None should not raise UnboundLocalError."""
|
||
am = self._make_auto_model(punc_model=None)
|
||
self._setup_mocks(am, mock_slice, mock_load, mock_prep)
|
||
|
||
results = am.inference_with_vad("dummy_input")
|
||
|
||
self.assertEqual(len(results), 1)
|
||
self.assertEqual(results[0]["text"], "hello world")
|
||
self.assertEqual(results[0]["key"], "test_utt")
|
||
|
||
@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_with_punc_model_none(self, mock_prep, mock_load, mock_slice):
|
||
"""sentence_timestamp=True with punc_model=None should not crash."""
|
||
am = self._make_auto_model(punc_model=None)
|
||
self._setup_mocks(am, mock_slice, mock_load, mock_prep)
|
||
|
||
# This path previously caused UnboundLocalError on punc_res
|
||
results = am.inference_with_vad("dummy_input", sentence_timestamp=True)
|
||
|
||
self.assertEqual(len(results), 1)
|
||
# sentence_info should be empty list since punc_res is unavailable
|
||
self.assertEqual(results[0].get("sentence_info"), [])
|
||
|
||
@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_vad_segments_when_asr_has_no_timestamps(
|
||
self, mock_prep, mock_load, mock_slice
|
||
):
|
||
"""SenseVoice should expose its per-VAD text as timestamped sentences."""
|
||
am = self._make_auto_model(punc_model=None)
|
||
results_seq = [
|
||
[{"key": "test_utt", "value": [[120, 800], [1050, 1900]]}],
|
||
[{"text": "first phrase"}],
|
||
[{"text": "second phrase"}],
|
||
]
|
||
am.inference = MagicMock(side_effect=lambda *args, **kwargs: results_seq.pop(0))
|
||
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(
|
||
results[0]["sentence_info"],
|
||
[
|
||
{
|
||
"start": 120,
|
||
"end": 800,
|
||
"text": "first phrase",
|
||
"sentence": "first phrase",
|
||
"timestamp": [],
|
||
},
|
||
{
|
||
"start": 1050,
|
||
"end": 1900,
|
||
"text": "second phrase",
|
||
"sentence": "second phrase",
|
||
"timestamp": [],
|
||
},
|
||
],
|
||
)
|
||
|
||
@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_vad_fallback_requires_punc_model_to_be_absent(
|
||
self, mock_prep, mock_load, mock_slice
|
||
):
|
||
am = self._make_auto_model(punc_model=MagicMock())
|
||
results_seq = [
|
||
[{"key": "test_utt", "value": [[120, 800]]}],
|
||
[{"text": "phrase", "timestamps": []}],
|
||
]
|
||
am.inference = MagicMock(side_effect=lambda *args, **kwargs: results_seq.pop(0))
|
||
mock_prep.return_value = (["test_utt"], [np.zeros(16000, 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"], [])
|
||
|
||
@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_punc_model_with_value_still_works(self, mock_prep, mock_load, mock_slice):
|
||
"""When punc_model is provided, punc_res should still be used normally."""
|
||
punc_mock = MagicMock()
|
||
am = self._make_auto_model(punc_model=punc_mock)
|
||
|
||
vad_result = [{"key": "test_utt", "value": [[0, 16000]]}]
|
||
asr_result = [{"text": "hello world", "timestamp": [[0, 500], [500, 1000]]}]
|
||
punc_result = [{"text": "Hello, world.", "punc_array": [1, 2]}]
|
||
|
||
call_count = [0]
|
||
results_seq = [vad_result, asr_result, punc_result]
|
||
|
||
def mock_inference(data, input_len=None, model=None, kwargs=None, **cfg):
|
||
idx = call_count[0]
|
||
call_count[0] += 1
|
||
return results_seq[idx]
|
||
|
||
am.inference = MagicMock(side_effect=mock_inference)
|
||
mock_prep.return_value = (["test_utt"], [np.zeros(16000, 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")
|
||
|
||
self.assertEqual(len(results), 1)
|
||
# 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.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_falls_back_to_vad_when_punctuation_alignment_fails(
|
||
self, mock_prep, mock_load, mock_slice
|
||
):
|
||
"""A malformed long-audio punctuation array must not collapse all VAD segments."""
|
||
am = self._make_auto_model(punc_model=MagicMock())
|
||
tag = "<|zh|><|NEUTRAL|><|Speech|><|woitn|>"
|
||
results_seq = [
|
||
[{"key": "test_utt", "value": [[100, 1100], [1300, 2300]]}],
|
||
[
|
||
{
|
||
"text": f"{tag}第一句",
|
||
"timestamp": [[0, 300], [300, 600], [600, 900]],
|
||
"words": ["第", "一", "句"],
|
||
}
|
||
],
|
||
[
|
||
{
|
||
"text": f"{tag}第二句",
|
||
"timestamp": [[0, 300], [300, 600], [600, 900]],
|
||
"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(40000, dtype=np.float32)])
|
||
mock_load.return_value = np.zeros(40000, 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"],
|
||
[
|
||
{
|
||
"start": 100,
|
||
"end": 1000,
|
||
"text": "第一句",
|
||
"sentence": "第一句",
|
||
"timestamp": [[100, 400], [400, 700], [700, 1000]],
|
||
},
|
||
{
|
||
"start": 1300,
|
||
"end": 2200,
|
||
"text": "第二句",
|
||
"sentence": "第二句",
|
||
"timestamp": [[1300, 1600], [1600, 1900], [1900, 2200]],
|
||
},
|
||
],
|
||
)
|
||
|
||
@patch(
|
||
"funasr.auto.auto_model._get_punc_tokens",
|
||
return_value=["aon", "storyi", "ca", "说", "差", "距"],
|
||
)
|
||
@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_splits_one_asr_word_across_punctuation_tokens(
|
||
self, mock_prep, mock_load, mock_slice, _mock_get_punc_tokens
|
||
):
|
||
"""Punctuation token boundaries may occur inside one timestamped ASR word."""
|
||
punc_model = MagicMock()
|
||
punc_model.punc_list = None
|
||
am = self._make_auto_model(punc_model=punc_model)
|
||
tag = "<|zh|><|NEUTRAL|><|Speech|><|woitn|>"
|
||
results_seq = [
|
||
[{"key": "test_utt", "value": [[0, 1200]]}],
|
||
[
|
||
{
|
||
"text": f"{tag}aon storyica说差距",
|
||
"timestamp": [
|
||
[0, 100],
|
||
[100, 900],
|
||
[900, 1000],
|
||
[1000, 1100],
|
||
[1100, 1200],
|
||
],
|
||
"words": ["aon", "storyica", "说", "差", "距"],
|
||
}
|
||
],
|
||
[
|
||
{
|
||
"text": "aon storyica.说差距。",
|
||
"punc_array": [1, 1, 3, 1, 1, 3],
|
||
}
|
||
],
|
||
]
|
||
am.inference = MagicMock(side_effect=lambda *args, **kwargs: results_seq.pop(0))
|
||
mock_prep.return_value = (["test_utt"], [np.zeros(19200, dtype=np.float32)])
|
||
mock_load.return_value = np.zeros(19200, dtype=np.float32)
|
||
mock_slice.return_value = ([np.zeros(19200, dtype=np.float32)], [19200])
|
||
|
||
results = am.inference_with_vad("dummy_input", sentence_timestamp=True)
|
||
|
||
self.assertEqual(
|
||
results[0]["sentence_info"],
|
||
[
|
||
{
|
||
"text": "aon storyica.",
|
||
"start": 0,
|
||
"end": 900,
|
||
"timestamp": [[0, 100], [100, 700], [700, 900]],
|
||
},
|
||
{
|
||
"text": "说差距。",
|
||
"start": 900,
|
||
"end": 1200,
|
||
"timestamp": [[900, 1000], [1000, 1100], [1100, 1200]],
|
||
},
|
||
],
|
||
)
|
||
|
||
@patch(
|
||
"funasr.auto.auto_model._get_punc_tokens",
|
||
return_value=["你", "好", "真", "的"],
|
||
)
|
||
@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_sentencepiece_word_boundary_marker(
|
||
self, mock_prep, mock_load, mock_slice, _mock_get_punc_tokens
|
||
):
|
||
punc_model = MagicMock()
|
||
punc_model.punc_list = None
|
||
am = self._make_auto_model(punc_model=punc_model)
|
||
tag = "<|zh|><|NEUTRAL|><|Speech|><|woitn|>"
|
||
results_seq = [
|
||
[{"key": "test_utt", "value": [[0, 2000]]}],
|
||
[
|
||
{
|
||
"text": f"{tag}你好真的",
|
||
"timestamp": [[0, 500], [500, 1000], [1000, 2000]],
|
||
"words": ["你", "好", "▁真的"],
|
||
}
|
||
],
|
||
[{"text": "你好，真的。", "punc_array": [1, 2, 1, 3]}],
|
||
]
|
||
am.inference = MagicMock(side_effect=lambda *args, **kwargs: results_seq.pop(0))
|
||
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])
|
||
|
||
results = am.inference_with_vad("dummy_input", sentence_timestamp=True)
|
||
|
||
self.assertEqual(
|
||
results[0]["sentence_info"],
|
||
[
|
||
{
|
||
"text": "你好，",
|
||
"start": 0,
|
||
"end": 1000,
|
||
"timestamp": [[0, 500], [500, 1000]],
|
||
},
|
||
{
|
||
"text": "真的。",
|
||
"start": 1000,
|
||
"end": 2000,
|
||
"timestamp": [[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.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_no_timestamp_speaker_fallback_keeps_sentence_and_adds_text(
|
||
self,
|
||
mock_prep,
|
||
mock_load,
|
||
mock_slice,
|
||
mock_sv_chunk,
|
||
mock_postprocess,
|
||
_mock_distribute_spk,
|
||
):
|
||
"""Speaker VAD fallback exposes documented text without dropping sentence."""
|
||
am = self._make_auto_model(
|
||
punc_model=None, spk_model=MagicMock(), spk_mode="vad_segment"
|
||
)
|
||
am.cb_model = MagicMock(return_value=np.array([0]))
|
||
inference_calls = []
|
||
results_seq = [
|
||
[{"key": "test_utt", "value": [[0, 1000]]}],
|
||
[{"text": "hello"}],
|
||
[{"spk_embedding": torch.tensor([[1.0, 0.0]])}],
|
||
]
|
||
|
||
def mock_inference(data, *args, **kwargs):
|
||
inference_calls.append(kwargs)
|
||
return results_seq.pop(0)
|
||
|
||
am.inference = MagicMock(side_effect=mock_inference)
|
||
mock_prep.return_value = (
|
||
["test_utt"],
|
||
[np.zeros(16000, dtype=np.float32)],
|
||
)
|
||
mock_load.return_value = np.zeros(16000, dtype=np.float32)
|
||
mock_slice.return_value = ([np.zeros(16000, dtype=np.float32)], [16000])
|
||
mock_sv_chunk.return_value = [
|
||
[0.0, 1.0, np.zeros(16000, dtype=np.float32)]
|
||
]
|
||
mock_postprocess.return_value = [{"start": 0.0, "end": 1.0, "spk": 0}]
|
||
|
||
results = am.inference_with_vad("dummy_input")
|
||
|
||
self.assertIs(inference_calls[1]["output_timestamp"], True)
|
||
self.assertIs(inference_calls[1]["return_time_stamps"], True)
|
||
sentence = results[0]["sentence_info"][0]
|
||
self.assertEqual(sentence["text"], "hello")
|
||
self.assertEqual(sentence["sentence"], "hello")
|
||
|
||
@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_keeps_unpunctuated_tail_without_asr_words(
|
||
self, mock_prep, mock_load, mock_slice
|
||
):
|
||
am = self._make_auto_model(punc_model=MagicMock())
|
||
results_seq = [
|
||
[{"key": "test_utt", "value": [[0, 300]]}],
|
||
[{"text": "hello world", "timestamp": [[0, 100], [100, 300]]}],
|
||
[{"text": "hello world", "punc_array": [1, 1]}],
|
||
]
|
||
am.inference = MagicMock(side_effect=lambda *args, **kwargs: results_seq.pop(0))
|
||
mock_prep.return_value = (["test_utt"], [np.zeros(4800, dtype=np.float32)])
|
||
mock_load.return_value = np.zeros(4800, dtype=np.float32)
|
||
mock_slice.return_value = ([np.zeros(4800, dtype=np.float32)], [4800])
|
||
|
||
result = am.inference_with_vad(
|
||
"dummy_input", sentence_timestamp=True, en_post_proc=True
|
||
)
|
||
|
||
self.assertEqual(
|
||
[item["text"].strip() for item in result[0]["sentence_info"]],
|
||
["hello world"],
|
||
)
|
||
self.assertEqual(result[0]["sentence_info"][0]["timestamp"], [[0, 100], [100, 300]])
|
||
|
||
@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()
|