From 518df6316166c6144eb796938bb10754033712e0 Mon Sep 17 00:00:00 2001 From: MasterPhooey Date: Sun, 26 Jul 2026 11:23:35 -0500 Subject: [PATCH] Release NVIDIA WakeWord Trainer v19 --- VERSION | 2 +- WHATS_NEW.md | 6 ++-- tests/test_auto_train.py | 65 +++++++++++++++++++++++++++++++++++++++- trainer_server.py | 41 ++++++++++++++++++++++++- 4 files changed, 108 insertions(+), 6 deletions(-) diff --git a/VERSION b/VERSION index 3c03207..d6b2404 100644 --- a/VERSION +++ b/VERSION @@ -1 +1 @@ -18 +19 diff --git a/WHATS_NEW.md b/WHATS_NEW.md index cde4096..258b87e 100644 --- a/WHATS_NEW.md +++ b/WHATS_NEW.md @@ -1,3 +1,3 @@ -- Added Parakeet ONNX as a second local Auto Training STT engine alongside Faster Whisper. -- Replaced manual model, device, and compute fields with a simple engine selector and managed language-aware models. -- Added CUDA-enabled ONNX Runtime with CPU fallback, runtime reporting, and model-cache cleanup when switching engines. +- Fixed first-run Parakeet ONNX setup failing when its empty model directory was mistaken for a complete offline model. +- Parakeet now downloads or resumes the required INT8 snapshot before loading through ONNX ASR. +- Complete local snapshots are reused without Hub access, preserving offline startup after the initial download. diff --git a/tests/test_auto_train.py b/tests/test_auto_train.py index 4ef9b5d..9332042 100644 --- a/tests/test_auto_train.py +++ b/tests/test_auto_train.py @@ -43,12 +43,14 @@ class AutoTrainTests(unittest.TestCase): trainer.PERSONAL_DIR, trainer.AUTO_TRAIN_CONFIG_FILE, trainer.AUTO_TRAIN_STATE_FILE, + trainer.AUTO_TRAIN_MODEL_DIR, ) trainer.CAPTURED_DIR = root / "captured_audio" trainer.NEGATIVE_DIR = root / "negative_samples" trainer.PERSONAL_DIR = root / "personal_samples" trainer.AUTO_TRAIN_CONFIG_FILE = root / "auto_train_config.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): directory.mkdir(parents=True) @@ -75,6 +77,7 @@ class AutoTrainTests(unittest.TestCase): trainer.PERSONAL_DIR, trainer.AUTO_TRAIN_CONFIG_FILE, trainer.AUTO_TRAIN_STATE_FILE, + trainer.AUTO_TRAIN_MODEL_DIR, ) = self.original_paths trainer.AUTO_TRAIN_CONFIG.clear() 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): fake_model = object() 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 ( - 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( trainer, "_parakeet_onnx_providers", @@ -173,6 +185,57 @@ class AutoTrainTests(unittest.TestCase): loaded = trainer._load_parakeet_onnx_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( trainer.DEFAULT_PARAKEET_ONNX_MODEL, str(trainer.AUTO_TRAIN_MODEL_DIR), diff --git a/trainer_server.py b/trainer_server.py index 5fc47a8..101c59b 100644 --- a/trainer_server.py +++ b/trainer_server.py @@ -113,6 +113,10 @@ DEFAULT_PARAKEET_ONNX_MODEL = os.environ.get( "AUTO_TRAIN_PARAKEET_ONNX_MODEL", "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" AUTO_TRAIN_DEFAULT_CONFIG: Dict[str, Any] = { @@ -953,6 +957,27 @@ def _load_parakeet_onnx_model(): cached = PARAKEET_ONNX_MODEL_CACHE.get(cache_key) if cached is not None: 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) previous = { 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["HUGGINGFACE_HUB_CACHE"] = str(AUTO_TRAIN_MODEL_DIR / "hub") 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( DEFAULT_PARAKEET_ONNX_MODEL, - str(AUTO_TRAIN_MODEL_DIR), + str(snapshot_root), quantization=DEFAULT_PARAKEET_ONNX_QUANTIZATION, providers=list(providers), )