wakebox 0.1.1__py3-none-any.whl

wakebox/__init__.py

@@ -0,0 +1,6 @@
+"""On-device wake word. No key."""
+
+from wakebox.engine import Engine, threshold_for_sensitivity
+from wakebox.microphone import Microphone
+
+__all__ = ["Engine", "Microphone", "threshold_for_sensitivity"]

wakebox/engine.py

@@ -0,0 +1,146 @@
+"""On-device wake word engine.
+
+The shared runtime is the frozen mel frontend and speech embedding. Each
+phrase file is a small head. Sensitivity is chosen here, not stored in the file.
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import cast
+
+import numpy as np
+import onnxruntime as ort
+
+SAMPLE_RATE = 16000
+MEL_WINDOW = 76
+MEL_STRIDE = 8
+EMBEDDING_DIM = 96
+HEAD_FRAMES = 16
+REFRACTORY_SEC = 1.0
+# Scores for a real phrase sit in a narrow band near the top. Map the
+# caller's 0–1 control onto that band. 0 is strict, 1 is loose.
+THRESHOLD_STRICT = 0.995
+THRESHOLD_LOOSE = 0.85
+
+
+def threshold_for_sensitivity(sensitivity: float) -> float:
+    value = float(np.clip(sensitivity, 0.0, 1.0))
+    return THRESHOLD_STRICT + (THRESHOLD_LOOSE - THRESHOLD_STRICT) * value
+
+
+def _session(model: str | Path | bytes) -> ort.InferenceSession:
+    options = ort.SessionOptions()
+    options.inter_op_num_threads = 1
+    options.intra_op_num_threads = 1
+    if isinstance(model, bytes):
+        return ort.InferenceSession(model, sess_options=options, providers=["CPUExecutionProvider"])
+    return ort.InferenceSession(str(model), sess_options=options, providers=["CPUExecutionProvider"])
+
+
+class Engine:
+    """Load phrase files and score 16 kHz mono int16 frames.
+
+    ``process`` returns the index of the phrase that fired, or ``None``.
+    """
+
+    def __init__(self, phrases: list[str | Path | bytes], sensitivities: list[float] | None = None):
+        if not phrases:
+            raise ValueError("load at least one phrase file")
+        if sensitivities is None:
+            sensitivities = [0.5] * len(phrases)
+        if len(sensitivities) != len(phrases):
+            raise ValueError("sensitivities must match the phrase files")
+        here = Path(__file__).resolve().parent / "models"
+        self._mel = _session(here / "melspectrogram.onnx")
+        self._embedding = _session(here / "embedding_model.onnx")
+        self._mel_input = self._mel.get_inputs()[0].name
+        self._embedding_input = self._embedding.get_inputs()[0].name
+        self._heads = [_session(phrase) for phrase in phrases]
+        self._head_input = self._heads[0].get_inputs()[0].name
+        self.sensitivities = [float(value) for value in sensitivities]
+        self._audio = np.zeros(0, dtype=np.int16)
+        self._scored = 0
+        self._previous = [0.0] * len(self._heads)
+        self._last_fire = [-1e9] * len(self._heads)
+        self._samples = 0
+
+    def set_sensitivity(self, index: int, sensitivity: float) -> None:
+        self.sensitivities[index] = float(sensitivity)
+
+    def process(self, frame: np.ndarray) -> int | None:
+        samples = np.asarray(frame, dtype=np.int16).reshape(-1)
+        if samples.size == 0:
+            return None
+        self._samples += int(samples.size)
+        self._audio = np.concatenate([self._audio, samples])
+        max_samples = SAMPLE_RATE * 4
+        if self._audio.size > max_samples:
+            self._audio = self._audio[-max_samples:]
+            self._scored = 0
+        embeddings = self._embed(self._audio)
+        fired: int | None = None
+        now = self._samples / SAMPLE_RATE
+        while self._scored < len(embeddings):
+            index = self._scored
+            self._scored += 1
+            if index + 1 < HEAD_FRAMES:
+                continue
+            window = embeddings[index + 1 - HEAD_FRAMES : index + 1]
+            choice = self._score(window, now)
+            if choice is not None:
+                fired = choice
+        return fired
+
+    def process_buffer(self, audio: np.ndarray, frame: int = 1280) -> list[int]:
+        """Feed a whole clip and return the phrase index of each detection."""
+        samples = np.asarray(audio, dtype=np.int16).reshape(-1)
+        hits = []
+        for start in range(0, len(samples), frame):
+            index = self.process(samples[start : start + frame])
+            if index is not None:
+                hits.append(index)
+        return hits
+
+    def _embed(self, audio: np.ndarray) -> np.ndarray:
+        spec = self._melspectrogram(audio)
+        if spec.shape[0] < MEL_WINDOW:
+            return np.zeros((0, EMBEDDING_DIM), dtype=np.float32)
+        windows = [
+            spec[i : i + MEL_WINDOW]
+            for i in range(0, spec.shape[0] - MEL_WINDOW + 1, MEL_STRIDE)
+        ]
+        batch = np.asarray(windows, dtype=np.float32)[..., None]
+        raw = self._embedding.run(None, {self._embedding_input: batch})[0]
+        return cast(np.ndarray, raw.reshape(batch.shape[0], EMBEDDING_DIM))
+
+    def _melspectrogram(self, audio: np.ndarray) -> np.ndarray:
+        data = np.ascontiguousarray(audio.astype(np.float32))[None, :]
+        spec = np.squeeze(self._mel.run(None, {self._mel_input: data})[0])
+        if spec.ndim == 1:
+            spec = spec[None, :]
+        if spec.shape[-1] != 32 and spec.shape[0] == 32:
+            spec = spec.T
+        return cast(np.ndarray, (spec / 10.0 + 2.0).astype(np.float32))
+
+    def _score(self, window: np.ndarray, now: float) -> int | None:
+        features = np.ascontiguousarray(window.astype(np.float32)[None, :])
+        best_index = None
+        best_score = -1.0
+        for index, head in enumerate(self._heads):
+            logit = np.asarray(head.run(None, {self._head_input: features})[0]).reshape(-1)
+            score = float(1.0 / (1.0 + np.exp(-float(logit[0]))))
+            threshold = threshold_for_sensitivity(self.sensitivities[index])
+            crossed = self._previous[index] < threshold <= score
+            self._previous[index] = score
+            if not crossed:
+                continue
+            if now - self._last_fire[index] < REFRACTORY_SEC:
+                continue
+            if score > best_score:
+                best_score = score
+                best_index = index
+        if best_index is None:
+            return None
+        self._last_fire[best_index] = now
+        return best_index

wakebox/microphone.py

@@ -0,0 +1,55 @@
+"""Microphone helper. The frame API is what tests drive."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import Protocol
+
+import numpy as np
+
+from wakebox.engine import Engine
+
+
+class _Capture(Protocol):
+    def start(self) -> None: ...
+
+    def stop(self) -> None: ...
+
+    def close(self) -> None: ...
+
+
+class Microphone:
+    """Read the default input and call back with the phrase index."""
+
+    def __init__(self, engine: Engine, on_detection: Callable[[int], None]):
+        self.engine = engine
+        self.on_detection = on_detection
+        self._stream: _Capture | None = None
+
+    def start(self) -> None:
+        import sounddevice as sd
+
+        def callback(indata: np.ndarray, frames: int, time: object, status: object) -> None:
+            del frames, time, status
+            samples = np.asarray(indata, dtype=np.int16).reshape(-1)
+            index = self.engine.process(samples)
+            if index is not None:
+                self.on_detection(index)
+
+        stream = sd.InputStream(
+            samplerate=16000,
+            channels=1,
+            dtype="int16",
+            blocksize=1280,
+            callback=callback,
+        )
+        self._stream = stream
+        stream.start()
+
+    def stop(self) -> None:
+        stream = self._stream
+        self._stream = None
+        if stream is None:
+            return
+        stream.stop()
+        stream.close()

wakebox-0.1.1.dist-info/METADATA

@@ -0,0 +1,31 @@
+Metadata-Version: 2.4
+Name: wakebox
+Version: 0.1.1
+Summary: On-device wake word. No key.
+License-Expression: MIT
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+Requires-Dist: numpy
+Requires-Dist: onnxruntime
+Requires-Dist: sounddevice
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: soundfile; extra == "dev"
+
+On-device wake word, no key.
+
+```bash
+pip install wakebox
+```
+
+Create a phrase file at [https://wakebox.ai](https://wakebox.ai), then:
+
+```python
+from wakebox import Engine
+import numpy as np
+
+engine = Engine(["phrase.onnx"], sensitivities=[0.5])
+index = engine.process(np.zeros(1280, dtype=np.int16))  # 0, or None
+```
+
+`process` takes 16 kHz mono int16 audio and returns the phrase index, or `None`. Sensitivity is a number from 0 to 1 on each file. It is not stored in the file.

wakebox-0.1.1.dist-info/RECORD

@@ -0,0 +1,10 @@
+wakebox/__init__.py,sha256=dSxcJf8M9xjNeQEmb8zCeFPF7phrnIB5_zimS7esPF0,204
+wakebox/engine.py,sha256=oWrUp1E7JI5cYcHJ61sIgbQNRl3LD9dnAbwTv4O-K-Y,5934
+wakebox/microphone.py,sha256=LJ0FYWM_hS_tw2EuLsTdaAF8QhKvogH-YCPnz4i9KEY,1430
+wakebox/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+wakebox/models/embedding_model.onnx,sha256=cNFkKQwdCV0dTuFJvF4AVDJQpzFrWfMdBWz_e9MHXB8,1326578
+wakebox/models/melspectrogram.onnx,sha256=uisOD4t7h1NposicsTNg_1O6xDbyiVzO2fR5-mXrF28,1087958
+wakebox-0.1.1.dist-info/METADATA,sha256=rhJNV08uvF8SRvR88ZeOHqX5Dn3jcH6gmgnqqdMb9k8,834
+wakebox-0.1.1.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
+wakebox-0.1.1.dist-info/top_level.txt,sha256=N2fNG_ATcdq0JdNMNUSKwYxHq3iJeM8VEKGpKZYx3B8,8
+wakebox-0.1.1.dist-info/RECORD,,

wakebox-0.1.1.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

wakebox-0.1.1.dist-info/top_level.txt

@@ -0,0 +1 @@
+wakebox