From 4286e88344bc1b913de0a787f322b2de31c1853a Mon Sep 17 00:00:00 2001 From: Adam Harrison-Fuller Date: Wed, 29 Jul 2026 03:04:48 +0100 Subject: [PATCH] Wyoming speech-to-text server for Home Assistant using NVIDIA Parakeet Loads parakeet-mlx in-process (no HTTP hop) and serves it over the Wyoming protocol. On an M4 Mac mini this transcribes typical voice commands in ~110ms versus ~1150ms for a whisper.cpp large-v3 setup, with identical accuracy on a ten-command benchmark. Two behaviours matter beyond speed: silence returns an empty string rather than whisper's "Thank you." hallucination, and there is no decoder context carried between requests. Notable implementation details, all covered by mutation-checked regression tests: - MLX streams are thread-local, so the model is loaded and evaluated on a single dedicated worker thread. Splitting those raises "There is no Stream(cpu, 1) in current thread". - parakeet_mlx.load_audio() shells out to ffmpeg, which is unnecessary here since Wyoming delivers 16kHz mono PCM. The mel is built directly via get_logmel(), whose input must be float32 -- it views the complex STFT output as the input dtype, so anything narrower doubles the mel bin count. - Wyoming's run loop has no except clause, so an exception escaping handle_event closes the connection without sending a Transcript and Home Assistant waits indefinitely. Failures are caught and returned as an empty transcript instead. Defaults to parakeet-tdt-0.6b-v2 rather than the newer multilingual v3 because v2 emits digits ("21 degrees") where v3 spells numbers out, and Home Assistant's local intent matching expects digits. Co-Authored-By: Claude Opus 5 --- .gitignore | 6 + LICENSE | 21 ++++ README.md | 158 ++++++++++++++++++++++++ install.sh | 187 +++++++++++++++++++++++++++++ pytest.ini | 3 + requirements.txt | 2 + script/run | 2 + test/conftest.py | 12 ++ test/make-clips.sh | 40 +++++++ test/test_engine.py | 193 ++++++++++++++++++++++++++++++ test/test_handler.py | 226 +++++++++++++++++++++++++++++++++++ test/test_main.py | 45 +++++++ test/wy-test.py | 30 +++++ uninstall.sh | 24 ++++ wyoming_parakeet/__init__.py | 0 wyoming_parakeet/__main__.py | 91 ++++++++++++++ wyoming_parakeet/engine.py | 64 ++++++++++ wyoming_parakeet/handler.py | 64 ++++++++++ 18 files changed, 1168 insertions(+) create mode 100644 .gitignore create mode 100644 LICENSE create mode 100644 README.md create mode 100755 install.sh create mode 100644 pytest.ini create mode 100644 requirements.txt create mode 100755 script/run create mode 100644 test/conftest.py create mode 100755 test/make-clips.sh create mode 100644 test/test_engine.py create mode 100644 test/test_handler.py create mode 100644 test/test_main.py create mode 100644 test/wy-test.py create mode 100755 uninstall.sh create mode 100644 wyoming_parakeet/__init__.py create mode 100644 wyoming_parakeet/__main__.py create mode 100644 wyoming_parakeet/engine.py create mode 100644 wyoming_parakeet/handler.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..cfa63b4 --- /dev/null +++ b/.gitignore @@ -0,0 +1,6 @@ +.venv/ +__pycache__/ +*.pyc +.pytest_cache/ +# Generated by test/make-clips.sh +test/clips/ diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..1b2bfce --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 Adam Harrison-Fuller + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 index 0000000..563808d --- /dev/null +++ b/README.md @@ -0,0 +1,158 @@ +# wyoming-parakeet + +A [Wyoming protocol](https://github.com/rhasspy/wyoming) speech-to-text server +for Home Assistant, backed by NVIDIA's +[Parakeet TDT](https://huggingface.co/nvidia/parakeet-tdt-0.6b-v2) running on +Apple Silicon via [parakeet-mlx](https://github.com/senstella/parakeet-mlx). + +The model is loaded **in-process** — there is no HTTP hop between the Wyoming +bridge and inference. + +## Why + +Measured on an M4 Mac mini against ten typical Home Assistant voice commands, +replacing a whisper.cpp setup: + +| Backend | Mean latency | Correct | Silent input | +|---|---|---|---| +| whisper.cpp `large-v3` | ~1150 ms | 10/10 | `"Thank you."` | +| whisper.cpp `large-v3-turbo` | ~570 ms | 10/10 | `"Thank you."` | +| **parakeet-tdt-0.6b-v2** | **~110 ms** | **10/10** | `""` | + +Two things matter beyond raw speed: + +- **Silence returns an empty string.** Whisper hallucinates `"Thank you."` on + digital silence, which reaches your conversation agent as a real utterance. +- **No cross-request contamination.** whisper.cpp's server carries decoder + context between requests unless you pass `-nc`, and will return the + *previous* utterance — in testing, roughly one time in five. + +## Requirements + +- Apple Silicon Mac (MLX is Metal/ANE-backed) +- Python 3.10+ **with the `lzma` module** — `librosa` pulls in `pooch`, which + imports it. Pythons built without `xz` (a common pyenv default) pass every + version check and then fail at import time with + `ModuleNotFoundError: _lzma`. Homebrew's Python is fine. +- ~2.3 GB disk for the model, ~600 MB for MLX wheels + +`ffmpeg` is **not** required — Wyoming already delivers 16 kHz mono PCM, so +the mel spectrogram is built directly. + +## Install + +```bash +git clone https://github.com/adamhf/wyoming-parakeet-mlx +cd wyoming-parakeet-mlx +./install.sh +``` + +This creates a virtualenv, runs the unit tests, pre-downloads the model, and +registers a LaunchDaemon on port 7892 that starts at boot without needing a +GUI login. It installs *in place*, so keep the checkout somewhere permanent. + +Options: `--port`, `--model`, `--user`, `--python`, `--no-daemon`, +`--no-download`. + +Then in Home Assistant: **Settings → Devices & Services → Add Integration → +Wyoming Protocol**, enter the host and port, and select the new engine as the +speech-to-text step of your Assist pipeline. + +Remove it with `./uninstall.sh`. + +## Model choice: v2, not v3 + +The default is `parakeet-tdt-0.6b-v2` even though v3 is newer and +multilingual, because of inverse text normalisation: + +| Spoken | v2 | v3 | +|---|---|---| +| "twenty one degrees" | `21 degrees` | `twenty-one degrees` | +| "thirty percent" | `30%` | `thirty percent` | + +Home Assistant's local intent matching (hassil) expects digits. If your +pipeline has `prefer_local_intents` enabled, a model that spells numbers out +still *looks* accurate while quietly pushing commands off the fast local path +onto your LLM fallback. v3 is the better choice if you need languages other +than English — just be aware of the trade. + +## Updating the model + +`HF_HUB_OFFLINE=1` is set in the daemon, so it never silently re-downloads or +changes model at boot, and starts fine without a network. Updating is +therefore deliberate: + +```bash +HF_HUB_OFFLINE= .venv/bin/python -c \ + "from parakeet_mlx import from_pretrained; from_pretrained('mlx-community/parakeet-tdt-0.6b-v3')" +``` + +Then re-run `./install.sh --model mlx-community/parakeet-tdt-0.6b-v3`. The old +model stays cached, so reverting is just another `./install.sh`. + +## Tests + +```bash +.venv/bin/python -m pytest +``` + +37 unit tests, well under a second. The model is mocked throughout, so they +need no GPU, no network and no 2.3 GB download — they cover the audio +marshalling and threading around it, which is where the real bugs were. Every +regression test was mutation-checked: the fix reverted, the test confirmed to +fail. + +Worth knowing about a few: + +- `test_load_and_inference_share_one_thread` — MLX streams are thread-local, + so the model must be loaded *and* evaluated on the same thread or `mx.eval()` + raises `There is no Stream(cpu, 1) in current thread`. The test deliberately + *overlaps* its calls: sequential calls can coincidentally reuse one thread + out of a multi-worker pool and pass a broken implementation. +- `test_audio_is_float32_not_bfloat16` — `get_logmel` views the complex STFT + output as the input dtype, so anything narrower silently doubles the mel bin + count and the matmul fails. +- `test_model_failure_still_sends_a_transcript` — Wyoming's run loop is + `try/finally` with no `except`, so an exception escaping `handle_event` + closes the connection having sent nothing, and Home Assistant waits for a + response that never arrives. The handler catches and returns an empty + transcript so it fails fast instead. +- `test_concurrent_handlers_do_not_share_audio` — guards against + reintroducing the cross-request contamination described above. + +### End to end + +Unit tests never touch the real model, so after any model or library change: + +```bash +./test/make-clips.sh # generates via macOS TTS +.venv/bin/python test/wy-test.py test/clips/*.wav +``` + +Expect all ten commands correct, `silence.wav` empty, ~100–150 ms each once +warm. The first request or two after a restart run slower (~200–350 ms) while +Metal compiles its kernels. **Check number formatting, not just the words** — +`cmd2`, `cmd4` and `cmd7` are the ones that catch a model with weak ITN. + +## Updating the library + +```bash +.venv/bin/pip install -U parakeet-mlx mlx mlx-metal wyoming +``` + +Re-run both test suites afterwards. This project calls `get_logmel()` directly +rather than `load_audio()` (which shells out to ffmpeg), so it depends on two +parakeet-mlx internals rather than public API — the tests above are what tell +you if either moved. + +## Logs + +```bash +tail -f /tmp/local.wyoming-parakeet.stderr +``` + +Each request logs audio duration, inference time and the transcript. + +## License + +MIT diff --git a/install.sh b/install.sh new file mode 100755 index 0000000..8831c39 --- /dev/null +++ b/install.sh @@ -0,0 +1,187 @@ +#!/usr/bin/env bash +# Install wyoming-parakeet and (optionally) register it as a LaunchDaemon. +set -euo pipefail + +REPO_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +VENV="$REPO_DIR/.venv" +LABEL="local.wyoming-parakeet" +PLIST="/Library/LaunchDaemons/$LABEL.plist" + +PORT=7892 +MODEL="mlx-community/parakeet-tdt-0.6b-v2" +SERVICE_USER="$(id -un)" +PYTHON="" +INSTALL_DAEMON=1 +DOWNLOAD_MODEL=1 + +usage() { + cat <&2; usage >&2; exit 2 ;; + esac +done + +die() { echo "error: $*" >&2; exit 1; } + +# --- preflight ------------------------------------------------------------- + +[[ "$(uname -s)" == "Darwin" ]] || die "macOS only (MLX is Apple Silicon)." +[[ "$(uname -m)" == "arm64" ]] || die "Apple Silicon required; this is $(uname -m)." +id -u "$SERVICE_USER" >/dev/null 2>&1 || die "no such user: $SERVICE_USER" + +# librosa pulls in pooch, which imports lzma. Pythons built without xz (a +# common pyenv default) satisfy every version check and then fail at import +# time with ModuleNotFoundError: _lzma -- so check for it up front. +usable_python() { + local py resolved + py="$1" + resolved="$(command -v "$py" 2>/dev/null)" || return 1 + [[ -x "$resolved" ]] || return 1 + "$resolved" -c 'import sys, lzma; sys.exit(0 if sys.version_info >= (3, 10) else 1)' \ + >/dev/null 2>&1 || return 1 + echo "$resolved" +} + +if [[ -n "$PYTHON" ]]; then + PYTHON="$(usable_python "$PYTHON")" \ + || die "$PYTHON is unusable: needs >=3.10 and the lzma module." +else + # Search PATH *and* the usual Homebrew prefixes explicitly. A + # non-interactive shell often has neither on PATH, which would otherwise + # leave us falling back to the system python3 (3.9, too old). + for candidate in \ + python3.14 python3.13 python3.12 python3.11 python3 \ + /opt/homebrew/bin/python3.1{4,3,2,1} /opt/homebrew/bin/python3 \ + /usr/local/bin/python3.1{4,3,2,1} /usr/local/bin/python3 + do + if PYTHON="$(usable_python "$candidate")"; then + break + fi + PYTHON="" + done + [[ -n "$PYTHON" ]] || die "no suitable python found (needs >=3.10 with the lzma module; try: brew install python@3.13)" +fi + +echo "==> Python: $PYTHON ($("$PYTHON" -V 2>&1))" +echo "==> Install: $REPO_DIR" +echo "==> Model: $MODEL" +echo "==> Port: $PORT" + +# --- venv ------------------------------------------------------------------ + +echo "==> Creating virtualenv" +"$PYTHON" -m venv "$VENV" +"$VENV/bin/pip" install --quiet --upgrade pip +echo "==> Installing dependencies (this pulls ~600MB of MLX wheels)" +"$VENV/bin/pip" install --quiet -r "$REPO_DIR/requirements.txt" + +echo "==> Running unit tests" +"$VENV/bin/pip" install --quiet pytest pytest-asyncio +( cd "$REPO_DIR" && "$VENV/bin/python" -m pytest -q ) + +if [[ "$DOWNLOAD_MODEL" -eq 1 ]]; then + echo "==> Downloading $MODEL (~2.3GB on first run)" + HF_HUB_OFFLINE= "$VENV/bin/python" - "$MODEL" <<'EOF' +import sys +from parakeet_mlx import from_pretrained +from_pretrained(sys.argv[1]) +EOF +fi + +if [[ "$INSTALL_DAEMON" -eq 0 ]]; then + echo + echo "Done (no daemon installed). Run it with:" + echo " $REPO_DIR/script/run --uri tcp://0.0.0.0:$PORT --model $MODEL" + exit 0 +fi + +# --- launchd --------------------------------------------------------------- + +echo "==> Installing LaunchDaemon (needs sudo)" +TMP_PLIST="$(mktemp -t wyoming-parakeet)" +trap 'rm -f "$TMP_PLIST"' EXIT +cat > "$TMP_PLIST" < + + + + Label + $LABEL + ProgramArguments + + $REPO_DIR/script/run + --uri + tcp://0.0.0.0:$PORT + --model + $MODEL + + EnvironmentVariables + + HF_HUB_OFFLINE1 + HOME$(eval echo "~$SERVICE_USER") + + UserName$SERVICE_USER + GroupNamestaff + InitGroups + WorkingDirectory$REPO_DIR + RunAtLoad + KeepAlive + ProcessTypeInteractive + LowPriorityIO + StandardOutPath/tmp/$LABEL.stdout + StandardErrorPath/tmp/$LABEL.stderr + + +EOF + +plutil -lint "$TMP_PLIST" >/dev/null || die "generated plist is malformed" +sudo cp "$TMP_PLIST" "$PLIST" +sudo chown root:wheel "$PLIST" +sudo chmod 644 "$PLIST" + +sudo launchctl bootout "system/$LABEL" 2>/dev/null || true +sudo launchctl bootstrap system "$PLIST" + +echo "==> Waiting for the service to come up" +for _ in $(seq 1 60); do + if nc -z 127.0.0.1 "$PORT" 2>/dev/null; then + echo " listening on $PORT" + break + fi + sleep 2 +done + +nc -z 127.0.0.1 "$PORT" 2>/dev/null || die "service did not start; check /tmp/$LABEL.stderr" + +cat < Devices & Services -> +Add Integration -> Wyoming Protocol, using this host and port $PORT, +then select the new engine as the speech-to-text step in your Assist pipeline. + + logs: tail -f /tmp/$LABEL.stderr + verify: $VENV/bin/python test/wy-test.py test/clips/*.wav + (run test/make-clips.sh first to generate the clips) +EOF diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..894ffc6 --- /dev/null +++ b/pytest.ini @@ -0,0 +1,3 @@ +[pytest] +testpaths = test +asyncio_mode = auto diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..b7da79f --- /dev/null +++ b/requirements.txt @@ -0,0 +1,2 @@ +parakeet-mlx>=0.5.2 +wyoming>=1.10.0 diff --git a/script/run b/script/run new file mode 100755 index 0000000..7e95d69 --- /dev/null +++ b/script/run @@ -0,0 +1,2 @@ +#!/usr/bin/env bash +exec "$(dirname "$0")/../.venv/bin/python" -m wyoming_parakeet "$@" diff --git a/test/conftest.py b/test/conftest.py new file mode 100644 index 0000000..c4f4d6a --- /dev/null +++ b/test/conftest.py @@ -0,0 +1,12 @@ +import numpy as np +import pytest + + +@pytest.fixture +def pcm(): + """Build int16 mono PCM bytes of a given sample count.""" + + def _pcm(n_samples: int, value: int = 8192) -> bytes: + return np.full(n_samples, value, dtype=np.int16).tobytes() + + return _pcm diff --git a/test/make-clips.sh b/test/make-clips.sh new file mode 100755 index 0000000..52430f7 --- /dev/null +++ b/test/make-clips.sh @@ -0,0 +1,40 @@ +#!/usr/bin/env bash +# Generate the end-to-end test clips using macOS TTS. They are not committed +# because they are trivially reproducible binaries. +set -euo pipefail + +DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/clips" +mkdir -p "$DIR" + +PHRASES=( + "turn off the kitchen lights" + "set the living room thermostat to twenty one degrees" + "what is the temperature in the bedroom" + "dim the hallway lights to thirty percent" + "is the back door locked" + "turn on the christmas tree in the conservatory" + "set a timer for twelve minutes" + "play radio six music in the kitchen" + "whats the octopus agile rate right now" + "close the blinds in the study and turn on the desk lamp" +) + +i=1 +for phrase in "${PHRASES[@]}"; do + say -o "$DIR/cmd$i.aiff" "$phrase" + afconvert -f WAVE -d LEI16@16000 -c 1 "$DIR/cmd$i.aiff" "$DIR/cmd$i.wav" + rm -f "$DIR/cmd$i.aiff" + i=$((i + 1)) +done + +# 3s of digital silence -- a good model returns an empty string here rather +# than hallucinating "Thank you." the way whisper does. +python3 - "$DIR/silence.wav" <<'PY' +import sys, wave +w = wave.open(sys.argv[1], "wb") +w.setparams((1, 2, 16000, 0, "NONE", "NONE")) +w.writeframes(b"\x00\x00" * 16000 * 3) +w.close() +PY + +echo "Wrote $(ls "$DIR"/*.wav | wc -l | tr -d ' ') clips to $DIR" diff --git a/test/test_engine.py b/test/test_engine.py new file mode 100644 index 0000000..a68b790 --- /dev/null +++ b/test/test_engine.py @@ -0,0 +1,193 @@ +"""Tests for ParakeetEngine. + +The model itself is mocked throughout -- these cover the audio marshalling +and threading around it, which is where the real bugs were. +""" +import threading +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest + +from wyoming_parakeet.engine import MIN_SAMPLES, SAMPLE_RATE, ParakeetEngine + + +def make_engine(text="turn off the kitchen lights"): + engine = ParakeetEngine("fake/model") + result = MagicMock() + result.text = text + engine.model = MagicMock() + engine.model.generate.return_value = [result] + return engine + + +def test_returns_model_text(pcm): + engine = make_engine("set a timer for 12 minutes") + with patch("wyoming_parakeet.engine.get_logmel"): + assert engine._transcribe(pcm(SAMPLE_RATE)) == "set a timer for 12 minutes" + + +def test_short_audio_short_circuits_without_touching_model(pcm): + """A clipped VAD flush must not reach the model at all.""" + engine = make_engine() + with patch("wyoming_parakeet.engine.get_logmel") as get_logmel: + assert engine._transcribe(pcm(MIN_SAMPLES - 1)) == "" + engine.model.generate.assert_not_called() + get_logmel.assert_not_called() + + +def test_audio_at_threshold_is_processed(pcm): + engine = make_engine() + with patch("wyoming_parakeet.engine.get_logmel"): + assert engine._transcribe(pcm(MIN_SAMPLES)) != "" + engine.model.generate.assert_called_once() + + +def test_empty_audio_returns_empty_string(): + engine = make_engine() + with patch("wyoming_parakeet.engine.get_logmel"): + assert engine._transcribe(b"") == "" + + +def test_no_results_returns_empty_string(pcm): + """Silence legitimately decodes to nothing; don't IndexError on it.""" + engine = make_engine() + engine.model.generate.return_value = [] + with patch("wyoming_parakeet.engine.get_logmel"): + assert engine._transcribe(pcm(SAMPLE_RATE)) == "" + + +def test_audio_is_float32_not_bfloat16(pcm): + """Regression: get_logmel views the complex STFT output as the input + dtype, so anything narrower than float32 silently doubles the mel bin + count and the downstream matmul fails.""" + engine = make_engine() + with patch("wyoming_parakeet.engine.get_logmel") as get_logmel: + engine._transcribe(pcm(SAMPLE_RATE)) + samples = get_logmel.call_args[0][0] + assert np.asarray(samples).dtype == np.float32 + + +def test_pcm_is_scaled_to_unit_range(pcm): + engine = make_engine() + with patch("wyoming_parakeet.engine.get_logmel") as get_logmel: + engine._transcribe(pcm(SAMPLE_RATE, value=16384)) + samples = np.asarray(get_logmel.call_args[0][0]) + assert samples.shape == (SAMPLE_RATE,) + assert np.allclose(samples, 0.5) + + +def test_full_scale_pcm_stays_within_unit_range(pcm): + engine = make_engine() + with patch("wyoming_parakeet.engine.get_logmel") as get_logmel: + engine._transcribe(pcm(SAMPLE_RATE, value=-32768)) + samples = np.asarray(get_logmel.call_args[0][0]) + assert np.abs(samples).max() <= 1.0 + + +def test_preprocessor_config_is_passed_through(pcm): + engine = make_engine() + with patch("wyoming_parakeet.engine.get_logmel") as get_logmel: + engine._transcribe(pcm(SAMPLE_RATE)) + assert get_logmel.call_args[0][1] is engine.model.preprocessor_config + + +@pytest.mark.asyncio +async def test_load_and_inference_share_one_thread(): + """Regression: MLX streams are thread-local, so a model loaded on one + thread cannot be evaluated from another -- mx.eval() raises + 'There is no Stream(cpu, 1) in current thread'.""" + import asyncio + + engine = ParakeetEngine("fake/model") + threads = [] + lock = threading.Lock() + + def record(): + with lock: + threads.append(threading.get_ident()) + # Hold the worker. Sequential calls can coincidentally reuse a single + # thread out of a multi-worker pool, so overlap them -- a pool wider + # than one will hand these to different threads and fail the assert. + threading.Event().wait(0.05) + + engine._load = record + engine._transcribe = lambda _pcm: (record(), "")[1] + + await engine.start() + await asyncio.gather(*(engine.transcribe(b"") for _ in range(4))) + + assert len(threads) == 5 + assert len(set(threads)) == 1, "load and inference must share one thread" + assert threads[0] != threading.get_ident(), "must not run on the event loop" + + +@pytest.mark.asyncio +async def test_requests_are_serialised(pcm): + """A single worker means overlapping satellites queue rather than + racing the ANE.""" + import asyncio + + engine = ParakeetEngine("fake/model") + concurrent = 0 + peak = 0 + + def slow(_pcm): + nonlocal concurrent, peak + concurrent += 1 + peak = max(peak, concurrent) + threading.Event().wait(0.05) + concurrent -= 1 + return "ok" + + engine._transcribe = slow + await asyncio.gather(*(engine.transcribe(pcm(SAMPLE_RATE)) for _ in range(4))) + assert peak == 1 + + +def test_transcription_carries_no_state_between_calls(pcm): + """Regression guard for the bug class that motivated leaving whisper.cpp: + its server reused decoder context across requests and would return the + previous utterance. Each call here must stand alone.""" + engine = make_engine() + audio = pcm(SAMPLE_RATE) + + with patch("wyoming_parakeet.engine.get_logmel") as get_logmel: + engine._transcribe(pcm(SAMPLE_RATE * 2, value=4096)) + engine._transcribe(audio) + first = np.asarray(get_logmel.call_args[0][0]) + + engine._transcribe(pcm(SAMPLE_RATE // 2, value=-2048)) + engine._transcribe(audio) + second = np.asarray(get_logmel.call_args[0][0]) + + assert np.array_equal(first, second) + # generate() must be handed only the current mel, with no prompt/context. + assert engine.model.generate.call_args[0][0] is get_logmel.return_value + assert engine.model.generate.call_args.kwargs == {} + + +@pytest.mark.asyncio +async def test_start_warms_the_model(): + """The first inference JITs Metal kernels. If warm-up is dropped, the + first voice command after every reboot pays that cost.""" + engine = ParakeetEngine("fake/model") + model = MagicMock() + model.generate.return_value = [] + + with patch("parakeet_mlx.from_pretrained", return_value=model) as load: + with patch("wyoming_parakeet.engine.get_logmel"): + await engine.start() + + load.assert_called_once_with("fake/model") + assert model.generate.call_count == 1, "start() should run one warm-up pass" + + +@pytest.mark.asyncio +async def test_engine_survives_a_failed_request(pcm): + """One bad utterance must not poison the worker for later ones.""" + engine = make_engine("recovered") + with patch("wyoming_parakeet.engine.get_logmel", side_effect=[RuntimeError("boom"), MagicMock()]): + with pytest.raises(RuntimeError): + await engine.transcribe(pcm(SAMPLE_RATE)) + assert await engine.transcribe(pcm(SAMPLE_RATE)) == "recovered" diff --git a/test/test_handler.py b/test/test_handler.py new file mode 100644 index 0000000..56c8e1a --- /dev/null +++ b/test/test_handler.py @@ -0,0 +1,226 @@ +"""Tests for the Wyoming event handling.""" +from unittest.mock import AsyncMock, MagicMock + +import pytest +from wyoming.asr import Transcribe, Transcript +from wyoming.audio import AudioChunk, AudioStop +from wyoming.info import AsrModel, AsrProgram, Attribution, Describe, Info + +from wyoming_parakeet.engine import SAMPLE_RATE +from wyoming_parakeet.handler import ParakeetEventHandler + +INFO = Info( + asr=[ + AsrProgram( + name="parakeet-mlx", + description="test", + attribution=Attribution(name="t", url="http://example.invalid"), + installed=True, + version="1.0.0", + models=[ + AsrModel( + name="fake/model", + description="test", + attribution=Attribution(name="t", url="http://example.invalid"), + installed=True, + version=None, + languages=["en"], + ) + ], + ) + ] +) + + +@pytest.fixture +def handler(): + engine = MagicMock() + engine.transcribe = AsyncMock(return_value="turn off the kitchen lights") + h = ParakeetEventHandler(INFO, MagicMock(), engine, MagicMock(), MagicMock()) + h.write_event = AsyncMock() + return h + + +def chunk(n_samples, rate=SAMPLE_RATE): + return AudioChunk( + rate=rate, width=2, channels=1, audio=b"\x01\x00" * n_samples + ).event() + + +async def test_describe_returns_info(handler): + assert await handler.handle_event(Describe().event()) is True + handler.write_event.assert_awaited_once() + assert Info.is_type(handler.write_event.await_args[0][0].type) + + +async def test_transcribe_event_is_accepted(handler): + assert await handler.handle_event(Transcribe(language="en").event()) is True + handler.write_event.assert_not_awaited() + + +async def test_unknown_event_keeps_connection_open(handler): + from wyoming.event import Event + + assert await handler.handle_event(Event(type="something-else")) is True + + +async def test_chunks_accumulate_before_stop(handler): + for _ in range(3): + assert await handler.handle_event(chunk(100)) is True + handler.engine.transcribe.assert_not_awaited() + assert len(handler.audio) == 3 * 100 * 2 + + +async def test_stop_transcribes_accumulated_audio(handler): + await handler.handle_event(chunk(SAMPLE_RATE // 2)) + await handler.handle_event(chunk(SAMPLE_RATE // 2)) + + assert await handler.handle_event(AudioStop().event()) is False + + handler.engine.transcribe.assert_awaited_once() + assert len(handler.engine.transcribe.await_args[0][0]) == SAMPLE_RATE * 2 + + +async def test_stop_writes_transcript(handler): + await handler.handle_event(chunk(SAMPLE_RATE)) + await handler.handle_event(AudioStop().event()) + + event = handler.write_event.await_args[0][0] + assert Transcript.is_type(event.type) + assert Transcript.from_event(event).text == "turn off the kitchen lights" + + +async def test_empty_transcript_is_still_sent(handler): + """Silence must produce an empty Transcript, not a dropped response -- + Home Assistant waits for one.""" + handler.engine.transcribe = AsyncMock(return_value="") + await handler.handle_event(chunk(SAMPLE_RATE)) + await handler.handle_event(AudioStop().event()) + + event = handler.write_event.await_args[0][0] + assert Transcript.is_type(event.type) + assert Transcript.from_event(event).text == "" + + +async def test_buffer_resets_after_stop(handler): + await handler.handle_event(chunk(SAMPLE_RATE)) + await handler.handle_event(AudioStop().event()) + assert handler.audio == b"" + + +async def test_stop_with_no_audio_does_not_crash(handler): + assert await handler.handle_event(AudioStop().event()) is False + handler.engine.transcribe.assert_awaited_once_with(b"") + + +async def test_resampled_input_is_converted_to_16k(handler): + """Satellites may send other rates; the converter must normalise them + before the engine sees them.""" + await handler.handle_event(chunk(48000, rate=48000)) + await handler.handle_event(AudioStop().event()) + + pcm = handler.engine.transcribe.await_args[0][0] + assert len(pcm) == SAMPLE_RATE * 2, "1s of 48kHz audio should become 1s at 16kHz" + + +# --- failure handling ------------------------------------------------------- +# wyoming's run loop is try/finally with no except, so an exception escaping +# handle_event closes the connection having sent nothing and Home Assistant +# waits for a response that never arrives. + + +async def test_model_failure_still_sends_a_transcript(handler): + handler.engine.transcribe = AsyncMock(side_effect=RuntimeError("metal exploded")) + + assert await handler.handle_event(AudioStop().event()) is False + + event = handler.write_event.await_args[0][0] + assert Transcript.is_type(event.type) + assert Transcript.from_event(event).text == "" + + +async def test_model_failure_does_not_propagate(handler): + """Must not escape into wyoming's run loop.""" + handler.engine.transcribe = AsyncMock(side_effect=RuntimeError("metal exploded")) + await handler.handle_event(chunk(SAMPLE_RATE)) + await handler.handle_event(AudioStop().event()) # would raise if unhandled + + +async def test_model_failure_is_logged_with_traceback(handler, caplog): + handler.engine.transcribe = AsyncMock(side_effect=RuntimeError("metal exploded")) + with caplog.at_level("ERROR"): + await handler.handle_event(AudioStop().event()) + assert "metal exploded" in caplog.text + + +async def test_buffer_resets_after_failure(handler): + handler.engine.transcribe = AsyncMock(side_effect=RuntimeError("metal exploded")) + await handler.handle_event(chunk(SAMPLE_RATE)) + await handler.handle_event(AudioStop().event()) + assert handler.audio == b"" + + +# --- isolation -------------------------------------------------------------- +# The whisper.cpp server this replaced leaked decoder context between requests +# and would return the *previous* utterance. Guard against reintroducing any +# shared per-request state. + + +async def test_concurrent_handlers_do_not_share_audio(): + engine = MagicMock() + engine.transcribe = AsyncMock(return_value="") + seen = [] + engine.transcribe.side_effect = lambda pcm: seen.append(pcm) or "" + + a = ParakeetEventHandler(INFO, MagicMock(), engine, MagicMock(), MagicMock()) + b = ParakeetEventHandler(INFO, MagicMock(), engine, MagicMock(), MagicMock()) + a.write_event = AsyncMock() + b.write_event = AsyncMock() + + # Interleave two conversations through one shared engine. + await a.handle_event(chunk(100)) + await b.handle_event(chunk(300)) + await a.handle_event(chunk(100)) + await a.handle_event(AudioStop().event()) + await b.handle_event(AudioStop().event()) + + assert [len(p) for p in seen] == [200 * 2, 300 * 2] + + +async def test_transcript_reflects_only_this_requests_audio(handler): + """A second utterance must not inherit the first one's text.""" + handler.engine.transcribe = AsyncMock(side_effect=["first", "second"]) + + await handler.handle_event(chunk(SAMPLE_RATE)) + await handler.handle_event(AudioStop().event()) + first = Transcript.from_event(handler.write_event.await_args[0][0]).text + + await handler.handle_event(chunk(SAMPLE_RATE)) + await handler.handle_event(AudioStop().event()) + second = Transcript.from_event(handler.write_event.await_args[0][0]).text + + assert (first, second) == ("first", "second") + + +# --- audio format normalisation --------------------------------------------- +# Satellites vary; the engine must always receive 16 kHz mono 16-bit. + + +async def test_stereo_input_is_downmixed(handler): + stereo = AudioChunk( + rate=SAMPLE_RATE, width=2, channels=2, audio=b"\x01\x00\x01\x00" * SAMPLE_RATE + ).event() + await handler.handle_event(stereo) + await handler.handle_event(AudioStop().event()) + + assert len(handler.engine.transcribe.await_args[0][0]) == SAMPLE_RATE * 2 + + +async def test_8bit_input_is_widened(handler): + narrow = AudioChunk( + rate=SAMPLE_RATE, width=1, channels=1, audio=b"\x40" * SAMPLE_RATE + ).event() + await handler.handle_event(narrow) + await handler.handle_event(AudioStop().event()) + + assert len(handler.engine.transcribe.await_args[0][0]) == SAMPLE_RATE * 2 diff --git a/test/test_main.py b/test/test_main.py new file mode 100644 index 0000000..cf0862b --- /dev/null +++ b/test/test_main.py @@ -0,0 +1,45 @@ +"""Tests for CLI wiring and the Info advertised to Home Assistant.""" +from wyoming.info import Info + +from wyoming_parakeet.__main__ import DEFAULT_MODEL, build_info, build_parser + + +def test_default_model_is_v2_not_v3(): + """Deliberate: v3 is newer and multilingual but spells numbers out + ("twenty-one degrees" vs "21 degrees"). Home Assistant's local intent + matching wants digits, and both pipelines run prefer_local_intents, so v3 + would still look accurate while pushing commands onto the LLM fallback. + If you are changing this, re-run test/wy-test.py and check cmd2/cmd4/cmd7.""" + assert DEFAULT_MODEL == "mlx-community/parakeet-tdt-0.6b-v2" + + +def test_model_defaults_are_applied(): + args = build_parser().parse_args(["--uri", "tcp://0.0.0.0:7892"]) + assert args.model == DEFAULT_MODEL + assert args.language == "en" + assert args.debug is False + + +def test_model_can_be_overridden(): + args = build_parser().parse_args( + ["--uri", "tcp://0.0.0.0:7892", "--model", "mlx-community/other"] + ) + assert args.model == "mlx-community/other" + + +def test_info_advertises_the_running_model(): + """Home Assistant's Wyoming config flow reads this; if it is malformed the + integration cannot be added at all.""" + info = build_info("mlx-community/other", "en") + assert Info.is_type(info.event().type) + + program = info.asr[0] + assert program.installed is True + assert program.models[0].name == "mlx-community/other" + assert program.models[0].languages == ["en"] + + +def test_info_round_trips_through_an_event(): + info = build_info(DEFAULT_MODEL, "en") + restored = Info.from_event(info.event()) + assert restored.asr[0].models[0].name == DEFAULT_MODEL diff --git a/test/wy-test.py b/test/wy-test.py new file mode 100644 index 0000000..264724c --- /dev/null +++ b/test/wy-test.py @@ -0,0 +1,30 @@ +import asyncio, os, sys, time, wave +from wyoming.client import AsyncTcpClient +from wyoming.asr import Transcribe, Transcript +from wyoming.audio import AudioChunk, AudioStart, AudioStop +from wyoming.info import Describe, Info + +async def once(path): + w = wave.open(path,"rb"); pcm = w.readframes(w.getnframes()); rate=w.getframerate() + async with AsyncTcpClient("127.0.0.1", int(os.environ.get("PORT", 7892))) as c: + await c.write_event(Transcribe(language="en").event()) + await c.write_event(AudioStart(rate=rate, width=2, channels=1).event()) + t=time.time() + for i in range(0, len(pcm), 4096): + await c.write_event(AudioChunk(rate=rate,width=2,channels=1,audio=pcm[i:i+4096]).event()) + await c.write_event(AudioStop().event()) + while True: + e = await c.read_event() + if e is None: return None, 0 + if Transcript.is_type(e.type): + return Transcript.from_event(e).text, (time.time()-t)*1000 + +async def main(): + async with AsyncTcpClient("127.0.0.1", int(os.environ.get("PORT", 7892))) as c: + await c.write_event(Describe().event()) + e = await c.read_event() + print("INFO ok:", Info.is_type(e.type)) + for p in sys.argv[1:]: + txt, ms = await once(p) + print(f"{p.split('/')[-1]:14s} {ms:7.0f}ms :: {txt!r}") +asyncio.run(main()) diff --git a/uninstall.sh b/uninstall.sh new file mode 100755 index 0000000..85f5232 --- /dev/null +++ b/uninstall.sh @@ -0,0 +1,24 @@ +#!/usr/bin/env bash +# Remove the wyoming-parakeet LaunchDaemon. Leaves the checkout and the +# downloaded model in ~/.cache/huggingface alone. +set -euo pipefail + +LABEL="local.wyoming-parakeet" +PLIST="/Library/LaunchDaemons/$LABEL.plist" + +echo "==> Stopping $LABEL" +sudo launchctl bootout "system/$LABEL" 2>/dev/null || echo " (not running)" + +if [[ -f "$PLIST" ]]; then + echo "==> Removing $PLIST" + sudo rm -f "$PLIST" +fi + +cat < argparse.ArgumentParser: + parser = argparse.ArgumentParser() + parser.add_argument("--uri", required=True, help="unix:// or tcp://") + parser.add_argument("--model", default=DEFAULT_MODEL, help="HuggingFace model id") + parser.add_argument("--language", default="en", help="Language code reported to HA") + parser.add_argument("--debug", action="store_true", help="Log DEBUG messages") + parser.add_argument("--log-format", default=logging.BASIC_FORMAT) + parser.add_argument("--version", action="version", version=__version__) + return parser + + +def build_info(model: str, language: str) -> Info: + """Describe this service to Home Assistant's Wyoming config flow.""" + return Info( + asr=[ + AsrProgram( + name="parakeet-mlx", + description="NVIDIA Parakeet TDT via MLX", + attribution=Attribution( + name="senstella", url="https://github.com/senstella/parakeet-mlx" + ), + installed=True, + version=__version__, + models=[ + AsrModel( + name=model, + description=model, + attribution=Attribution( + name="NVIDIA", + url="https://huggingface.co/nvidia/parakeet-tdt-0.6b-v2", + ), + installed=True, + version=None, + languages=[language], + ) + ], + ) + ] + ) + + +async def main() -> None: + args = build_parser().parse_args() + + logging.basicConfig( + level=logging.DEBUG if args.debug else logging.INFO, format=args.log_format + ) + + _LOGGER.info("Loading %s", args.model) + engine = ParakeetEngine(args.model) + await engine.start() + + wyoming_info = build_info(args.model, args.language) + server = AsyncServer.from_uri(args.uri) + _LOGGER.info("Ready on %s", args.uri) + await server.run( + lambda *a, **kw: ParakeetEventHandler(wyoming_info, args, engine, *a, **kw) + ) + + +def run() -> None: + try: + asyncio.run(main()) + except KeyboardInterrupt: + pass + + +if __name__ == "__main__": + run() diff --git a/wyoming_parakeet/engine.py b/wyoming_parakeet/engine.py new file mode 100644 index 0000000..e00cf89 --- /dev/null +++ b/wyoming_parakeet/engine.py @@ -0,0 +1,64 @@ +"""Parakeet model wrapper pinned to a single MLX worker thread.""" +import asyncio +import logging +import time +from concurrent.futures import ThreadPoolExecutor + +import mlx.core as mx +import numpy as np +from parakeet_mlx.audio import get_logmel + +_LOGGER = logging.getLogger(__name__) + +SAMPLE_RATE = 16000 +# Shorter than this and the encoder has nothing useful to chew on; Wyoming +# clients occasionally flush a near-empty buffer when VAD clips too tightly. +MIN_SAMPLES = SAMPLE_RATE // 10 + + +class ParakeetEngine: + """Owns the model and guarantees every MLX call happens on one thread. + + MLX streams are thread-local, so a model loaded on the main thread cannot + be evaluated from an arbitrary executor thread -- mx.eval() raises + "There is no Stream(cpu, 1) in current thread". Loading and inference both + run on this single worker, which also serialises requests for free. + """ + + def __init__(self, model_name: str): + self.model_name = model_name + self.model = None + self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="mlx") + + def _run(self, fn, *args): + return asyncio.get_running_loop().run_in_executor(self._executor, fn, *args) + + def _load(self) -> None: + from parakeet_mlx import from_pretrained + + started = time.monotonic() + self.model = from_pretrained(self.model_name) + _LOGGER.info("Loaded %s in %.1fs", self.model_name, time.monotonic() - started) + # First inference JITs Metal kernels; pay that now, not on the user's + # first voice command. + started = time.monotonic() + self._transcribe(b"\x00\x00" * SAMPLE_RATE) + _LOGGER.info("Warmed up in %.1fs", time.monotonic() - started) + + def _transcribe(self, pcm: bytes) -> str: + samples = np.frombuffer(pcm, dtype=np.int16).astype(np.float32) / 32768.0 + if samples.size < MIN_SAMPLES: + return "" + # parakeet_mlx.load_audio() shells out to ffmpeg, which we neither have + # nor need: Wyoming already hands us 16 kHz mono PCM, so build the mel + # directly. float32 is required -- get_logmel views the complex STFT + # output as the input dtype, so bfloat16 silently doubles the bin count. + mel = get_logmel(mx.array(samples), self.model.preprocessor_config) + results = self.model.generate(mel) + return results[0].text if results else "" + + async def start(self) -> None: + await self._run(self._load) + + async def transcribe(self, pcm: bytes) -> str: + return await self._run(self._transcribe, pcm) diff --git a/wyoming_parakeet/handler.py b/wyoming_parakeet/handler.py new file mode 100644 index 0000000..7d5e445 --- /dev/null +++ b/wyoming_parakeet/handler.py @@ -0,0 +1,64 @@ +"""Wyoming event handler backed by an in-process parakeet-mlx model.""" +import logging +import time + +from wyoming.asr import Transcribe, Transcript +from wyoming.audio import AudioChunk, AudioChunkConverter, AudioStop +from wyoming.event import Event +from wyoming.info import Describe, Info +from wyoming.server import AsyncEventHandler + +from .engine import SAMPLE_RATE + +_LOGGER = logging.getLogger(__name__) + + +class ParakeetEventHandler(AsyncEventHandler): + def __init__(self, wyoming_info: Info, cli_args, engine, *args, **kwargs): + super().__init__(*args, **kwargs) + self.cli_args = cli_args + self.wyoming_info_event = wyoming_info.event() + self.engine = engine + self.audio = bytes() + self.converter = AudioChunkConverter(rate=SAMPLE_RATE, width=2, channels=1) + + async def handle_event(self, event: Event) -> bool: + if AudioChunk.is_type(event.type): + if not self.audio: + _LOGGER.debug("Receiving audio") + self.audio += self.converter.convert(AudioChunk.from_event(event)).audio + return True + + if AudioStop.is_type(event.type): + duration = len(self.audio) / (SAMPLE_RATE * 2) + started = time.monotonic() + try: + text = await self.engine.transcribe(self.audio) + _LOGGER.info( + "%.2fs audio -> %.0fms :: %r", + duration, + (time.monotonic() - started) * 1000, + text, + ) + except Exception: + # wyoming's run loop has no except clause, so letting this + # propagate closes the connection without ever sending a + # Transcript -- Home Assistant then waits for a response that + # will never arrive. Fail fast with an empty result instead. + _LOGGER.exception("Transcription failed after %.2fs audio", duration) + text = "" + finally: + self.audio = bytes() + + await self.write_event(Transcript(text=text).event()) + return False + + if Transcribe.is_type(event.type): + return True + + if Describe.is_type(event.type): + await self.write_event(self.wyoming_info_event) + _LOGGER.debug("Sent info") + return True + + return True