mirror of
https://github.com/TaterTotterson/microWakeWord-Trainer-Nvidia-Docker.git
synced 2026-08-12 07:55:33 -06:00
Release NVIDIA WakeWord Trainer v19
This commit is contained in:
@@ -1,3 +1,3 @@
|
|||||||
- Added Parakeet ONNX as a second local Auto Training STT engine alongside Faster Whisper.
|
- Fixed first-run Parakeet ONNX setup failing when its empty model directory was mistaken for a complete offline model.
|
||||||
- Replaced manual model, device, and compute fields with a simple engine selector and managed language-aware models.
|
- Parakeet now downloads or resumes the required INT8 snapshot before loading through ONNX ASR.
|
||||||
- Added CUDA-enabled ONNX Runtime with CPU fallback, runtime reporting, and model-cache cleanup when switching engines.
|
- Complete local snapshots are reused without Hub access, preserving offline startup after the initial download.
|
||||||
|
|||||||
@@ -43,12 +43,14 @@ class AutoTrainTests(unittest.TestCase):
|
|||||||
trainer.PERSONAL_DIR,
|
trainer.PERSONAL_DIR,
|
||||||
trainer.AUTO_TRAIN_CONFIG_FILE,
|
trainer.AUTO_TRAIN_CONFIG_FILE,
|
||||||
trainer.AUTO_TRAIN_STATE_FILE,
|
trainer.AUTO_TRAIN_STATE_FILE,
|
||||||
|
trainer.AUTO_TRAIN_MODEL_DIR,
|
||||||
)
|
)
|
||||||
trainer.CAPTURED_DIR = root / "captured_audio"
|
trainer.CAPTURED_DIR = root / "captured_audio"
|
||||||
trainer.NEGATIVE_DIR = root / "negative_samples"
|
trainer.NEGATIVE_DIR = root / "negative_samples"
|
||||||
trainer.PERSONAL_DIR = root / "personal_samples"
|
trainer.PERSONAL_DIR = root / "personal_samples"
|
||||||
trainer.AUTO_TRAIN_CONFIG_FILE = root / "auto_train_config.json"
|
trainer.AUTO_TRAIN_CONFIG_FILE = root / "auto_train_config.json"
|
||||||
trainer.AUTO_TRAIN_STATE_FILE = root / "auto_train_state.json"
|
trainer.AUTO_TRAIN_STATE_FILE = root / "auto_train_state.json"
|
||||||
|
trainer.AUTO_TRAIN_MODEL_DIR = root / "auto_train_models"
|
||||||
for directory in (trainer.CAPTURED_DIR, trainer.NEGATIVE_DIR, trainer.PERSONAL_DIR):
|
for directory in (trainer.CAPTURED_DIR, trainer.NEGATIVE_DIR, trainer.PERSONAL_DIR):
|
||||||
directory.mkdir(parents=True)
|
directory.mkdir(parents=True)
|
||||||
|
|
||||||
@@ -75,6 +77,7 @@ class AutoTrainTests(unittest.TestCase):
|
|||||||
trainer.PERSONAL_DIR,
|
trainer.PERSONAL_DIR,
|
||||||
trainer.AUTO_TRAIN_CONFIG_FILE,
|
trainer.AUTO_TRAIN_CONFIG_FILE,
|
||||||
trainer.AUTO_TRAIN_STATE_FILE,
|
trainer.AUTO_TRAIN_STATE_FILE,
|
||||||
|
trainer.AUTO_TRAIN_MODEL_DIR,
|
||||||
) = self.original_paths
|
) = self.original_paths
|
||||||
trainer.AUTO_TRAIN_CONFIG.clear()
|
trainer.AUTO_TRAIN_CONFIG.clear()
|
||||||
trainer.AUTO_TRAIN_CONFIG.update(self.original_config)
|
trainer.AUTO_TRAIN_CONFIG.update(self.original_config)
|
||||||
@@ -160,8 +163,17 @@ class AutoTrainTests(unittest.TestCase):
|
|||||||
def test_parakeet_loader_prefers_cuda_then_cpu(self):
|
def test_parakeet_loader_prefers_cuda_then_cpu(self):
|
||||||
fake_model = object()
|
fake_model = object()
|
||||||
fake_onnx_asr = SimpleNamespace(load_model=Mock(return_value=fake_model))
|
fake_onnx_asr = SimpleNamespace(load_model=Mock(return_value=fake_model))
|
||||||
|
fake_huggingface_hub = SimpleNamespace(
|
||||||
|
snapshot_download=Mock(return_value=str(trainer.AUTO_TRAIN_MODEL_DIR))
|
||||||
|
)
|
||||||
with (
|
with (
|
||||||
patch.dict(sys.modules, {"onnx_asr": fake_onnx_asr}),
|
patch.dict(
|
||||||
|
sys.modules,
|
||||||
|
{
|
||||||
|
"onnx_asr": fake_onnx_asr,
|
||||||
|
"huggingface_hub": fake_huggingface_hub,
|
||||||
|
},
|
||||||
|
),
|
||||||
patch.object(
|
patch.object(
|
||||||
trainer,
|
trainer,
|
||||||
"_parakeet_onnx_providers",
|
"_parakeet_onnx_providers",
|
||||||
@@ -173,6 +185,57 @@ class AutoTrainTests(unittest.TestCase):
|
|||||||
loaded = trainer._load_parakeet_onnx_model()
|
loaded = trainer._load_parakeet_onnx_model()
|
||||||
|
|
||||||
self.assertIs(loaded, fake_model)
|
self.assertIs(loaded, fake_model)
|
||||||
|
fake_huggingface_hub.snapshot_download.assert_called_once_with(
|
||||||
|
repo_id=trainer.DEFAULT_PARAKEET_ONNX_REPO,
|
||||||
|
local_dir=str(trainer.AUTO_TRAIN_MODEL_DIR),
|
||||||
|
allow_patterns=[
|
||||||
|
"config.json",
|
||||||
|
"vocab.txt",
|
||||||
|
"encoder-model.int8.onnx",
|
||||||
|
"encoder-model.int8.onnx.data",
|
||||||
|
"decoder_joint-model.int8.onnx",
|
||||||
|
"decoder_joint-model.int8.onnx.data",
|
||||||
|
],
|
||||||
|
)
|
||||||
|
fake_onnx_asr.load_model.assert_called_once_with(
|
||||||
|
trainer.DEFAULT_PARAKEET_ONNX_MODEL,
|
||||||
|
str(trainer.AUTO_TRAIN_MODEL_DIR),
|
||||||
|
quantization="int8",
|
||||||
|
providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_parakeet_loader_reuses_complete_snapshot_offline(self):
|
||||||
|
fake_model = object()
|
||||||
|
fake_onnx_asr = SimpleNamespace(load_model=Mock(return_value=fake_model))
|
||||||
|
fake_huggingface_hub = SimpleNamespace(snapshot_download=Mock())
|
||||||
|
trainer.AUTO_TRAIN_MODEL_DIR.mkdir(parents=True, exist_ok=True)
|
||||||
|
for filename in (
|
||||||
|
"config.json",
|
||||||
|
"vocab.txt",
|
||||||
|
"encoder-model.int8.onnx",
|
||||||
|
"decoder_joint-model.int8.onnx",
|
||||||
|
):
|
||||||
|
(trainer.AUTO_TRAIN_MODEL_DIR / filename).touch()
|
||||||
|
with (
|
||||||
|
patch.dict(
|
||||||
|
sys.modules,
|
||||||
|
{
|
||||||
|
"onnx_asr": fake_onnx_asr,
|
||||||
|
"huggingface_hub": fake_huggingface_hub,
|
||||||
|
},
|
||||||
|
),
|
||||||
|
patch.object(
|
||||||
|
trainer,
|
||||||
|
"_parakeet_onnx_providers",
|
||||||
|
return_value=["CUDAExecutionProvider", "CPUExecutionProvider"],
|
||||||
|
),
|
||||||
|
):
|
||||||
|
with trainer.PARAKEET_ONNX_MODEL_LOCK:
|
||||||
|
trainer.PARAKEET_ONNX_MODEL_CACHE.clear()
|
||||||
|
loaded = trainer._load_parakeet_onnx_model()
|
||||||
|
|
||||||
|
self.assertIs(loaded, fake_model)
|
||||||
|
fake_huggingface_hub.snapshot_download.assert_not_called()
|
||||||
fake_onnx_asr.load_model.assert_called_once_with(
|
fake_onnx_asr.load_model.assert_called_once_with(
|
||||||
trainer.DEFAULT_PARAKEET_ONNX_MODEL,
|
trainer.DEFAULT_PARAKEET_ONNX_MODEL,
|
||||||
str(trainer.AUTO_TRAIN_MODEL_DIR),
|
str(trainer.AUTO_TRAIN_MODEL_DIR),
|
||||||
|
|||||||
@@ -113,6 +113,10 @@ DEFAULT_PARAKEET_ONNX_MODEL = os.environ.get(
|
|||||||
"AUTO_TRAIN_PARAKEET_ONNX_MODEL",
|
"AUTO_TRAIN_PARAKEET_ONNX_MODEL",
|
||||||
"nemo-parakeet-tdt-0.6b-v3",
|
"nemo-parakeet-tdt-0.6b-v3",
|
||||||
)
|
)
|
||||||
|
DEFAULT_PARAKEET_ONNX_REPO = os.environ.get(
|
||||||
|
"AUTO_TRAIN_PARAKEET_ONNX_REPO",
|
||||||
|
"istupakov/parakeet-tdt-0.6b-v3-onnx",
|
||||||
|
)
|
||||||
DEFAULT_PARAKEET_ONNX_QUANTIZATION = "int8"
|
DEFAULT_PARAKEET_ONNX_QUANTIZATION = "int8"
|
||||||
|
|
||||||
AUTO_TRAIN_DEFAULT_CONFIG: Dict[str, Any] = {
|
AUTO_TRAIN_DEFAULT_CONFIG: Dict[str, Any] = {
|
||||||
@@ -953,6 +957,27 @@ def _load_parakeet_onnx_model():
|
|||||||
cached = PARAKEET_ONNX_MODEL_CACHE.get(cache_key)
|
cached = PARAKEET_ONNX_MODEL_CACHE.get(cache_key)
|
||||||
if cached is not None:
|
if cached is not None:
|
||||||
return cached
|
return cached
|
||||||
|
suffix = (
|
||||||
|
f".{DEFAULT_PARAKEET_ONNX_QUANTIZATION}"
|
||||||
|
if DEFAULT_PARAKEET_ONNX_QUANTIZATION
|
||||||
|
else ""
|
||||||
|
)
|
||||||
|
model_patterns = [
|
||||||
|
"config.json",
|
||||||
|
"vocab.txt",
|
||||||
|
f"encoder-model{suffix}.onnx",
|
||||||
|
f"encoder-model{suffix}.onnx.data",
|
||||||
|
f"decoder_joint-model{suffix}.onnx",
|
||||||
|
f"decoder_joint-model{suffix}.onnx.data",
|
||||||
|
]
|
||||||
|
required_model_files = [
|
||||||
|
"config.json",
|
||||||
|
"vocab.txt",
|
||||||
|
f"encoder-model{suffix}.onnx",
|
||||||
|
f"decoder_joint-model{suffix}.onnx",
|
||||||
|
]
|
||||||
|
if not DEFAULT_PARAKEET_ONNX_QUANTIZATION:
|
||||||
|
required_model_files.append("encoder-model.onnx.data")
|
||||||
AUTO_TRAIN_MODEL_DIR.mkdir(parents=True, exist_ok=True)
|
AUTO_TRAIN_MODEL_DIR.mkdir(parents=True, exist_ok=True)
|
||||||
previous = {
|
previous = {
|
||||||
key: os.environ.get(key)
|
key: os.environ.get(key)
|
||||||
@@ -962,9 +987,23 @@ def _load_parakeet_onnx_model():
|
|||||||
os.environ["HF_HUB_CACHE"] = str(AUTO_TRAIN_MODEL_DIR / "hub")
|
os.environ["HF_HUB_CACHE"] = str(AUTO_TRAIN_MODEL_DIR / "hub")
|
||||||
os.environ["HUGGINGFACE_HUB_CACHE"] = str(AUTO_TRAIN_MODEL_DIR / "hub")
|
os.environ["HUGGINGFACE_HUB_CACHE"] = str(AUTO_TRAIN_MODEL_DIR / "hub")
|
||||||
try:
|
try:
|
||||||
|
snapshot_root = AUTO_TRAIN_MODEL_DIR
|
||||||
|
if not all(
|
||||||
|
(AUTO_TRAIN_MODEL_DIR / filename).is_file()
|
||||||
|
for filename in required_model_files
|
||||||
|
):
|
||||||
|
from huggingface_hub import snapshot_download
|
||||||
|
|
||||||
|
snapshot_root = Path(
|
||||||
|
snapshot_download(
|
||||||
|
repo_id=DEFAULT_PARAKEET_ONNX_REPO,
|
||||||
|
local_dir=str(AUTO_TRAIN_MODEL_DIR),
|
||||||
|
allow_patterns=model_patterns,
|
||||||
|
)
|
||||||
|
)
|
||||||
model = onnx_asr.load_model(
|
model = onnx_asr.load_model(
|
||||||
DEFAULT_PARAKEET_ONNX_MODEL,
|
DEFAULT_PARAKEET_ONNX_MODEL,
|
||||||
str(AUTO_TRAIN_MODEL_DIR),
|
str(snapshot_root),
|
||||||
quantization=DEFAULT_PARAKEET_ONNX_QUANTIZATION,
|
quantization=DEFAULT_PARAKEET_ONNX_QUANTIZATION,
|
||||||
providers=list(providers),
|
providers=list(providers),
|
||||||
)
|
)
|
||||||
|
|||||||
Reference in New Issue
Block a user