Release NVIDIA WakeWord Trainer v20

VERSION

@@ -1 +1 @@
-18
+20

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 the v19 container startup failure caused by malformed indentation in the Parakeet ONNX loader.
+- Preserved automatic download, resume, and offline reuse of the required Parakeet INT8 model snapshot.
+- Revalidated both CUDA and CPU Parakeet provider paths with the complete trainer test suite.

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),

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),
             )