mkbk
245 lines
9.4 KiB
Python
245 lines
9.4 KiB
Python
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
import json
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from .models import ErrorCode, ProviderError
|
|
|
|
DEFAULT_VAD_URL = "https://github.com/k2-fsa/sherpa-onnx/releases/download/asr-models/silero_vad.onnx"
|
|
DEFAULT_STT_URL = (
|
|
"https://github.com/k2-fsa/sherpa-onnx/releases/download/asr-models/"
|
|
"sherpa-onnx-streaming-zipformer-ctc-zh-int8-2025-06-30.tar.bz2"
|
|
)
|
|
DEFAULT_STT_DIR = "sherpa-onnx-streaming-zipformer-ctc-zh-int8-2025-06-30"
|
|
DEFAULT_KWS_URL = (
|
|
"https://github.com/k2-fsa/sherpa-onnx/releases/download/kws-models/"
|
|
"sherpa-onnx-kws-zipformer-wenetspeech-3.3M-2024-01-01-mobile.tar.bz2"
|
|
)
|
|
DEFAULT_KWS_DIR = "sherpa-onnx-kws-zipformer-wenetspeech-3.3M-2024-01-01-mobile"
|
|
DEFAULT_KWS_KEYWORDS = "x iǎo j ié x iǎo j ié @小杰小杰\n"
|
|
DEFAULT_DENOISER_URL = (
|
|
"https://github.com/k2-fsa/sherpa-onnx/releases/download/"
|
|
"speech-enhancement-models/gtcrn_simple.onnx"
|
|
)
|
|
DEFAULT_DENOISER_PATH = "denoise/gtcrn_simple.onnx"
|
|
|
|
REQUIRED_MODEL_FILES = (
|
|
f"wake/{DEFAULT_KWS_DIR}/tokens.txt",
|
|
f"wake/{DEFAULT_KWS_DIR}/encoder-epoch-12-avg-2-chunk-16-left-64.int8.onnx",
|
|
f"wake/{DEFAULT_KWS_DIR}/decoder-epoch-12-avg-2-chunk-16-left-64.onnx",
|
|
f"wake/{DEFAULT_KWS_DIR}/joiner-epoch-12-avg-2-chunk-16-left-64.int8.onnx",
|
|
"wake/keywords.txt",
|
|
"vad/silero_vad.onnx",
|
|
f"stt/{DEFAULT_STT_DIR}/tokens.txt",
|
|
f"stt/{DEFAULT_STT_DIR}/model.int8.onnx",
|
|
DEFAULT_DENOISER_PATH,
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class SpeechModelStatus:
|
|
models_dir: Path
|
|
manifest_path: Path
|
|
manifest_present: bool
|
|
required_files: tuple[str, ...]
|
|
missing_files: tuple[str, ...]
|
|
sherpa_onnx_available: bool
|
|
|
|
@property
|
|
def ok(self) -> bool:
|
|
return not self.missing_files and self.sherpa_onnx_available
|
|
|
|
def to_json(self) -> dict[str, Any]:
|
|
return {
|
|
"ok": self.ok,
|
|
"models_dir": str(self.models_dir),
|
|
"manifest_path": str(self.manifest_path),
|
|
"manifest_present": self.manifest_present,
|
|
"required_files": list(self.required_files),
|
|
"missing_files": list(self.missing_files),
|
|
"sherpa_onnx_available": self.sherpa_onnx_available,
|
|
}
|
|
|
|
|
|
def default_manifest() -> dict[str, Any]:
|
|
return {
|
|
"schema": "owner_voice_pet.speech_models.v1",
|
|
"sample_rate": 16000,
|
|
"sources": {
|
|
"wake": DEFAULT_KWS_URL,
|
|
"vad": DEFAULT_VAD_URL,
|
|
"stt": DEFAULT_STT_URL,
|
|
"denoiser": DEFAULT_DENOISER_URL,
|
|
},
|
|
"providers": {
|
|
"wake": {
|
|
"type": "sherpa-onnx-keyword-spotter",
|
|
"model_dir": f"wake/{DEFAULT_KWS_DIR}",
|
|
"tokens": f"wake/{DEFAULT_KWS_DIR}/tokens.txt",
|
|
"encoder": f"wake/{DEFAULT_KWS_DIR}/encoder-epoch-12-avg-2-chunk-16-left-64.int8.onnx",
|
|
"decoder": f"wake/{DEFAULT_KWS_DIR}/decoder-epoch-12-avg-2-chunk-16-left-64.onnx",
|
|
"joiner": f"wake/{DEFAULT_KWS_DIR}/joiner-epoch-12-avg-2-chunk-16-left-64.int8.onnx",
|
|
"keywords": "wake/keywords.txt",
|
|
},
|
|
"vad": {
|
|
"type": "silero-vad",
|
|
"path": "vad/silero_vad.onnx",
|
|
},
|
|
"stt": {
|
|
"type": "sherpa-onnx-streaming-zipformer2-ctc",
|
|
"model_dir": f"stt/{DEFAULT_STT_DIR}",
|
|
"tokens": f"stt/{DEFAULT_STT_DIR}/tokens.txt",
|
|
"model": f"stt/{DEFAULT_STT_DIR}/model.int8.onnx",
|
|
},
|
|
"denoiser": {
|
|
"type": "sherpa-onnx-gtcrn",
|
|
"path": DEFAULT_DENOISER_PATH,
|
|
},
|
|
},
|
|
"required_files": list(REQUIRED_MODEL_FILES),
|
|
}
|
|
|
|
|
|
def write_default_manifest(models_dir: str | Path) -> Path:
|
|
root = Path(models_dir)
|
|
root.mkdir(parents=True, exist_ok=True)
|
|
manifest_path = root / "manifest.json"
|
|
manifest_path.write_text(
|
|
json.dumps(default_manifest(), ensure_ascii=False, indent=2) + "\n",
|
|
encoding="utf-8",
|
|
)
|
|
return manifest_path
|
|
|
|
|
|
def load_manifest(models_dir: str | Path) -> dict[str, Any]:
|
|
manifest_path = Path(models_dir) / "manifest.json"
|
|
if not manifest_path.exists():
|
|
return default_manifest()
|
|
return json.loads(manifest_path.read_text(encoding="utf-8"))
|
|
|
|
|
|
def required_model_files(models_dir: str | Path) -> tuple[str, ...]:
|
|
manifest = load_manifest(models_dir)
|
|
files = manifest.get("required_files", list(REQUIRED_MODEL_FILES))
|
|
return tuple(str(item) for item in files)
|
|
|
|
|
|
def vad_model_path(models_dir: str | Path) -> Path:
|
|
root = Path(models_dir)
|
|
manifest = load_manifest(root)
|
|
path = manifest.get("providers", {}).get("vad", {}).get("path", "vad/silero_vad.onnx")
|
|
return root / str(path)
|
|
|
|
|
|
def denoiser_model_path(models_dir: str | Path) -> Path:
|
|
root = Path(models_dir)
|
|
manifest = load_manifest(root)
|
|
path = manifest.get("providers", {}).get("denoiser", {}).get("path", DEFAULT_DENOISER_PATH)
|
|
return root / str(path)
|
|
|
|
|
|
def wake_model_paths(models_dir: str | Path) -> dict[str, Path]:
|
|
root = Path(models_dir)
|
|
manifest = load_manifest(root)
|
|
wake = manifest.get("providers", {}).get("wake", {})
|
|
return {
|
|
"model_dir": root / str(wake.get("model_dir", f"wake/{DEFAULT_KWS_DIR}")),
|
|
"tokens": root / str(wake.get("tokens", f"wake/{DEFAULT_KWS_DIR}/tokens.txt")),
|
|
"encoder": root
|
|
/ str(wake.get("encoder", f"wake/{DEFAULT_KWS_DIR}/encoder-epoch-12-avg-2-chunk-16-left-64.int8.onnx")),
|
|
"decoder": root
|
|
/ str(wake.get("decoder", f"wake/{DEFAULT_KWS_DIR}/decoder-epoch-12-avg-2-chunk-16-left-64.onnx")),
|
|
"joiner": root
|
|
/ str(wake.get("joiner", f"wake/{DEFAULT_KWS_DIR}/joiner-epoch-12-avg-2-chunk-16-left-64.int8.onnx")),
|
|
"keywords": root / str(wake.get("keywords", "wake/keywords.txt")),
|
|
}
|
|
|
|
|
|
def stt_model_paths(model_path: str | Path) -> dict[str, Any]:
|
|
root = Path(model_path)
|
|
if (root / "manifest.json").exists() or (root / "stt").exists():
|
|
manifest = load_manifest(root)
|
|
stt = manifest.get("providers", {}).get("stt", {})
|
|
stt_type = str(stt.get("type", "sherpa-onnx-streaming-transducer"))
|
|
if stt_type == "sherpa-onnx-streaming-zipformer2-ctc":
|
|
return {
|
|
"type": stt_type,
|
|
"model_dir": root / str(stt.get("model_dir", f"stt/{DEFAULT_STT_DIR}")),
|
|
"tokens": root / str(stt.get("tokens", f"stt/{DEFAULT_STT_DIR}/tokens.txt")),
|
|
"model": root / str(stt.get("model", f"stt/{DEFAULT_STT_DIR}/model.int8.onnx")),
|
|
}
|
|
return {
|
|
"type": stt_type,
|
|
"model_dir": root / str(stt.get("model_dir", f"stt/{DEFAULT_STT_DIR}")),
|
|
"tokens": root / str(stt.get("tokens", f"stt/{DEFAULT_STT_DIR}/tokens.txt")),
|
|
"encoder": root / str(stt.get("encoder", f"stt/{DEFAULT_STT_DIR}/encoder-epoch-99-avg-1.int8.onnx")),
|
|
"decoder": root / str(stt.get("decoder", f"stt/{DEFAULT_STT_DIR}/decoder-epoch-99-avg-1.onnx")),
|
|
"joiner": root / str(stt.get("joiner", f"stt/{DEFAULT_STT_DIR}/joiner-epoch-99-avg-1.int8.onnx")),
|
|
}
|
|
return {
|
|
"type": "sherpa-onnx-streaming-transducer",
|
|
"model_dir": root,
|
|
"tokens": root / "tokens.txt",
|
|
"encoder": root / "encoder-epoch-99-avg-1.int8.onnx",
|
|
"decoder": root / "decoder-epoch-99-avg-1.onnx",
|
|
"joiner": root / "joiner-epoch-99-avg-1.int8.onnx",
|
|
}
|
|
|
|
|
|
def check_speech_models(models_dir: str | Path, require_sherpa: bool = True) -> SpeechModelStatus:
|
|
root = Path(models_dir)
|
|
manifest_path = root / "manifest.json"
|
|
required = required_model_files(root)
|
|
missing = tuple(relative for relative in required if not (root / relative).exists())
|
|
sherpa_available = True
|
|
if require_sherpa:
|
|
sherpa_available = importlib.util.find_spec("sherpa_onnx") is not None
|
|
return SpeechModelStatus(
|
|
models_dir=root,
|
|
manifest_path=manifest_path,
|
|
manifest_present=manifest_path.exists(),
|
|
required_files=required,
|
|
missing_files=missing,
|
|
sherpa_onnx_available=sherpa_available,
|
|
)
|
|
|
|
|
|
def model_status_errors(status: SpeechModelStatus) -> list[ProviderError]:
|
|
errors: list[ProviderError] = []
|
|
if not status.sherpa_onnx_available:
|
|
errors.append(
|
|
ProviderError(
|
|
ErrorCode.STT_TRANSCRIBE_FAILED,
|
|
"sherpa_onnx is not installed in the active Python environment",
|
|
False,
|
|
"sherpa-onnx",
|
|
"model-check",
|
|
)
|
|
)
|
|
denoise_missing = tuple(item for item in status.missing_files if item.startswith("denoise/"))
|
|
speech_missing = tuple(item for item in status.missing_files if not item.startswith("denoise/"))
|
|
if speech_missing:
|
|
errors.append(
|
|
ProviderError(
|
|
ErrorCode.STT_MODEL_MISSING,
|
|
"missing speech model files: " + ", ".join(speech_missing),
|
|
False,
|
|
"speech-models",
|
|
"model-check",
|
|
)
|
|
)
|
|
if denoise_missing:
|
|
errors.append(
|
|
ProviderError(
|
|
ErrorCode.NOISE_FILTER_MODEL_MISSING,
|
|
"missing denoiser model files: " + ", ".join(denoise_missing),
|
|
False,
|
|
"speech-models",
|
|
"model-check",
|
|
)
|
|
)
|
|
return errors
|