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 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,91 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Wyoming ASR server for NVIDIA Parakeet via parakeet-mlx (Apple Silicon)."""
|
||||
import argparse
|
||||
import asyncio
|
||||
import logging
|
||||
|
||||
from wyoming.info import AsrModel, AsrProgram, Attribution, Info
|
||||
from wyoming.server import AsyncServer
|
||||
|
||||
from .engine import ParakeetEngine
|
||||
from .handler import ParakeetEventHandler
|
||||
|
||||
_LOGGER = logging.getLogger(__name__)
|
||||
__version__ = "1.0.0"
|
||||
|
||||
# v2 is English-only but does better inverse text normalisation than the
|
||||
# multilingual v3 ("21 degrees" / "30%" rather than "twenty-one degrees" /
|
||||
# "thirty percent"), which is what Home Assistant's local intent matching
|
||||
# expects. Don't switch to v3 without re-checking that.
|
||||
DEFAULT_MODEL = "mlx-community/parakeet-tdt-0.6b-v2"
|
||||
|
||||
|
||||
def build_parser() -> 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()
|
||||
@@ -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)
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user