Upload

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

bolt_pet/audio/wake_word.py

@@ -0,0 +1,185 @@
+"""Wake-word detection via a custom-trained openWakeWord model.
+
+Uses `thunderbolt.onnx` — trained specifically for "thunderbolt", the same
+way the main repo's `desk_client/bolt_desk.py` uses `bolt.onnx` for "hey
+bolt". Same runtime (openWakeWord, ONNX inference), same per-frame
+predict()/reset() pattern; the only difference is the model file
+(WAKE_MODEL_FILE) and threshold (WAKE_WORD_THRESHOLD), both configurable via
+.env if a differently-trained model is swapped in later.
+"""
+
+from __future__ import annotations
+
+from collections import deque
+from typing import Callable, Optional, Protocol, Union
+
+import numpy as np
+
+from .. import config
+
+
+class WakeModel(Protocol):
+    def predict(self, frame: np.ndarray) -> dict: ...
+    def reset(self) -> None: ...
+
+
+class NearMissLog:
+    """Rolling record of frames that *almost* fired the wake word.
+
+    WAKE_WORD_THRESHOLD is otherwise tuned by guessing at a number in .env
+    and seeing whether the pet ignores you. Keeping the near misses (scores
+    within WAKE_NEAR_MISS_MARGIN below the threshold) turns that into
+    evidence: the tray's wake-word tuner shows what your actual "thunderbolt"
+    scores, so you can set the threshold just under it.
+
+    Pure bookkeeping — the caller supplies timestamps, so it's testable.
+    """
+
+    def __init__(self, limit: int = None, margin: float = None):
+        self._entries: deque[tuple[float, float, float]] = deque(  # (timestamp, score, threshold)
+            maxlen=max(1, config.WAKE_NEAR_MISS_LIMIT if limit is None else limit)
+        )
+        self._margin = config.WAKE_NEAR_MISS_MARGIN if margin is None else margin
+        self._peak = 0.0
+
+    @property
+    def peak(self) -> float:
+        """Highest score seen since the last reset — the "how close did I
+        get?" readout while you test the wake phrase."""
+        return self._peak
+
+    def observe(self, score: float, threshold: float, timestamp: float) -> bool:
+        """Record *score*; returns True if it counted as a near miss."""
+        self._peak = max(self._peak, score)
+        if score >= threshold or score < threshold - self._margin:
+            return False
+        self._entries.append((timestamp, score, threshold))
+        return True
+
+    def entries(self) -> list[tuple[float, float, float]]:
+        return list(self._entries)
+
+    def clear(self) -> None:
+        self._entries.clear()
+        self._peak = 0.0
+
+
+def _construct_model(model_cls, model_path: str):
+    """openwakeword's Model() constructor keyword has drifted across
+    releases (wakeword_models -> wakeword_model_paths) and some builds
+    reject inference_framework entirely — the main repo's ai/wake_word.py
+    hit the same drift and works around it the same way: try each known
+    calling convention in turn."""
+    attempts = [
+        lambda: model_cls(wakeword_model_paths=[model_path], inference_framework="onnx"),
+        lambda: model_cls(wakeword_model_paths=[model_path]),
+        lambda: model_cls(wakeword_models=[model_path], inference_framework="onnx"),
+        lambda: model_cls(wakeword_models=[model_path]),
+        lambda: model_cls([model_path]),
+    ]
+    last_exc: Optional[TypeError] = None
+    for attempt in attempts:
+        try:
+            return attempt()
+        except TypeError as exc:
+            last_exc = exc
+    raise RuntimeError(
+        f"Could not construct openwakeword.Model with any known calling convention "
+        f"(last error: {last_exc})"
+    )
+
+
+class _OpenWakeWordModel:
+    """Lazily loads the ONNX model on first use so importing this module
+    (and unit-testing listen_for_wake_word with a fake model) never requires
+    onnxruntime/openwakeword or the model file to be present."""
+
+    def __init__(self):
+        self._model = None
+
+    def _ensure_model(self):
+        if self._model is None:
+            from openwakeword.model import Model
+            # from openwakeword.utils import download_models
+
+            # # The pip package doesn't bundle its melspectrogram/embedding
+            # # feature-extraction sub-models — fetch them once on first use
+            # # (no-op if already cached in openwakeword's own resources dir).
+            # # A non-empty, non-matching model_names list keeps this from
+            # # also pulling every official pretrained wakeword model.
+            # download_models(model_names=["thunderbolt"])
+
+            self._model = _construct_model(Model, config.WAKE_MODEL_PATH)
+        return self._model
+
+    def predict(self, frame: np.ndarray) -> dict:
+        return self._ensure_model().predict(frame)
+
+    def reset(self) -> None:
+        if self._model is not None:
+            self._model.reset()
+
+
+_default_model = _OpenWakeWordModel()
+
+
+def listen_for_wake_word(
+    stream,
+    should_continue=lambda: True,
+    model: Optional[WakeModel] = None,
+    threshold: Union[float, Callable[[], float], None] = None,
+    on_tick=None,
+    on_score: Optional[Callable[[float, float], None]] = None,
+) -> bool:
+    """Block until the wake word fires (returns True) or *should_continue*
+    goes false (returns False).
+
+    Feeds every frame to *model* (the thunderbolt openWakeWord model by
+    default) and treats any class score >= *threshold* as a detection,
+    resetting the model's internal state afterward so the next call starts
+    clean — same pattern as desk_client/bolt_desk.py's main loop.
+
+    *threshold* may be a number or a zero-argument callable. The callable
+    form exists because this function blocks for minutes at a time: the
+    tray's wake-word tuner slider has to be able to change sensitivity
+    *during* a listen, not only at the start of the next one.
+
+    *on_tick*, if given, is called once per ``WAKE_CHECK_INTERVAL_SECONDS``
+    (not every frame — prediction is cheap enough to run on every frame, but
+    this is the only point control returns to the caller while otherwise
+    blocked here for a possibly long time, so it's how a caller drives
+    periodic work, e.g. the heartbeat/announcement poll in controller.py,
+    during quiet stretches with no wake word).
+
+    *on_score*, if given, gets ``(best_score, threshold)`` every frame — used
+    to log near misses for threshold tuning.
+    """
+    model = model or _default_model
+    if threshold is None:
+        threshold = config.WAKE_WORD_THRESHOLD
+    resolve_threshold = threshold if callable(threshold) else (lambda: threshold)
+
+    frame_len = config.FRAME_LEN
+    check_every_frames = max(1, int(config.WAKE_CHECK_INTERVAL_SECONDS * config.SAMPLE_RATE / frame_len))
+    frames_since_tick = 0
+
+    while should_continue():
+        chunk, _ = stream.read(frame_len)
+        frame = np.asarray(chunk)[:, 0]
+        scores = model.predict(frame)
+        current_threshold = resolve_threshold()
+        best = max(scores.values()) if scores else 0.0
+
+        frames_since_tick += 1
+        if frames_since_tick >= check_every_frames:
+            frames_since_tick = 0
+            if on_tick is not None:
+                on_tick()
+
+        if on_score is not None:
+            on_score(best, current_threshold)
+
+        if scores and best >= current_threshold:
+            model.reset()
+            return True
+    return False