vee
107 lines
4.1 KiB
Python
107 lines
4.1 KiB
Python
"""Download official SenseVoiceSmall and Silero VAD sherpa-onnx models."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import shutil
|
|
import tarfile
|
|
import tempfile
|
|
import urllib.request
|
|
from pathlib import Path
|
|
|
|
MODEL_URL = (
|
|
"https://github.com/k2-fsa/sherpa-onnx/releases/download/asr-models/"
|
|
"sherpa-onnx-sense-voice-zh-en-ja-ko-yue-2024-07-17.tar.bz2"
|
|
)
|
|
MODEL_SHA256 = "f6b2a72ebcb1ac7a764d4cfccd886e6bcb2a95c4657c2199d0ba95ed4b9ea71a"
|
|
REPOSITORY_ROOT = Path(__file__).resolve().parents[1]
|
|
DEFAULT_TARGET = REPOSITORY_ROOT / "artifacts" / "models" / "sensevoice-small"
|
|
VAD_URL = (
|
|
"https://github.com/k2-fsa/sherpa-onnx/releases/download/asr-models/"
|
|
"silero_vad.onnx"
|
|
)
|
|
VAD_SHA256 = "9e2449e1087496d8d4caba907f23e0bd3f78d91fa552479bb9c23ac09cbb1fd6"
|
|
DEFAULT_VAD_TARGET = (
|
|
REPOSITORY_ROOT / "artifacts" / "models" / "silero-vad" / "silero_vad.onnx"
|
|
)
|
|
|
|
|
|
def _download(url: str, destination: Path) -> None:
|
|
with urllib.request.urlopen(url, timeout=60) as response, destination.open("wb") as output:
|
|
shutil.copyfileobj(response, output, length=1024 * 1024)
|
|
|
|
|
|
def _sha256(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as source:
|
|
for chunk in iter(lambda: source.read(1024 * 1024), b""):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def install(target: Path) -> None:
|
|
target = target.resolve()
|
|
if (target / "model.onnx").is_file() and (target / "tokens.txt").is_file():
|
|
print(f"SenseVoiceSmall is already installed at {target}")
|
|
return
|
|
target.mkdir(parents=True, exist_ok=True)
|
|
with tempfile.TemporaryDirectory(prefix="sensevoice-install-") as temporary_directory:
|
|
temporary_root = Path(temporary_directory)
|
|
archive = temporary_root / "sensevoice-small.tar.bz2"
|
|
print(f"Downloading the official SenseVoiceSmall model to {archive} ...", flush=True)
|
|
_download(MODEL_URL, archive)
|
|
actual_hash = _sha256(archive)
|
|
if actual_hash != MODEL_SHA256:
|
|
raise RuntimeError(
|
|
f"SenseVoiceSmall archive checksum mismatch: {actual_hash} != {MODEL_SHA256}"
|
|
)
|
|
extracted = temporary_root / "extracted"
|
|
extracted.mkdir()
|
|
with tarfile.open(archive, "r:bz2") as package:
|
|
package.extractall(extracted, filter="data")
|
|
model_file = next(extracted.rglob("model.onnx"), None)
|
|
if model_file is None or not (model_file.parent / "tokens.txt").is_file():
|
|
raise RuntimeError("SenseVoiceSmall archive does not contain the expected model files")
|
|
for item in model_file.parent.iterdir():
|
|
destination = target / item.name
|
|
if item.is_dir():
|
|
shutil.copytree(item, destination, dirs_exist_ok=True)
|
|
else:
|
|
shutil.copy2(item, destination)
|
|
print(f"SenseVoiceSmall installed at {target}")
|
|
|
|
|
|
def install_vad(target: Path) -> None:
|
|
target = target.resolve()
|
|
if target.is_file() and _sha256(target) == VAD_SHA256:
|
|
print(f"Silero VAD is already installed at {target}")
|
|
return
|
|
target.parent.mkdir(parents=True, exist_ok=True)
|
|
with tempfile.TemporaryDirectory(prefix="silero-vad-install-") as temporary_directory:
|
|
downloaded = Path(temporary_directory) / "silero_vad.onnx"
|
|
print(f"Downloading the official Silero VAD model to {downloaded} ...", flush=True)
|
|
_download(VAD_URL, downloaded)
|
|
actual_hash = _sha256(downloaded)
|
|
if actual_hash != VAD_SHA256:
|
|
raise RuntimeError(
|
|
f"Silero VAD checksum mismatch: {actual_hash} != {VAD_SHA256}"
|
|
)
|
|
shutil.copy2(downloaded, target)
|
|
print(f"Silero VAD installed at {target}")
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--target", type=Path, default=DEFAULT_TARGET)
|
|
parser.add_argument("--vad-target", type=Path, default=DEFAULT_VAD_TARGET)
|
|
parser.add_argument("--skip-vad", action="store_true")
|
|
args = parser.parse_args()
|
|
install(args.target)
|
|
if not args.skip_vad:
|
|
install_vad(args.vad_target)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|