diff --git a/README.en.md b/README.en.md index 1ffa5f5..f7117d3 100644 --- a/README.en.md +++ b/README.en.md @@ -52,6 +52,19 @@ uv sync --extra whisper --extra web Available extras: `whisper`, `sensevoice`, `volcengine`, `web`, `server`. +#### Prepare the SenseVoice ONNX model + +When using SenseVoice, download [`iic/SenseVoiceSmall-onnx`](https://modelscope.cn/models/iic/SenseVoiceSmall-onnx) to a local directory. That ONNX repository does not include the required `chn_jpn_yue_eng_ko_spectok.bpe.model`; download the file with that name from [`iic/SenseVoiceSmall`](https://modelscope.cn/models/iic/SenseVoiceSmall) and place it in the same directory. The final directory must contain at least: + +```text +config.yaml +am.mvn +chn_jpn_yue_eng_ko_spectok.bpe.model +model_quant.onnx # official quantized model; model.onnx is also supported +``` + +Point the SenseVoice model-directory prompt at this directory. bili2text automatically selects quantized or non-quantized loading from `model_quant.onnx` or `model.onnx`, and reports incomplete directories before loading the runtime. + ### Set Up A setup wizard runs automatically the first time, or you can launch it manually: diff --git a/README.md b/README.md index 5b17613..4c83cd1 100644 --- a/README.md +++ b/README.md @@ -59,6 +59,19 @@ uv sync --extra whisper --extra web 可选的 extras：`whisper`、`sensevoice`、`volcengine`、`web`、`server`。可以暂时不用安装，详看下方的初始化文档。 +#### 准备 SenseVoice ONNX 模型 + +选择 SenseVoice 时，先下载 [`iic/SenseVoiceSmall-onnx`](https://modelscope.cn/models/iic/SenseVoiceSmall-onnx) 到本地目录。该 ONNX 仓库缺少运行时必需的 `chn_jpn_yue_eng_ko_spectok.bpe.model`，还需要从 [`iic/SenseVoiceSmall`](https://modelscope.cn/models/iic/SenseVoiceSmall) 下载这个同名文件，并放进同一目录。最终目录至少应包含： + +```text +config.yaml +am.mvn +chn_jpn_yue_eng_ko_spectok.bpe.model +model_quant.onnx # 官方量化模型；也支持非量化 model.onnx +``` + +配置向导中的 SenseVoice 模型目录应指向这个目录。bili2text 会根据 `model_quant.onnx` 或 `model.onnx` 自动选择量化模式，并在文件不完整时列出缺失项。 + ### 初始化配置 第一次运行时会自动弹出配置向导，也可以手动运行： diff --git a/src/b2t/transcribers/sensevoice_local.py b/src/b2t/transcribers/sensevoice_local.py index 28650f6..ac9c8e6 100644 --- a/src/b2t/transcribers/sensevoice_local.py +++ b/src/b2t/transcribers/sensevoice_local.py @@ -6,11 +6,16 @@ from b2t.i18n import dependency_sync_guidance from b2t.transcribers.base import Transcriber +SENSEVOICE_TOKENIZER = "chn_jpn_yue_eng_ko_spectok.bpe.model" +SENSEVOICE_SUPPORT_FILES = ("config.yaml", "am.mvn", SENSEVOICE_TOKENIZER) + class SenseVoiceSmallTranscriber(Transcriber): name = "sensevoice" - def __init__(self, *, model_dir: Path, language: str = "auto", use_itn: bool = True) -> None: + def __init__( + self, *, model_dir: Path, language: str = "auto", use_itn: bool = True + ) -> None: self.model_dir = model_dir self.language = language self.use_itn = use_itn @@ -28,7 +33,9 @@ def transcribe( progress.running("transcribing", message="transcribing", indeterminate=True) try: - from funasr_onnx.utils.postprocess_utils import rich_transcription_postprocess + from funasr_onnx.utils.postprocess_utils import ( + rich_transcription_postprocess, + ) except ImportError as exc: raise RuntimeError( "SenseVoice support is not installed. " @@ -58,7 +65,11 @@ def _ensure_model(self) -> Any: return self._model if not self.model_dir.exists(): - raise RuntimeError(f"SenseVoice model directory does not exist: {self.model_dir}") + raise RuntimeError( + f"SenseVoice model directory does not exist: {self.model_dir}" + ) + + quantize = _validate_model_dir(self.model_dir) try: from funasr_onnx import SenseVoiceSmall @@ -68,7 +79,7 @@ def _ensure_model(self) -> Any: f"{dependency_sync_guidance('en-US')}" ) from exc - self._model = SenseVoiceSmall(str(self.model_dir)) + self._model = SenseVoiceSmall(str(self.model_dir), quantize=quantize) return self._model @@ -76,3 +87,28 @@ def _extract_text(item: object) -> str: if isinstance(item, dict): return str(item.get("text", "")) return str(item) + + +def _validate_model_dir(model_dir: Path) -> bool: + missing = [ + filename + for filename in SENSEVOICE_SUPPORT_FILES + if not (model_dir / filename).is_file() + ] + if missing: + missing_text = ", ".join(missing) + raise RuntimeError( + f"SenseVoice model directory is incomplete ({model_dir}). Missing: {missing_text}. " + "The ONNX repository does not include the SentencePiece tokenizer; copy " + f"{SENSEVOICE_TOKENIZER} from iic/SenseVoiceSmall into this directory." + ) + + if (model_dir / "model_quant.onnx").is_file(): + return True + if (model_dir / "model.onnx").is_file(): + return False + + raise RuntimeError( + f"SenseVoice model directory is incomplete ({model_dir}). " + "Missing model_quant.onnx or model.onnx." + ) diff --git a/tests/test_sensevoice_local.py b/tests/test_sensevoice_local.py new file mode 100644 index 0000000..c1a612a --- /dev/null +++ b/tests/test_sensevoice_local.py @@ -0,0 +1,59 @@ +from __future__ import annotations + +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from b2t.transcribers.sensevoice_local import SenseVoiceSmallTranscriber + +REQUIRED_SUPPORT_FILES = ( + "config.yaml", + "am.mvn", + "chn_jpn_yue_eng_ko_spectok.bpe.model", +) + + +def _write_support_files(model_dir: Path) -> None: + for filename in REQUIRED_SUPPORT_FILES: + (model_dir / filename).write_text("test", encoding="utf-8") + + +@pytest.mark.parametrize( + ("model_filename", "expected_quantize"), + (("model_quant.onnx", True), ("model.onnx", False)), +) +def test_sensevoice_selects_quantization_from_available_model( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + model_filename: str, + expected_quantize: bool, +) -> None: + _write_support_files(tmp_path) + (tmp_path / model_filename).write_bytes(b"onnx") + calls: list[tuple[str, bool]] = [] + + def fake_model(model_dir: str, *, quantize: bool): + calls.append((model_dir, quantize)) + return object() + + monkeypatch.setitem( + sys.modules, "funasr_onnx", SimpleNamespace(SenseVoiceSmall=fake_model) + ) + + transcriber = SenseVoiceSmallTranscriber(model_dir=tmp_path) + transcriber._ensure_model() + + assert calls == [(str(tmp_path), expected_quantize)] + + +def test_sensevoice_reports_missing_tokenizer_before_model_load(tmp_path: Path) -> None: + (tmp_path / "model_quant.onnx").write_bytes(b"onnx") + (tmp_path / "config.yaml").write_text("test", encoding="utf-8") + (tmp_path / "am.mvn").write_text("test", encoding="utf-8") + + transcriber = SenseVoiceSmallTranscriber(model_dir=tmp_path) + + with pytest.raises(RuntimeError, match="chn_jpn_yue_eng_ko_spectok.bpe.model"): + transcriber._ensure_model()