Release NVIDIA WakeWord Trainer v12

This commit is contained in:
MasterPhooey
2026-07-17 08:21:50 -05:00
parent 3d341d0617
commit 7d77f71dc3
13 changed files with 1682 additions and 34 deletions

238
tests/test_auto_train.py Normal file
View File

@@ -0,0 +1,238 @@
import io
import json
import sys
import tempfile
import unittest
import wave
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import Mock, patch
import trainer_server as trainer
def silent_wav_bytes(duration_s: float = 0.25) -> bytes:
output = io.BytesIO()
with wave.open(output, "wb") as wav_file:
wav_file.setnchannels(1)
wav_file.setsampwidth(2)
wav_file.setframerate(16000)
wav_file.writeframes(b"\x00\x00" * int(16000 * duration_s))
return output.getvalue()
class AutoTrainTests(unittest.TestCase):
def setUp(self):
self.tempdir = tempfile.TemporaryDirectory()
root = Path(self.tempdir.name)
self.original_paths = (
trainer.CAPTURED_DIR,
trainer.NEGATIVE_DIR,
trainer.PERSONAL_DIR,
trainer.AUTO_TRAIN_CONFIG_FILE,
trainer.AUTO_TRAIN_STATE_FILE,
)
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"
for directory in (trainer.CAPTURED_DIR, trainer.NEGATIVE_DIR, trainer.PERSONAL_DIR):
directory.mkdir(parents=True)
self.original_config = dict(trainer.AUTO_TRAIN_CONFIG)
self.original_state = dict(trainer.AUTO_TRAIN_STATE)
trainer.AUTO_TRAIN_CONFIG.clear()
trainer.AUTO_TRAIN_CONFIG.update(
trainer._normalize_auto_train_config(
{
"enabled": True,
"wake_phrase": "hey tater",
"language": "en",
"tater_url": "http://127.0.0.1:8501",
"stt_device": "auto",
"stt_compute_type": "auto",
}
)
)
trainer.AUTO_TRAIN_STATE.clear()
trainer.AUTO_TRAIN_STATE.update(trainer.AUTO_TRAIN_DEFAULT_STATE)
def tearDown(self):
(
trainer.CAPTURED_DIR,
trainer.NEGATIVE_DIR,
trainer.PERSONAL_DIR,
trainer.AUTO_TRAIN_CONFIG_FILE,
trainer.AUTO_TRAIN_STATE_FILE,
) = self.original_paths
trainer.AUTO_TRAIN_CONFIG.clear()
trainer.AUTO_TRAIN_CONFIG.update(self.original_config)
trainer.AUTO_TRAIN_STATE.clear()
trainer.AUTO_TRAIN_STATE.update(self.original_state)
self.tempdir.cleanup()
def add_capture(self, name: str = "wake.wav", wake_word: str = "hey_tater") -> Path:
audio_path = trainer.CAPTURED_DIR / name
audio_path.write_bytes(silent_wav_bytes())
trainer._write_sidecar_json(
audio_path,
{
"original_name": name,
"wake_word": wake_word,
"event_type": "wake_detected",
"review_status": "pending",
},
)
return audio_path
def test_phrase_matching_normalizes_case_punctuation_and_underscores(self):
self.assertTrue(trainer._transcript_contains_wake_phrase("Okay, HEY TATER!", "hey_tater"))
self.assertFalse(trainer._transcript_contains_wake_phrase("Turn on the television", "hey tater"))
def test_phrase_miss_moves_wake_trigger_to_negative_samples(self):
self.add_capture()
with patch.object(trainer, "_transcribe_capture_with_faster_whisper", return_value="turn on the kitchen lights"):
trainer._auto_review_capture("wake.wav")
self.assertFalse((trainer.CAPTURED_DIR / "wake.wav").exists())
negatives = list(trainer.NEGATIVE_DIR.glob("*.wav"))
self.assertEqual(len(negatives), 1)
metadata = trainer._load_sidecar_json(negatives[0])
self.assertTrue(metadata["auto_negative"])
self.assertEqual(metadata["review_status"], "auto_approved_negative")
self.assertEqual(metadata["transcript"], "turn on the kitchen lights")
self.assertEqual(trainer.AUTO_TRAIN_STATE["pending_negative_count"], 1)
def test_matching_phrase_stays_in_manual_review_inbox(self):
audio_path = self.add_capture()
with patch.object(trainer, "_transcribe_capture_with_faster_whisper", return_value="hey tater turn on the lights"):
trainer._auto_review_capture("wake.wav")
self.assertTrue(audio_path.exists())
self.assertFalse(list(trainer.NEGATIVE_DIR.glob("*.wav")))
metadata = trainer._load_sidecar_json(audio_path)
self.assertEqual(metadata["auto_review_status"], "wake_phrase_detected")
self.assertEqual(trainer.AUTO_TRAIN_STATE["pending_negative_count"], 0)
def test_capture_for_another_wake_word_is_not_transcribed(self):
audio_path = self.add_capture(wake_word="computer")
with patch.object(trainer, "_transcribe_capture_with_faster_whisper") as transcribe:
trainer._auto_review_capture("wake.wav")
transcribe.assert_not_called()
self.assertTrue(audio_path.exists())
metadata = trainer._load_sidecar_json(audio_path)
self.assertEqual(metadata["auto_review_status"], "different_wake_phrase")
def test_due_schedule_starts_training_after_minimum_negatives(self):
trainer.AUTO_TRAIN_CONFIG["schedule_hours"] = 24
trainer.AUTO_TRAIN_CONFIG["minimum_new_negatives"] = 3
trainer.AUTO_TRAIN_STATE["pending_negative_count"] = 3
trainer.AUTO_TRAIN_STATE["next_run_at"] = "2000-01-01T00:00:00+00:00"
with patch.object(trainer, "_start_auto_training", return_value={"ok": True, "started": True}) as start:
trainer._maybe_run_scheduled_auto_training()
start.assert_called_once_with()
self.assertTrue(trainer.AUTO_TRAIN_STATE["next_run_at"])
def test_tater_refresh_repushes_settings_with_selector_and_token(self):
trainer.AUTO_TRAIN_CONFIG.update(
{
"notify_satellites": True,
"tater_url": "http://127.0.0.1:8501",
"tater_selector": "kitchen-sat",
"tater_api_token": "secret-token",
}
)
class Response:
def __enter__(self):
return self
def __exit__(self, *_args):
return False
def read(self):
return b'{"push":{"count":2}}'
with patch.object(trainer, "urlopen", return_value=Response()) as open_url:
result = trainer._notify_tater_satellites()
self.assertTrue(result["ok"])
self.assertEqual(result["count"], 2)
request = open_url.call_args.args[0]
self.assertEqual(request.full_url, "http://127.0.0.1:8501/api/tater/satellite/v1/settings")
self.assertEqual(request.get_header("X-tater-token"), "secret-token")
self.assertEqual(json.loads(request.data), {"selector": "kitchen-sat", "settings": {}})
def test_advertised_url_uses_non_loopback_browser_host(self):
request = SimpleNamespace(
base_url="http://192.168.1.50:8789/",
url=SimpleNamespace(hostname="192.168.1.50", scheme="http", port=8789),
)
self.assertEqual(trainer._advertised_base_url(request), "http://192.168.1.50:8789")
def test_advertised_url_replaces_localhost_with_discovered_lan_host(self):
request = SimpleNamespace(
base_url="http://127.0.0.1:8789/",
url=SimpleNamespace(hostname="127.0.0.1", scheme="http", port=8789),
)
with patch.object(trainer, "_discover_lan_ipv4", return_value="192.168.1.60"):
self.assertEqual(trainer._advertised_base_url(request), "http://192.168.1.60:8789")
def test_configured_public_url_takes_precedence(self):
trainer.AUTO_TRAIN_CONFIG["advertised_base_url"] = "http://trainer.local:8789"
request = SimpleNamespace(
base_url="http://127.0.0.1:8789/",
url=SimpleNamespace(hostname="127.0.0.1", scheme="http", port=8789),
)
self.assertEqual(trainer._advertised_base_url(request), "http://trainer.local:8789")
def test_faster_whisper_auto_runtime_prefers_cuda_and_float16(self):
fake_ctranslate2 = SimpleNamespace(get_cuda_device_count=lambda: 1)
with patch.dict(sys.modules, {"ctranslate2": fake_ctranslate2}):
self.assertEqual(
trainer._resolve_faster_whisper_runtime("auto", "auto"),
("cuda", "float16"),
)
def test_faster_whisper_auto_runtime_falls_back_to_cpu_int8(self):
fake_ctranslate2 = SimpleNamespace(get_cuda_device_count=lambda: 0)
with patch.dict(sys.modules, {"ctranslate2": fake_ctranslate2}):
self.assertEqual(
trainer._resolve_faster_whisper_runtime("auto", "auto"),
("cpu", "int8"),
)
def test_faster_whisper_transcription_joins_segments_and_records_runtime(self):
fake_model = SimpleNamespace()
fake_model.transcribe = Mock(
return_value=(
iter([SimpleNamespace(text=" turn on "), SimpleNamespace(text="the lights ")]),
SimpleNamespace(),
)
)
with (
patch.object(trainer, "_resolve_faster_whisper_runtime", return_value=("cuda", "float16")),
patch.object(trainer, "_load_faster_whisper_model", return_value=fake_model),
):
transcript = trainer._transcribe_capture_with_faster_whisper(
Path("wake.wav"),
model="small.en",
language="en",
)
self.assertEqual(transcript, "turn on the lights")
fake_model.transcribe.assert_called_once_with(
"wake.wav",
language="en",
beam_size=1,
condition_on_previous_text=False,
)
self.assertEqual(trainer.AUTO_TRAIN_STATE["last_stt_device"], "cuda")
self.assertEqual(trainer.AUTO_TRAIN_STATE["last_stt_compute_type"], "float16")
if __name__ == "__main__":
unittest.main()

View File

@@ -0,0 +1,88 @@
import importlib.util
import unittest
from pathlib import Path
SCRIPT_PATH = (
Path(__file__).resolve().parents[1]
/ "cli"
/ "calibrate_detector.py"
)
SPEC = importlib.util.spec_from_file_location("calibrate_detector", SCRIPT_PATH)
calibrate_detector = importlib.util.module_from_spec(SPEC)
assert SPEC.loader is not None
SPEC.loader.exec_module(calibrate_detector)
def candidate(cutoff, window, recall, false_accepts_per_hour):
return {
"probability_cutoff": cutoff,
"sliding_window_size": window,
"recall": recall,
"false_accepts_per_hour": false_accepts_per_hour,
}
class CalibrationSelectionTests(unittest.TestCase):
def test_defaults_are_conservative(self):
self.assertEqual(calibrate_detector.DEFAULT_WINDOW_SIZES, [5, 6, 7])
self.assertEqual(calibrate_detector.DEFAULT_CUTOFF_MIN, 0.95)
self.assertEqual(calibrate_detector.DEFAULT_RECALL_MARGIN, 0.005)
def test_prefers_zero_false_accepts_within_recall_margin(self):
candidates = [
candidate(0.95, 5, 0.99894, 0.103408),
candidate(0.95, 6, 0.99744, 0.0),
candidate(0.95, 7, 0.99554, 0.0),
]
best, selected_limit = calibrate_detector._select_best_candidate(
candidates,
target_faph=0.25,
recall_margin=0.005,
)
self.assertEqual(best["sliding_window_size"], 6)
self.assertEqual(best["false_accepts_per_hour"], 0.0)
self.assertEqual(selected_limit, 0.25)
def test_does_not_trade_away_recall_beyond_margin(self):
candidates = [
candidate(0.95, 5, 0.99, 0.1),
candidate(0.99, 6, 0.90, 0.0),
]
best, _ = calibrate_detector._select_best_candidate(
candidates,
target_faph=0.25,
recall_margin=0.005,
)
self.assertEqual(best["sliding_window_size"], 5)
def test_uses_strictest_available_false_accept_tier(self):
candidates = [
candidate(0.95, 5, 0.99, 0.6),
candidate(0.99, 6, 0.99, 1.5),
]
best, selected_limit = calibrate_detector._select_best_candidate(
candidates,
target_faph=0.25,
recall_margin=0.005,
)
self.assertEqual(best["false_accepts_per_hour"], 0.6)
self.assertEqual(selected_limit, 0.75)
def test_rejects_negative_recall_margin(self):
with self.assertRaises(ValueError):
calibrate_detector._select_best_candidate(
[candidate(0.95, 6, 0.99, 0.0)],
target_faph=0.25,
recall_margin=-0.001,
)
if __name__ == "__main__":
unittest.main()