Fix detection on onnxruntime 1.27.0 / arm64 (KleidiAI batch-1 conv bug)

.github/workflows/ci.yml

@@ -31,3 +31,21 @@ jobs:
 
       - name: Run tests
         run: uv run pytest --tb=short -q
+
+  # onnxruntime's arm64 kernels differ from x86 (KleidiAI); run the pretrained-model check on Apple silicon too.
+  inference-arm64:
+    runs-on: macos-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v4
+
+      - name: Set up Python
+        run: uv python install 3.12
+
+      - name: Install dependencies
+        run: uv sync
+
+      - name: Run inference tests
+        run: uv run pytest tests/test_pretrained.py --noconftest --tb=short -q

src/livekit/wakeword/eval/evaluate.py

@@ -10,6 +10,7 @@
 import onnxruntime as ort
 
 from ..config import WakeWordConfig
+from ..session import session_options
 
 logger = logging.getLogger(__name__)
 
@@ -194,7 +195,9 @@ def run_eval(config: WakeWordConfig, model_path: str | Path) -> dict[str, float]
     if not model_path.exists():
         raise FileNotFoundError(f"Model not found: {model_path}")
 
-    session = ort.InferenceSession(str(model_path), providers=["CPUExecutionProvider"])
+    session = ort.InferenceSession(
+        str(model_path), session_options(), providers=["CPUExecutionProvider"]
+    )
     logger.info(f"Loaded model from {model_path}")
 
     # Load validation data

src/livekit/wakeword/inference/model.py

@@ -10,6 +10,7 @@
 
 from ..models.feature_extractor import MelSpectrogramFrontend, SpeechEmbedding
 from ..resources import get_embedding_model_path, get_mel_model_path
+from ..session import session_options
 
 logger = logging.getLogger(__name__)
 
@@ -89,7 +90,7 @@ def load_model(self, model_path: str | Path, model_name: str | None = None, sess
         session = ort.InferenceSession(
             str(model_path),
             providers=["CPUExecutionProvider"],
-            sess_options=sess_options
+            sess_options=session_options(sess_options),
         )
         input_name = session.get_inputs()[0].name
         self._classifiers[model_name] = (session, input_name)

src/livekit/wakeword/models/feature_extractor.py

@@ -15,6 +15,8 @@
 import numpy as np
 from onnxruntime.capi.onnxruntime_pybind11_state import SessionOptions
 
+from ..session import session_options
+
 logger = logging.getLogger(__name__)
 
 class MelSpectrogramFrontend:
@@ -43,7 +45,7 @@ def _init_onnx(self, onnx_path: str | Path, sess_options: SessionOptions = None)
         self._onnx_session = ort.InferenceSession(
             str(onnx_path),
             providers=["CPUExecutionProvider"],
-            sess_options=sess_options,
+            sess_options=session_options(sess_options),
         )
         self._input_name = self._onnx_session.get_inputs()[0].name
         logger.info(f"Loaded mel ONNX model from {onnx_path}")
@@ -103,7 +105,7 @@ def __init__(self, onnx_path: str | Path, sess_options: SessionOptions = None):
         self._session = ort.InferenceSession(
             str(onnx_path),
             providers=["CPUExecutionProvider"],
-            sess_options=sess_options,
+            sess_options=session_options(sess_options),
         )
         self._input_name = self._session.get_inputs()[0].name
         logger.info(f"Loaded embedding ONNX model from {onnx_path}")

src/livekit/wakeword/session.py

@@ -0,0 +1,20 @@
+"""onnxruntime session options shared by every model the package runs."""
+
+from __future__ import annotations
+
+from onnxruntime import SessionOptions
+
+# onnxruntime 1.27.0 on arm64 (e.g. Apple silicon) runs batch-1 convolutions through its KleidiAI
+# kernels, which return wrong results. The speech-embedding model runs one window at a time, so
+# every wake word scored ~0 there. Older onnxruntime versions accept this setting and ignore it.
+DISABLE_KLEIDIAI = "mlas.disable_kleidiai"
+
+
+def session_options(sess_options: SessionOptions | None = None) -> SessionOptions:
+    """``sess_options`` (or new defaults) with KleidiAI off, unless the caller set that entry."""
+    options = sess_options if sess_options is not None else SessionOptions()
+    try:
+        options.get_session_config_entry(DISABLE_KLEIDIAI)
+    except RuntimeError:  # not set
+        options.add_session_config_entry(DISABLE_KLEIDIAI, "1")
+    return options

tests/test_pretrained.py

@@ -0,0 +1,54 @@
+"""The shipped hey_livekit model on reference clips: inference end to end, with no torch.
+
+The clips are the ones the Rust and Swift ports test against. On onnxruntime 1.27.0 / arm64 the
+positive clip scored 0.004 until the KleidiAI kernels were turned off (livekit.wakeword.session).
+"""
+
+from __future__ import annotations
+
+import wave
+from pathlib import Path
+
+import numpy as np
+import pytest
+from onnxruntime import SessionOptions
+
+from livekit.wakeword import WakeWordModel
+from livekit.wakeword.session import DISABLE_KLEIDIAI, session_options
+
+MODEL = Path(__file__).resolve().parents[1] / "examples" / "resources" / "hey_livekit.onnx"
+FIXTURES = Path(__file__).parent / "fixtures"
+THRESHOLD = 0.5
+
+
+def clip(name: str) -> np.ndarray:
+    with wave.open(str(FIXTURES / name)) as w:
+        assert w.getframerate() == 16_000 and w.getnchannels() == 1 and w.getsampwidth() == 2
+        return np.frombuffer(w.readframes(w.getnframes()), dtype=np.int16)
+
+
+@pytest.fixture(scope="module")
+def model() -> WakeWordModel:
+    return WakeWordModel(models=[MODEL])
+
+
+def test_positive_clip_detected(model: WakeWordModel) -> None:
+    assert model.predict(clip("positive.wav"))["hey_livekit"] >= THRESHOLD
+
+
+def test_negative_clip_rejected(model: WakeWordModel) -> None:
+    assert model.predict(clip("negative.wav"))["hey_livekit"] < THRESHOLD
+
+
+def test_caller_options_are_kept() -> None:
+    options = SessionOptions()
+    options.intra_op_num_threads = 1
+    model = WakeWordModel(models=[MODEL], sess_options=options)
+    assert options.intra_op_num_threads == 1
+    assert model.predict(clip("positive.wav"))["hey_livekit"] >= THRESHOLD
+
+
+def test_caller_can_opt_back_in() -> None:
+    options = SessionOptions()
+    options.add_session_config_entry(DISABLE_KLEIDIAI, "0")
+    assert session_options(options).get_session_config_entry(DISABLE_KLEIDIAI) == "0"