diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..79d631c --- /dev/null +++ b/.env.example @@ -0,0 +1,26 @@ +# PepperWizard project env file. +# +# Copy to `.env` at the project root: +# cp .env.example .env +# +# Robot connection. +# Physical robot: set NAOQI_IP to the robot's IP and NAOQI_PORT=9559. +# Simulation: the dev overlay (docker-compose.dev.yml) runs a qiBullet sim +# inside pepper-box:latest. Set NAOQI_IP=127.0.0.1 to trigger sim mode. +NAOQI_IP=192.168.123.50 +NAOQI_PORT=9559 + +# For LLM talk mode, export ANTHROPIC_API_KEY in your shell before running or add to the .env. + +# `docker compose ...`. Uncomment to load the dev overlay by default. +# COMPOSE_FILE=docker-compose.yml:docker-compose.dev.yml +# COMPOSE_PROFILES=gpu + +# GPU exposure for stt-service. +# Defaults (both lines unset/commented) → CPU + Whisper. +# To run Parakeet on an NVIDIA host: uncomment both lines below. +# Requires: +# - nvidia-container-toolkit installed on the host +# - the `nvidia` Docker runtime registered (check with `docker info | grep -i runtime`) +# STT_DOCKER_RUNTIME=nvidia +# STT_GPU_DEVICES=all diff --git a/.gitignore b/.gitignore index d969dfb..b7c6a14 100644 --- a/.gitignore +++ b/.gitignore @@ -12,7 +12,6 @@ TODO.txt # Environment and Venv .venv_test/ -robot.env .env # Python Packaging diff --git a/README.md b/README.md index f9918dd..ac4a0d9 100644 --- a/README.md +++ b/README.md @@ -55,18 +55,20 @@ This fetches the SDK from Aldebaran's CDN (with a Wayback Machine fallback), ver ### 2. Point PepperWizard at your robot -The connection is configured via a `robot.env` file. Copy the example and edit if your robot isn't at the lab default: +The connection is configured via a `.env` file at the repo root. Copy the example and edit if your robot hasa different IP: ```bash -cp robot.env.example robot.env +cp .env.example .env ``` ```bash -# robot.env +# .env NAOQI_IP=192.168.123.50 # robot IP (use 127.0.0.1 for a local NAOqi sim) -NAOQI_PORT=9559 # 9559 on physical robots, sim-specific otherwise +NAOQI_PORT=9559 ``` +The same `.env` also includes opt-in switches for GPU STT and dev-overlay shortcuts — see the example for the full list. + ### 3. Build and run (MVP) ```bash @@ -97,7 +99,7 @@ The overlay mounts both sibling repos into the `pepper-wizard` container for liv #### Shortening the dev-overlay commands (optional) -Add a `.env` file (gitignored) to the repo root: +Uncomment these lines in your `.env` (created in step 2): ``` COMPOSE_FILE=docker-compose.yml:docker-compose.dev.yml @@ -116,7 +118,7 @@ Omit `COMPOSE_PROFILES=gpu` on hosts without NVIDIA runtime. #### Running against the simulator -Set `NAOQI_IP=127.0.0.1` in `robot.env` to trigger sim mode — `pepper-robot-env`'s entrypoint detects the local IP and boots qiBullet instead of pynaoqi. On first boot it auto-seeds the qiBullet asset cache (Pepper URDF + meshes) into `../PepperBox/.qibullet/`. +Set `NAOQI_IP=127.0.0.1` in `.env` to trigger sim mode — `pepper-robot-env`'s entrypoint detects the local IP and boots qiBullet instead of pynaoqi. On first boot it auto-seeds the qiBullet asset cache (Pepper URDF + meshes) into `../PepperBox/.qibullet/`. The cache directory is auto-created by Docker as `root`-owned on first mount, which blocks the container's `pepperdev` user (UID 1000) from writing. If the entrypoint prints a permission-denied message, chown the host directory once and recreate the container: diff --git a/docker-compose.dev.yml b/docker-compose.dev.yml index 1f29da6..ba871b6 100644 --- a/docker-compose.dev.yml +++ b/docker-compose.dev.yml @@ -8,10 +8,6 @@ # GPU-heavy services (perception) are gated behind the `gpu` profile so the # rest of the stack runs on hosts without an NVIDIA runtime. Enable with: # docker compose -f docker-compose.yml -f docker-compose.dev.yml --profile gpu up -# -# The overlay swaps `pepper-robot-env` from the published NAOqi bridge image to -# the sim-capable `pepper-box:latest` (headless qiBullet) and mounts /dev/dri so -# pybullet's GL renderer works on any open-source driver (AMD, Intel, NVIDIA). services: pepper-robot-env: @@ -36,11 +32,11 @@ services: image: ghcr.io/action-prediction-lab/pepper-box:latest network_mode: host volumes: - - ./robot.env:/home/pepperdev/py3-naoqi-bridge/robot.env + - ./.env:/home/pepperdev/py3-naoqi-bridge/robot.env - ../PepperBox/py3-naoqi-bridge:/home/pepperdev/py3-naoqi-bridge - ${HOME}/.pepperbox/pynaoqi-python2.7-2.5.7.1-linux64:/opt/pynaoqi-python2.7-2.5.7.1-linux64:ro env_file: - - robot.env + - .env environment: - PYTHONUNBUFFERED=1 command: [ "python2", "/home/pepperdev/py3-naoqi-bridge/state_service.py" ] @@ -50,11 +46,11 @@ services: image: ghcr.io/action-prediction-lab/pepper-box:latest network_mode: host volumes: - - ./robot.env:/home/pepperdev/py3-naoqi-bridge/robot.env + - ./.env:/home/pepperdev/py3-naoqi-bridge/robot.env - ../PepperBox/py3-naoqi-bridge:/home/pepperdev/py3-naoqi-bridge - ${HOME}/.pepperbox/pynaoqi-python2.7-2.5.7.1-linux64:/opt/pynaoqi-python2.7-2.5.7.1-linux64:ro env_file: - - robot.env + - .env environment: - PYTHONUNBUFFERED=1 command: [ "python2", "/home/pepperdev/py3-naoqi-bridge/audio_publisher.py" ] diff --git a/docker-compose.yml b/docker-compose.yml index ba69cf5..c858a77 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -5,8 +5,8 @@ # docker compose -f docker-compose.yml -f docker-compose.dev.yml up # # First-run setup: -# ./setup.sh -# cp robot.env.example robot.env && $EDITOR robot.env +# ./setup.sh +# cp .env.example .env && $EDITOR .env # docker compose up services: @@ -14,10 +14,10 @@ services: image: ghcr.io/action-prediction-lab/pepper-box:latest network_mode: host volumes: - - ./robot.env:/home/pepperdev/py3-naoqi-bridge/robot.env + - ./.env:/home/pepperdev/py3-naoqi-bridge/robot.env - ${HOME}/.pepperbox/pynaoqi-python2.7-2.5.7.1-linux64:/opt/pynaoqi-python2.7-2.5.7.1-linux64:ro env_file: - - robot.env + - .env environment: - PYTHONUNBUFFERED=1 command: [ "/home/pepperdev/entrypoint.sh" ] @@ -27,7 +27,7 @@ services: image: ghcr.io/action-prediction-lab/pepper-wizard:latest build: . env_file: - - robot.env + - .env environment: - HF_HUB_OFFLINE=1 - TRANSFORMERS_OFFLINE=1 @@ -50,6 +50,7 @@ services: build: context: . dockerfile: stt-service/Dockerfile + runtime: ${STT_DOCKER_RUNTIME:-runc} network_mode: host volumes: - ./stt-service:/app @@ -59,13 +60,15 @@ services: environment: - PULSE_SERVER=unix:/run/user/1000/pulse/native - PYTHONUNBUFFERED=1 - - HF_HOME=/opt/whisper-cache + - HF_HOME=/opt/asr-cache + - NVIDIA_VISIBLE_DEVICES=${STT_GPU_DEVICES:-none} + - NVIDIA_DRIVER_CAPABILITIES=compute,utility healthcheck: # Round-trips a real ZMQ ping so pepper-wizard only starts after STT - # is answering on :5562 (and Whisper model is loaded). + # is answering on :5562 (and the backend model is loaded). test: ["CMD", "python", "-c", "import sys, zmq; ctx=zmq.Context.instance(); s=ctx.socket(zmq.REQ); s.setsockopt(zmq.RCVTIMEO,2000); s.setsockopt(zmq.SNDTIMEO,2000); s.setsockopt(zmq.LINGER,0); s.connect('tcp://localhost:5562'); s.send_json({'action':'ping'}); sys.exit(0 if s.recv_json().get('status')=='ok' else 1)"] interval: 2s timeout: 3s retries: 15 - start_period: 30s + start_period: 60s restart: on-failure diff --git a/pepper_wizard/cli.py b/pepper_wizard/cli.py index 202f4a4..2c765cf 100644 --- a/pepper_wizard/cli.py +++ b/pepper_wizard/cli.py @@ -804,6 +804,7 @@ def llm_talk_session(robot_client, config, verbose=False): print_formatted_text(HTML( "Error: Could not connect to the STT service." )) + logger.error("STTEnableFailed", {"phase": "ping"}) stt_client.close() return @@ -811,6 +812,7 @@ def llm_talk_session(robot_client, config, verbose=False): print_formatted_text(HTML( "Error: Could not enable streaming on STT service." )) + logger.error("STTEnableFailed", {"phase": "enable_streaming"}) stt_client.close() return logger.info("StreamingEnabled", {}) @@ -910,6 +912,7 @@ def _handle_vad_event(evt, *, review_mode, llm, stt, robot_client, logger, sessi print_formatted_text( HTML("[VAD error] {}: {}").format(evt.get("error", ""), evt.get("detail", "")) ) + logger.error("VADError", {"error": evt.get("error", ""), "detail": evt.get("detail", "")}) return text = (evt.get("text") or "").strip() diff --git a/pepper_wizard/config.py b/pepper_wizard/config.py index 78d24c9..1643667 100644 --- a/pepper_wizard/config.py +++ b/pepper_wizard/config.py @@ -103,9 +103,11 @@ def load_llm_config(file_path): def load_stt_config(file_path): """Loads speech-to-text configuration from a JSON file.""" defaults = { + "engine": "whisper", "zmq_address": "tcp://localhost:5562", "review_mode": True, - "model_size": "base.en", + "whisper_model": "base.en", + "parakeet_model": "nvidia/parakeet-tdt-0.6b-v2", "sample_rate": 16000, "push_to_talk_key": "space", } diff --git a/pepper_wizard/config/stt.json b/pepper_wizard/config/stt.json index 1df1f7a..4e11226 100644 --- a/pepper_wizard/config/stt.json +++ b/pepper_wizard/config/stt.json @@ -1,10 +1,12 @@ { + "engine": "parakeet", "zmq_address": "tcp://localhost:5562", "audio_zmq_address": "tcp://localhost:5563", "transcription_zmq_address": "tcp://localhost:5564", "review_mode": true, "llm_review_mode": true, - "model_size": "base.en", + "whisper_model": "base.en", + "parakeet_model": "nvidia/parakeet-tdt-0.6b-v2", "sample_rate": 16000, "push_to_talk_key": "space", "vad": { diff --git a/pepper_wizard/probe/cli.py b/pepper_wizard/probe/cli.py index 930591c..a593cdf 100644 --- a/pepper_wizard/probe/cli.py +++ b/pepper_wizard/probe/cli.py @@ -41,7 +41,7 @@ def main(argv: Optional[Sequence[str]] = None) -> int: ) parser.add_argument( "--robot-env", type=str, default=None, - help="Path to robot.env (default: ./robot.env if present).", + help="Path to the project env file (default: ./.env if present).", ) args = parser.parse_args(argv) diff --git a/pepper_wizard/probe/detect.py b/pepper_wizard/probe/detect.py index 1664bb5..4d9e77a 100644 --- a/pepper_wizard/probe/detect.py +++ b/pepper_wizard/probe/detect.py @@ -94,7 +94,7 @@ def detect_robot( ) -> DetectorResult: raw = {} if ip is None or port is None: - env_path = Path(robot_env_path) if robot_env_path else Path("robot.env") + env_path = Path(robot_env_path) if robot_env_path else Path(".env") raw["robot_env_path"] = str(env_path) raw["robot_env_exists"] = env_path.exists() if env_path.exists(): @@ -108,7 +108,7 @@ def detect_robot( if not ip or not port: return DetectorResult( value="missing-config", - detail="NAOQI_IP/NAOQI_PORT not set (no robot.env or incomplete)", + detail="NAOQI_IP/NAOQI_PORT not set (no .env or incomplete)", raw=raw, ) diff --git a/pepper_wizard/probe/profile.py b/pepper_wizard/probe/profile.py index b537b83..1139afa 100644 --- a/pepper_wizard/probe/profile.py +++ b/pepper_wizard/probe/profile.py @@ -29,7 +29,7 @@ def probe(cls, robot_env_path: Optional[str] = None) -> "Profile": def recommend(self) -> Recommendation: missing = [] if self.robot.value == "missing-config": - missing.append("Robot endpoint not configured — set NAOQI_IP/NAOQI_PORT in robot.env") + missing.append("Robot endpoint not configured — set NAOQI_IP/NAOQI_PORT in .env") elif self.robot.value == "unreachable": missing.append(f"Robot unreachable — {self.robot.detail}") if self.audio.value == "none": diff --git a/pepper_wizard/stt_client.py b/pepper_wizard/stt_client.py index e5eb2dc..1a5bbb0 100644 --- a/pepper_wizard/stt_client.py +++ b/pepper_wizard/stt_client.py @@ -4,6 +4,8 @@ import zmq +from .logger import get_logger + class STTClient: """Thin wrapper around the ZMQ REQ socket to the STT service.""" @@ -23,6 +25,7 @@ def __init__(self, zmq_address: str, timeout_ms: int = 30000): self.socket.setsockopt(zmq.LINGER, 0) self.socket.connect(zmq_address) self._connected = False + self._log = get_logger("STTClient") def ping(self) -> bool: """Check if the STT service is reachable.""" @@ -46,6 +49,7 @@ def start_recording(self) -> bool: return reply.get("status") == "recording" except zmq.ZMQError as e: print(f"[STTClient] Start error: {e}") + self._log.error("STTClientError", {"action": "start", "error": str(e)}) return False def stop_and_transcribe(self) -> dict: @@ -64,6 +68,7 @@ def stop_and_transcribe(self) -> dict: return reply except zmq.ZMQError as e: print(f"[STTClient] Stop/transcribe error: {e}") + self._log.error("STTClientError", {"action": "stop", "error": str(e)}) return {"transcription": "", "error": str(e)} def enable_streaming(self) -> bool: @@ -111,6 +116,7 @@ def _simple_action(self, action: str, expected_status: str) -> bool: return reply.get("status") == expected_status except zmq.ZMQError as e: print(f"[STTClient] {action} error: {e}") + self._log.error("STTClientError", {"action": action, "error": str(e)}) return False @property diff --git a/robot.env.example b/robot.env.example deleted file mode 100644 index 2ccecde..0000000 --- a/robot.env.example +++ /dev/null @@ -1,15 +0,0 @@ -# Robot connection for the PepperWizard → NAOqi bridge. -# -# Copy this file to `robot.env` and adjust if your robot isn't the lab's default: -# cp robot.env.example robot.env -# -# Physical robot: set NAOQI_IP to the robot's IP and NAOQI_PORT=9559. -# Simulation: the dev overlay (docker-compose.dev.yml) runs a qiBullet sim -# inside pepper-box:latest. Set NAOQI_IP=127.0.0.1 to trigger sim mode; -# NAOQI_PORT is unused in that case but must still be present for the -# env_file parser. - -NAOQI_IP=192.168.123.50 -NAOQI_PORT=9559 - -# For LLM talk mode, export ANTHROPIC_API_KEY in your shell before running `docker compose up`. \ No newline at end of file diff --git a/stt-service/Dockerfile b/stt-service/Dockerfile index a234a96..22a5a03 100644 --- a/stt-service/Dockerfile +++ b/stt-service/Dockerfile @@ -1,14 +1,14 @@ -FROM python:3.9-slim +FROM python:3.11-slim WORKDIR /app -# Install system dependencies for sounddevice (PortAudio) and PulseAudio RUN apt-get update && apt-get install -y \ libportaudio2 \ libsndfile1 \ libpulse0 \ libasound2-plugins \ pulseaudio-utils \ + ffmpeg \ && rm -rf /var/lib/apt/lists/* # Configure ALSA to route through PulseAudio (PortAudio uses ALSA backend) @@ -21,19 +21,15 @@ RUN echo '\ }\n' > /etc/asound.conf COPY stt-service/requirements.txt . -RUN pip install --no-cache-dir -r requirements.txt +RUN pip install --no-cache-dir --index-url https://download.pytorch.org/whl/cu124 torch torchaudio \ + && pip install --no-cache-dir -r requirements.txt -# Bake the Whisper model into the image so first-boot is fast and offline-capable. -# Reads model_size from pepper_wizard/config/stt.json -# shared with the runtime config (mounted at /app/pepper_config/stt.json). -# HF_HOME points outside /app so the runtime bind-mount of ./stt-service:/app -ENV HF_HOME=/opt/whisper-cache + +ENV HF_HOME=/opt/asr-cache COPY pepper_wizard/config/stt.json /tmp/stt.json +COPY stt-service/bake_models.py stt-service/config_loader.py stt-service/vad_segmenter.py /tmp/ RUN mkdir -p ${HF_HOME} && \ - python -c "import json; from faster_whisper import WhisperModel; \ -m = json.load(open('/tmp/stt.json')).get('model_size', 'base.en'); \ -print(f'[build] baking whisper model: {m}'); \ -WhisperModel(m, device='cpu', compute_type='int8')" + PYTHONPATH=/tmp python -u /tmp/bake_models.py /tmp/stt.json COPY stt-service/. . diff --git a/stt-service/backends.py b/stt-service/backends.py new file mode 100644 index 0000000..51eedb2 --- /dev/null +++ b/stt-service/backends.py @@ -0,0 +1,126 @@ +"""Engine abstraction for stt-service. + +Defines the STTBackend Protocol and WhisperBackend / ParakeetBackend +implementations. make_backend(config) factory selects between +them based on stt.json's `engine` key. +""" +import sys +from typing import Any, Iterable, Protocol, Tuple + +import numpy as np +from faster_whisper import WhisperModel + +from config_loader import ( + DEFAULT_PARAKEET_MODEL, + _resolve_whisper_model, +) + + +class Segment(Protocol): + text: str + + +class STTBackend(Protocol): + def transcribe( + self, + audio: np.ndarray, + beam_size: int = 3, + language: str = "en", + vad_filter: bool = False, + ) -> Tuple[Iterable[Segment], Any]: + ... + + +class EngineLoadError(RuntimeError): + """Raised when the configured engine cannot be loaded. + + Caught at the top of main() to produce a loud, actionable exit. + """ + + +class WhisperBackend: + """Wraps faster_whisper.WhisperModel. Delegates transcribe() directly.""" + + def __init__(self, model_size: str, device: str = "cpu", compute_type: str = "int8"): + print(f"[STTService] Loading whisper model '{model_size}' ({device}, {compute_type})...", + file=sys.stderr) + self._model = WhisperModel(model_size, device=device, compute_type=compute_type) + print(f"[STTService] Whisper model loaded.", file=sys.stderr) + + def transcribe( + self, + audio: np.ndarray, + beam_size: int = 3, + language: str = "en", + vad_filter: bool = False, + ) -> Tuple[Iterable[Segment], Any]: + return self._model.transcribe( + audio, beam_size=beam_size, language=language, vad_filter=vad_filter, + ) + + +class ParakeetBackend: + """Wraps NeMo ASRModel. Lazy-imports nemo inside __init__ so Whisper-only + deployments don't pay NeMo's import cost. + """ + + def __init__(self, model_name: str, device: str = "cuda"): + print(f"[STTService] Loading parakeet model '{model_name}' ({device})...", + file=sys.stderr) + import nemo.collections.asr # noqa: F401 + from nemo.collections.asr.models import ASRModel + # Stage on CPU then cast to FP16 before moving to GPU. + # Params selected for Ada500 + model = ASRModel.from_pretrained(model_name, map_location="cpu") + if device != "cpu": + model = model.half().to(device) + self._model = model + print(f"[STTService] Parakeet model loaded.", file=sys.stderr) + + def transcribe( + self, + audio: np.ndarray, + beam_size: int = 3, + language: str = "en", + vad_filter: bool = False, + ) -> Tuple[Iterable[Segment], Any]: + # NeMo expects a list of audio inputs; we batch one item. + hypotheses = self._model.transcribe([audio]) + text = hypotheses[0].text if hypotheses else "" + seg = type("Seg", (), {"text": text})() + return [seg], None + + +def make_backend(config: dict) -> STTBackend: + """Select and construct the configured engine. + + Raises EngineLoadError on misconfiguration or precondition failure. Caller + is expected to print the message and exit non-zero before binding ZMQ. + """ + engine = config.get("engine", "whisper") + + if engine == "whisper": + return WhisperBackend( + model_size=_resolve_whisper_model(config), + device="cpu", + compute_type="int8", + ) + + if engine == "parakeet": + import torch + if not torch.cuda.is_available(): + raise EngineLoadError( + "engine: parakeet requires CUDA. " + "torch.cuda.is_available() returned False. " + "Fix the nvidia runtime: ensure STT_DOCKER_RUNTIME=nvidia " + "and STT_GPU_DEVICES=all are set in .env (check `nvidia-smi` " + "on the host first), or set `engine: whisper` in stt.json." + ) + return ParakeetBackend( + model_name=config.get("parakeet_model", DEFAULT_PARAKEET_MODEL), + device="cuda", + ) + + raise EngineLoadError( + f"Unknown engine: {engine!r}. Valid values: 'whisper', 'parakeet'." + ) diff --git a/stt-service/bake_models.py b/stt-service/bake_models.py new file mode 100644 index 0000000..f2f51ec --- /dev/null +++ b/stt-service/bake_models.py @@ -0,0 +1,57 @@ +"""Build-time entry point that pre-populates ${HF_HOME} with both stt model engine +checkpoints. Reads model names from stt.json and falls back to library +defaults if a key is absent. Intended to run inside the Dockerfile. + +Bakes Whisper and Parakeet so the resulting +image carries weights for both at runtime. +""" +import sys + +from config_loader import ( + DEFAULT_PARAKEET_MODEL, + DEFAULT_WHISPER_MODEL, + _load_config, + _resolve_whisper_model, +) + + +def _log(msg: str) -> None: + # flush=True so Docker buildkit can stream progress in real time. + print(f"[bake_models] {msg}", flush=True) + + +def _bake_whisper(model_size: str) -> None: + _log(f"importing faster_whisper") + from faster_whisper import WhisperModel + _log(f"Baking whisper checkpoint: {model_size}") + WhisperModel(model_size, device="cpu", compute_type="int8") + _log("whisper bake done") + + +def _bake_parakeet(model_name: str) -> None: + _log("importing nemo.collections.asr.models (slow, ~30s)") + from nemo.collections.asr.models import ASRModel + _log(f"Baking parakeet checkpoint: {model_name}") + ASRModel.from_pretrained(model_name) + _log("parakeet bake done") + + +def main() -> int: + if len(sys.argv) < 2: + print("[bake_models] usage: bake_models.py ", + file=sys.stderr) + return 2 + + config_path = sys.argv[1] + config = _load_config(config_path) + + whisper_model = _resolve_whisper_model(config) + parakeet_model = config.get("parakeet_model", DEFAULT_PARAKEET_MODEL) + + _bake_whisper(whisper_model) + _bake_parakeet(parakeet_model) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/stt-service/config_loader.py b/stt-service/config_loader.py new file mode 100644 index 0000000..3dc060a --- /dev/null +++ b/stt-service/config_loader.py @@ -0,0 +1,51 @@ +"""Configuration loading and resolution for stt-service. + +Shared by main.py at runtime and bake_models.py at image build time. +""" +import json +import sys + +from vad_segmenter import VadConfig + + +DEFAULT_WHISPER_MODEL = "base.en" +DEFAULT_PARAKEET_MODEL = "nvidia/parakeet-tdt-0.6b-v2" + +DEFAULT_VAD = { + "threshold": 0.5, + "min_silence_ms": 700, + "min_utterance_ms": 300, + "max_utterance_ms": 15000, + "preroll_ms": 200, +} + + +def _load_config(path: str = "/app/pepper_config/stt.json") -> dict: + """Read stt.json. Return {} on missing file or malformed JSON.""" + try: + with open(path) as f: + return json.load(f) + except (FileNotFoundError, json.JSONDecodeError, OSError): + return {} + + +def _resolve_whisper_model(config: dict) -> str: + """Resolve the Whisper checkpoint, honouring the legacy `model_size` alias.""" + if "whisper_model" in config: + return config["whisper_model"] + if "model_size" in config: + print( + "[STTService] DEPRECATED: 'model_size' in stt.json is now " + "'whisper_model'. The old key still works but will be removed " + "in a future release.", + file=sys.stderr, + ) + return config["model_size"] + return DEFAULT_WHISPER_MODEL + + +def _load_vad_config(config: dict) -> VadConfig: + """Build a VadConfig by merging the `vad` sub-dict over DEFAULT_VAD.""" + params = dict(DEFAULT_VAD) + params.update(config.get("vad", {})) + return VadConfig(**params) diff --git a/stt-service/main.py b/stt-service/main.py index 4a2166b..d295f66 100644 --- a/stt-service/main.py +++ b/stt-service/main.py @@ -2,8 +2,8 @@ STT Service — Speech-to-Text microservice for PepperWizard. Captures audio from the host microphone (via PortAudio/ALSA routed through PulseAudio) and transcribes it -using a configurable Whisper model. Communicates with pepper-wizard over -ZMQ REQ/REP. +using a configurable backend (faster-whisper or NVIDIA Parakeet via NeMo). +Communicates with pepper-wizard over ZMQ REQ/REP. Protocol: REQ: {"action": "start"} → REP: {"status": "recording"} @@ -12,32 +12,33 @@ """ import argparse -import json -import os +import sys import time import threading from datetime import datetime, timezone +from typing import Optional import numpy as np import sounddevice as sd import zmq -from faster_whisper import WhisperModel -from vad_segmenter import VadSegmenter, VadConfig +from backends import EngineLoadError, STTBackend, make_backend +from config_loader import _load_config, _load_vad_config from events import UtteranceEvent, encode_event, encode_error +from vad_segmenter import VadSegmenter, VadConfig class StreamingWorker(threading.Thread): """Consumes robot-mic audio on a SUB, segments via VAD, transcribes with - Whisper (sequentially), and publishes JSON utterance events on a PUB socket.""" + the configured backend (sequentially), and publishes JSON utterance events on a PUB socket.""" def __init__(self, audio_addr: str, pub_addr: str, - vad_config: VadConfig, whisper, is_muted): + vad_config: VadConfig, backend: STTBackend, is_muted): super().__init__(daemon=True) self._audio_addr = audio_addr self._pub_addr = pub_addr self._vad_config = vad_config - self._whisper = whisper + self._backend = backend self._is_muted = is_muted self._stop_evt = threading.Event() @@ -62,12 +63,12 @@ def on_utterance(pcm_bytes: bytes) -> None: t_start = utt_start[0] or t_end audio_f32 = np.frombuffer(pcm_bytes, dtype=np.int16).astype(np.float32) / 32768.0 try: - segments, _ = self._whisper.transcribe( + segments, _ = self._backend.transcribe( audio_f32, beam_size=3, language="en", vad_filter=False, ) text = " ".join(s.text.strip() for s in segments).strip() except Exception as e: - pub.send_string(encode_error("whisper_failed", str(e), t_start)) + pub.send_string(encode_error("transcribe_failed", str(e), t_start)) utt_start[0] = None return duration_s = len(audio_f32) / 16000.0 @@ -141,51 +142,13 @@ def stop(self) -> np.ndarray: return audio -DEFAULT_VAD = { - "threshold": 0.5, - "min_silence_ms": 700, - "min_utterance_ms": 300, - "max_utterance_ms": 15000, - "preroll_ms": 200, -} - -DEFAULT_MODEL_SIZE = "small.en" - - -def _load_vad_config(path: str = "/app/pepper_config/stt.json") -> VadConfig: - """Load VAD parameters from stt.json if present; fall back to DEFAULT_VAD.""" - params = dict(DEFAULT_VAD) - try: - with open(path) as f: - data = json.load(f) - params.update(data.get("vad", {})) - except (FileNotFoundError, json.JSONDecodeError, OSError): - pass - return VadConfig(**params) - - -def _load_model_size(path: str = "/app/pepper_config/stt.json") -> str: - """Load Whisper model_size from stt.json if present; fall back to DEFAULT_MODEL_SIZE.""" - try: - with open(path) as f: - data = json.load(f) - return data.get("model_size", DEFAULT_MODEL_SIZE) - except (FileNotFoundError, json.JSONDecodeError, OSError): - return DEFAULT_MODEL_SIZE - - class STTService: """ZMQ REQ/REP service that records and transcribes on demand.""" - def __init__(self, model_size: str, zmq_port: int, sample_rate: int): + def __init__(self, backend: STTBackend, zmq_port: int, sample_rate: int, + config: Optional[dict] = None): self.recorder = AudioRecorder(sample_rate=sample_rate) - - # Load the Whisper model (CPU, int8 quantised) - print(f"[STTService] Loading whisper model '{model_size}' (cpu, int8)...") - self.model = WhisperModel( - model_size, device="cpu", compute_type="int8" - ) - print(f"[STTService] Model loaded.") + self.backend = backend # ZMQ setup self.context = zmq.Context() @@ -196,16 +159,16 @@ def __init__(self, model_size: str, zmq_port: int, sample_rate: int): self._muted = False self._worker = None - self._vad_config = _load_vad_config() + self._vad_config = _load_vad_config(config or {}) self._audio_addr = "tcp://localhost:5563" self._pub_addr = "tcp://*:5564" def transcribe(self, audio: np.ndarray) -> str: - """Run whisper transcription on a float32 audio array.""" + """Run transcription on a float32 audio array.""" if len(audio) == 0: return "" - segments, _info = self.model.transcribe( + segments, _info = self.backend.transcribe( audio, beam_size=3, language="en", @@ -249,7 +212,6 @@ def _handle_action(self, msg: dict) -> dict: duration = len(audio) / self.recorder.sample_rate if duration < 0.3: - # Too short to be meaningful speech return { "transcription": "", "duration": round(duration, 2), @@ -268,12 +230,15 @@ def _handle_action(self, msg: dict) -> dict: } elif action == "enable_streaming": + # Clear any stale mute state from a previous session + # whose unmute() never reached us (e.g. wizard exited mid-dispatch). + self._muted = False if self._worker is None: self._worker = StreamingWorker( audio_addr=self._audio_addr, pub_addr=self._pub_addr, vad_config=self._vad_config, - whisper=self.model, + backend=self.backend, is_muted=lambda: self._muted, ) self._worker.start() @@ -284,6 +249,9 @@ def _handle_action(self, msg: dict) -> dict: self._worker.stop() self._worker.join(timeout=2.0) self._worker = None + # mute/unmute are per-dispatch state; clear so the next streaming + # session does not inherit a stuck-True flag from a dropped unmute(). + self._muted = False return {"status": "idle"} elif action == "mute": @@ -320,7 +288,7 @@ def main(): parser = argparse.ArgumentParser(description="PepperWizard STT Service") parser.add_argument( "--model", type=str, default=None, - help="Whisper model size (overrides stt.json model_size; default taken from stt.json)", + help="Whisper model size (overrides stt.json whisper_model; ignored when engine is not 'whisper').", ) parser.add_argument( "--port", type=int, default=5562, @@ -332,12 +300,29 @@ def main(): ) args = parser.parse_args() - model_size = args.model if args.model is not None else _load_model_size() + config = _load_config() + + if args.model is not None: + if config.get("engine", "whisper") == "whisper": + config["whisper_model"] = args.model + else: + print( + f"[STTService] --model={args.model!r} ignored: engine is " + f"{config.get('engine')!r}, not 'whisper'.", + file=sys.stderr, + ) + + try: + backend = make_backend(config) + except EngineLoadError as e: + print(f"[STTService] {e}", file=sys.stderr) + sys.exit(1) service = STTService( - model_size=model_size, + backend=backend, zmq_port=args.port, sample_rate=args.sample_rate, + config=config, ) service.run() diff --git a/stt-service/requirements.txt b/stt-service/requirements.txt index cae9c93..f78d40f 100644 --- a/stt-service/requirements.txt +++ b/stt-service/requirements.txt @@ -3,3 +3,5 @@ silero-vad sounddevice pyzmq numpy +nemo_toolkit[asr]>=2.0 +torch>=2.1 diff --git a/stt-service/tests/test_backends.py b/stt-service/tests/test_backends.py new file mode 100644 index 0000000..43e18cd --- /dev/null +++ b/stt-service/tests/test_backends.py @@ -0,0 +1,129 @@ +"""Unit tests for backends.py: Protocol, WhisperBackend, ParakeetBackend, factory.""" +import unittest +from unittest import mock + +import numpy as np + + +class TestWhisperBackend(unittest.TestCase): + @mock.patch("backends.WhisperModel") + def test_constructs_whisper_model_with_args(self, FakeWhisperModel): + from backends import WhisperBackend + WhisperBackend(model_size="tiny.en", device="cpu", compute_type="int8") + FakeWhisperModel.assert_called_once_with( + "tiny.en", device="cpu", compute_type="int8" + ) + + @mock.patch("backends.WhisperModel") + def test_transcribe_delegates_to_underlying_model(self, FakeWhisperModel): + from backends import WhisperBackend + fake_model = FakeWhisperModel.return_value + class Seg: + text = "hello" + fake_model.transcribe.return_value = ([Seg()], "info") + + backend = WhisperBackend("tiny.en", device="cpu", compute_type="int8") + audio = np.zeros(16000, dtype=np.float32) + segments, info = backend.transcribe(audio, beam_size=3, language="en", vad_filter=True) + + fake_model.transcribe.assert_called_once_with( + audio, beam_size=3, language="en", vad_filter=True, + ) + self.assertEqual(list(segments)[0].text, "hello") + self.assertEqual(info, "info") + + +class TestParakeetBackend(unittest.TestCase): + def test_lazy_imports_nemo_inside_init(self): + import sys + from backends import ParakeetBackend + # ParakeetBackend.__init__ imports nemo lazily; the module class import + # alone must not have imported nemo.collections.asr. + self.assertNotIn("nemo.collections.asr", sys.modules, + "nemo.collections.asr must not be imported until ParakeetBackend() is constructed") + + def test_constructs_and_transcribes_with_mocked_nemo(self): + from backends import ParakeetBackend + fake_nemo_module = mock.MagicMock() + fake_asr_model_cls = fake_nemo_module.collections.asr.models.ASRModel + fake_loaded = fake_asr_model_cls.from_pretrained.return_value + fake_half = fake_loaded.half.return_value + fake_on_device = fake_half.to.return_value + fake_hyp = mock.MagicMock() + fake_hyp.text = "hello pepper" + fake_on_device.transcribe.return_value = [fake_hyp] + + with mock.patch.dict("sys.modules", { + "nemo": fake_nemo_module, + "nemo.collections": fake_nemo_module.collections, + "nemo.collections.asr": fake_nemo_module.collections.asr, + "nemo.collections.asr.models": fake_nemo_module.collections.asr.models, + }): + backend = ParakeetBackend( + model_name="nvidia/parakeet-tdt-0.6b-v2", device="cuda", + ) + audio = np.zeros(16000, dtype=np.float32) + segments, info = backend.transcribe(audio) + seg_list = list(segments) + + fake_asr_model_cls.from_pretrained.assert_called_once_with( + "nvidia/parakeet-tdt-0.6b-v2", map_location="cpu", + ) + fake_loaded.half.assert_called_once_with() + fake_half.to.assert_called_once_with("cuda") + fake_on_device.transcribe.assert_called_once() + self.assertEqual(len(seg_list), 1) + self.assertEqual(seg_list[0].text, "hello pepper") + self.assertIsNone(info) + + +class TestMakeBackend(unittest.TestCase): + @mock.patch("backends.WhisperModel") + def test_whisper_branch_returns_whisper_backend(self, FakeWhisperModel): + from backends import WhisperBackend, make_backend + backend = make_backend({"engine": "whisper", "whisper_model": "base.en"}) + self.assertIsInstance(backend, WhisperBackend) + FakeWhisperModel.assert_called_once_with( + "base.en", device="cpu", compute_type="int8", + ) + + @mock.patch("torch.cuda.is_available", return_value=False) + def test_parakeet_without_cuda_raises_loud_error(self, _is_avail): + from backends import EngineLoadError, make_backend + with self.assertRaises(EngineLoadError) as ctx: + make_backend({"engine": "parakeet", "parakeet_model": "nvidia/parakeet-tdt-0.6b-v2"}) + msg = str(ctx.exception) + # Loud-failure contract: the message must be actionable. + self.assertIn("CUDA", msg) + self.assertIn("torch.cuda.is_available()", msg) + self.assertIn("STT_DOCKER_RUNTIME", msg) + self.assertIn("STT_GPU_DEVICES", msg) + self.assertIn("engine: whisper", msg) + + @mock.patch("torch.cuda.is_available", return_value=True) + def test_parakeet_with_cuda_returns_parakeet_backend(self, _is_avail): + from backends import ParakeetBackend, make_backend + fake_nemo_module = mock.MagicMock() + fake_asr_model_cls = fake_nemo_module.collections.asr.models.ASRModel + with mock.patch.dict("sys.modules", { + "nemo": fake_nemo_module, + "nemo.collections": fake_nemo_module.collections, + "nemo.collections.asr": fake_nemo_module.collections.asr, + "nemo.collections.asr.models": fake_nemo_module.collections.asr.models, + }): + backend = make_backend({ + "engine": "parakeet", + "parakeet_model": "nvidia/parakeet-tdt-0.6b-v2", + }) + self.assertIsInstance(backend, ParakeetBackend) + fake_asr_model_cls.from_pretrained.assert_called_once_with( + "nvidia/parakeet-tdt-0.6b-v2", map_location="cpu", + ) + + def test_unknown_engine_raises(self): + from backends import EngineLoadError, make_backend + with self.assertRaises(EngineLoadError) as ctx: + make_backend({"engine": "bogus"}) + self.assertIn("Unknown engine: 'bogus'", str(ctx.exception)) + self.assertIn("'whisper'", str(ctx.exception)) + self.assertIn("'parakeet'", str(ctx.exception)) diff --git a/stt-service/tests/test_events.py b/stt-service/tests/test_events.py index 7280e2b..5fdca5b 100644 --- a/stt-service/tests/test_events.py +++ b/stt-service/tests/test_events.py @@ -19,7 +19,7 @@ def test_utterance_event_round_trip(self): def test_error_event(self): t0 = datetime(2026, 4, 19, 15, 44, 12, tzinfo=timezone.utc) - payload = json.loads(encode_error("whisper_failed", "cuda oom", t0)) - self.assertEqual(payload["error"], "whisper_failed") + payload = json.loads(encode_error("transcribe_failed", "cuda oom", t0)) + self.assertEqual(payload["error"], "transcribe_failed") self.assertEqual(payload["detail"], "cuda oom") self.assertEqual(payload["t_start"], "2026-04-19T15:44:12.000Z") diff --git a/stt-service/tests/test_model_size_loader.py b/stt-service/tests/test_model_size_loader.py deleted file mode 100644 index 3c2729b..0000000 --- a/stt-service/tests/test_model_size_loader.py +++ /dev/null @@ -1,42 +0,0 @@ -"""Unit tests for _load_model_size in main.py.""" -import json -import os -import tempfile -import unittest - -from main import _load_model_size, DEFAULT_MODEL_SIZE - - -class TestLoadModelSize(unittest.TestCase): - def test_missing_file_returns_default(self): - self.assertEqual( - _load_model_size(path="/nonexistent/path/stt.json"), - DEFAULT_MODEL_SIZE, - ) - - def test_custom_value_overrides_default(self): - with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f: - json.dump({"model_size": "medium.en"}, f) - tmp_path = f.name - try: - self.assertEqual(_load_model_size(path=tmp_path), "medium.en") - finally: - os.unlink(tmp_path) - - def test_missing_key_returns_default(self): - with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f: - json.dump({"other_key": "ignored"}, f) - tmp_path = f.name - try: - self.assertEqual(_load_model_size(path=tmp_path), DEFAULT_MODEL_SIZE) - finally: - os.unlink(tmp_path) - - def test_malformed_json_falls_back_to_default(self): - with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f: - f.write("not valid json {{") - tmp_path = f.name - try: - self.assertEqual(_load_model_size(path=tmp_path), DEFAULT_MODEL_SIZE) - finally: - os.unlink(tmp_path) diff --git a/stt-service/tests/test_streaming_worker.py b/stt-service/tests/test_streaming_worker.py index a59d675..3bb64ec 100644 --- a/stt-service/tests/test_streaming_worker.py +++ b/stt-service/tests/test_streaming_worker.py @@ -28,7 +28,7 @@ def _load_fixture_bytes() -> bytes: return wf.readframes(wf.getnframes()) -class FakeWhisper: +class FakeSTTBackend: def transcribe(self, audio, beam_size=3, language="en", vad_filter=False): class Seg: text = "hello pepper" @@ -50,7 +50,7 @@ def test_publishes_transcription_after_utterance(self): audio_addr="tcp://localhost:16563", pub_addr="tcp://*:16564", vad_config=cfg, - whisper=FakeWhisper(), + backend=FakeSTTBackend(), is_muted=lambda: False, ) w.start() diff --git a/stt-service/tests/test_stt_actions.py b/stt-service/tests/test_stt_actions.py index 1414fa9..ccdcd81 100644 --- a/stt-service/tests/test_stt_actions.py +++ b/stt-service/tests/test_stt_actions.py @@ -9,8 +9,10 @@ class TestSTTActions(unittest.TestCase): @classmethod def setUpClass(cls): - # Share the STTService (and its Whisper load) across the class. - cls.svc = STTService(model_size="tiny.en", zmq_port=15562, sample_rate=16000) + # Share the STTService (and its WhisperBackend load) across the class. + from backends import WhisperBackend + backend = WhisperBackend("tiny.en", device="cpu", compute_type="int8") + cls.svc = STTService(backend=backend, zmq_port=15562, sample_rate=16000) @classmethod def tearDownClass(cls): diff --git a/stt-service/tests/test_vad_loader.py b/stt-service/tests/test_vad_loader.py index f49da72..5cf69fe 100644 --- a/stt-service/tests/test_vad_loader.py +++ b/stt-service/tests/test_vad_loader.py @@ -1,19 +1,16 @@ -"""Unit tests for _load_vad_config in main.py.""" -import json -import os -import tempfile +"""Unit tests for _load_vad_config in config_loader.py.""" import unittest -from main import _load_vad_config, DEFAULT_VAD +from config_loader import DEFAULT_VAD, _load_vad_config from vad_segmenter import VadConfig class TestLoadVadConfig(unittest.TestCase): - """Tests for _load_vad_config with explicit path injection (no real mount required).""" + """Tests for _load_vad_config, which builds a VadConfig from a parsed config dict.""" - def test_missing_file_returns_defaults(self): - """When the config file does not exist, defaults are used.""" - cfg = _load_vad_config(path="/nonexistent/path/stt.json") + def test_empty_config_returns_defaults(self): + """An empty config dict yields VadConfig with all defaults.""" + cfg = _load_vad_config({}) self.assertIsInstance(cfg, VadConfig) self.assertAlmostEqual(cfg.threshold, DEFAULT_VAD["threshold"]) self.assertEqual(cfg.min_silence_ms, DEFAULT_VAD["min_silence_ms"]) @@ -22,25 +19,16 @@ def test_missing_file_returns_defaults(self): self.assertEqual(cfg.preroll_ms, DEFAULT_VAD["preroll_ms"]) def test_custom_values_override_defaults(self): - """Values in the vad block override the defaults.""" - custom = { - "vad": { - "threshold": 0.7, - "min_silence_ms": 500, - "preroll_ms": 100, + """Values in the vad block override the defaults; unset keys retain defaults.""" + cfg = _load_vad_config( + { + "vad": { + "threshold": 0.7, + "min_silence_ms": 500, + "preroll_ms": 100, + } } - } - with tempfile.NamedTemporaryFile( - mode="w", suffix=".json", delete=False - ) as f: - json.dump(custom, f) - tmp_path = f.name - - try: - cfg = _load_vad_config(path=tmp_path) - finally: - os.unlink(tmp_path) - + ) self.assertAlmostEqual(cfg.threshold, 0.7) self.assertEqual(cfg.min_silence_ms, 500) self.assertEqual(cfg.preroll_ms, 100) @@ -49,33 +37,13 @@ def test_custom_values_override_defaults(self): self.assertEqual(cfg.max_utterance_ms, DEFAULT_VAD["max_utterance_ms"]) def test_empty_vad_block_returns_defaults(self): - """A stt.json with an empty vad block returns full defaults.""" - with tempfile.NamedTemporaryFile( - mode="w", suffix=".json", delete=False - ) as f: - json.dump({"vad": {}}, f) - tmp_path = f.name - - try: - cfg = _load_vad_config(path=tmp_path) - finally: - os.unlink(tmp_path) - + """A config with an empty vad block returns full defaults.""" + cfg = _load_vad_config({"vad": {}}) self.assertAlmostEqual(cfg.threshold, DEFAULT_VAD["threshold"]) self.assertEqual(cfg.min_silence_ms, DEFAULT_VAD["min_silence_ms"]) - def test_malformed_json_falls_back_to_defaults(self): - """A corrupt JSON file silently falls back to defaults.""" - with tempfile.NamedTemporaryFile( - mode="w", suffix=".json", delete=False - ) as f: - f.write("not valid json {{") - tmp_path = f.name - - try: - cfg = _load_vad_config(path=tmp_path) - finally: - os.unlink(tmp_path) - + def test_config_without_vad_key_returns_defaults(self): + """A config dict missing the `vad` key falls back to defaults.""" + cfg = _load_vad_config({"other_key": "ignored"}) self.assertIsInstance(cfg, VadConfig) self.assertAlmostEqual(cfg.threshold, DEFAULT_VAD["threshold"]) diff --git a/stt-service/tests/test_whisper_model_loader.py b/stt-service/tests/test_whisper_model_loader.py new file mode 100644 index 0000000..e7fa11d --- /dev/null +++ b/stt-service/tests/test_whisper_model_loader.py @@ -0,0 +1,94 @@ +"""Unit tests for _load_config + _resolve_whisper_model in config_loader.py.""" +import io +import json +import os +import tempfile +import unittest +from contextlib import redirect_stderr + +from config_loader import ( + DEFAULT_WHISPER_MODEL, + _load_config, + _resolve_whisper_model, +) + + +def _write_tmp_json(payload): + with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f: + if isinstance(payload, str): + f.write(payload) + else: + json.dump(payload, f) + return f.name + + +class TestLoadConfig(unittest.TestCase): + def test_missing_file_returns_empty_dict(self): + self.assertEqual(_load_config("/nonexistent/path/stt.json"), {}) + + def test_malformed_json_returns_empty_dict(self): + path = _write_tmp_json("not valid json {{") + try: + self.assertEqual(_load_config(path), {}) + finally: + os.unlink(path) + + +class TestResolveWhisperModel(unittest.TestCase): + def test_prefers_whisper_model_over_model_size(self): + config = {"whisper_model": "small.en", "model_size": "tiny.en"} + with redirect_stderr(io.StringIO()) as err: + self.assertEqual(_resolve_whisper_model(config), "small.en") + self.assertNotIn("DEPRECATED", err.getvalue()) + + def test_falls_back_to_model_size_with_stderr_warning(self): + config = {"model_size": "medium.en"} + with redirect_stderr(io.StringIO()) as err: + self.assertEqual(_resolve_whisper_model(config), "medium.en") + self.assertIn("DEPRECATED", err.getvalue()) + self.assertIn("model_size", err.getvalue()) + self.assertIn("whisper_model", err.getvalue()) + + def test_defaults_when_both_keys_absent(self): + self.assertEqual(_resolve_whisper_model({}), DEFAULT_WHISPER_MODEL) + self.assertEqual(DEFAULT_WHISPER_MODEL, "base.en") + + +class TestEndToEnd(unittest.TestCase): + """Ports the four cases from the deleted test_model_size_loader.py.""" + + def test_missing_file_resolves_to_default(self): + self.assertEqual( + _resolve_whisper_model(_load_config("/nonexistent/path/stt.json")), + DEFAULT_WHISPER_MODEL, + ) + + def test_custom_whisper_model_value(self): + path = _write_tmp_json({"whisper_model": "medium.en"}) + try: + self.assertEqual( + _resolve_whisper_model(_load_config(path)), + "medium.en", + ) + finally: + os.unlink(path) + + def test_missing_key_resolves_to_default(self): + path = _write_tmp_json({"other_key": "ignored"}) + try: + self.assertEqual( + _resolve_whisper_model(_load_config(path)), + DEFAULT_WHISPER_MODEL, + ) + finally: + os.unlink(path) + + def test_malformed_json_resolves_to_default(self): + path = _write_tmp_json("not valid json {{") + try: + self.assertEqual( + _resolve_whisper_model(_load_config(path)), + DEFAULT_WHISPER_MODEL, + ) + finally: + os.unlink(path) diff --git a/stt-service/tools/bench_backends.py b/stt-service/tools/bench_backends.py new file mode 100644 index 0000000..68db559 --- /dev/null +++ b/stt-service/tools/bench_backends.py @@ -0,0 +1,116 @@ +"""Backend latency benchmark. + +Times Whisper-CPU vs Parakeet-GPU on the same speech, after a +warm-up pass. +Reports cold-load time, warm-call latency, and the resulting transcription +so output drift between engines is visible. + +Run inside the stt-service container: + docker compose run --rm stt-service python3 tools/bench_backends.py +""" +import os +import statistics +import sys +import time +import wave + +import numpy as np + +# Ensure /app is on sys.path so the script can run from any cwd. +sys.path.insert(0, "/app") + +from backends import ParakeetBackend, WhisperBackend + + +FIXTURE = "/app/tests/fixture_speech_16k.wav" +N_TIMED = 10 + + +def load_audio(path): + with wave.open(path, "rb") as wf: + assert wf.getframerate() == 16000 + assert wf.getnchannels() == 1 + pcm = np.frombuffer(wf.readframes(wf.getnframes()), dtype=np.int16) + audio = pcm.astype(np.float32) / 32768.0 + return audio, len(audio) / 16000.0 + + +def time_backend(label, backend, audio): + print(f"\n=== {label} ===", flush=True) + print("warm-up call...", flush=True) + t0 = time.perf_counter() + segments, _ = backend.transcribe(audio, beam_size=3, language="en", vad_filter=False) + text = " ".join(s.text.strip() for s in segments).strip() + warmup_ms = (time.perf_counter() - t0) * 1000 + print(f" warm-up: {warmup_ms:.1f} ms") + print(f" text: '{text}'") + + print(f"timed runs (n={N_TIMED})...", flush=True) + samples = [] + for i in range(N_TIMED): + t0 = time.perf_counter() + segments, _ = backend.transcribe(audio, beam_size=3, language="en", vad_filter=False) + _ = " ".join(s.text.strip() for s in segments).strip() + samples.append((time.perf_counter() - t0) * 1000) + return { + "warmup_ms": warmup_ms, + "mean_ms": statistics.mean(samples), + "median_ms": statistics.median(samples), + "min_ms": min(samples), + "max_ms": max(samples), + "stdev_ms": statistics.stdev(samples), + "text": text, + } + + +def main(): + audio, duration_s = load_audio(FIXTURE) + print(f"Fixture: {FIXTURE}") + print(f"Duration: {duration_s:.2f} s, {len(audio)} samples @ 16 kHz") + + print("\n--- loading WhisperBackend (CPU, int8, base.en) ---", flush=True) + t0 = time.perf_counter() + whisper = WhisperBackend("base.en", device="cpu", compute_type="int8") + whisper_load_ms = (time.perf_counter() - t0) * 1000 + print(f" load: {whisper_load_ms:.1f} ms") + whisper_stats = time_backend("WhisperBackend (CPU int8 base.en)", whisper, audio) + whisper_stats["load_ms"] = whisper_load_ms + + print("\n--- loading ParakeetBackend (CUDA FP16, parakeet-tdt-0.6b-v2) ---", flush=True) + t0 = time.perf_counter() + parakeet = ParakeetBackend("nvidia/parakeet-tdt-0.6b-v2", device="cuda") + parakeet_load_ms = (time.perf_counter() - t0) * 1000 + print(f" load: {parakeet_load_ms:.1f} ms") + parakeet_stats = time_backend("ParakeetBackend (CUDA FP16 0.6b-v2)", parakeet, audio) + parakeet_stats["load_ms"] = parakeet_load_ms + + audio_ms = duration_s * 1000 + print("\n\n========== RESULTS ==========") + print(f"Audio length: {duration_s:.2f} s ({audio_ms:.0f} ms)") + print(f"Timed runs per engine: {N_TIMED}\n") + + def show(label, s): + rtf = s["mean_ms"] / audio_ms + print(f"{label}") + print(f" cold load: {s['load_ms']:7.0f} ms") + print(f" warm-up: {s['warmup_ms']:7.0f} ms") + print(f" mean: {s['mean_ms']:7.1f} ms (RTF {rtf:.3f}x audio length)") + print(f" median: {s['median_ms']:7.1f} ms") + print(f" min / max: {s['min_ms']:7.1f} / {s['max_ms']:.1f} ms") + print(f" stdev: {s['stdev_ms']:7.1f} ms") + print(f" text: '{s['text']}'\n") + + show("Whisper CPU int8 base.en (baseline)", whisper_stats) + show("Parakeet CUDA FP16 0.6b-v2 (v2.3.0)", parakeet_stats) + + delta_mean = parakeet_stats["mean_ms"] - whisper_stats["mean_ms"] + speedup = whisper_stats["mean_ms"] / parakeet_stats["mean_ms"] + direction = "faster" if delta_mean < 0 else "slower" + print("--- delta (Parakeet vs Whisper baseline) ---") + print(f" mean Δ: {delta_mean:+.1f} ms → {abs(delta_mean):.0f} ms {direction}") + print(f" speedup: {speedup:.2f}x (mean Whisper / mean Parakeet)") + print(f" load Δ: {parakeet_stats['load_ms'] - whisper_stats['load_ms']:+.0f} ms") + + +if __name__ == "__main__": + main() diff --git a/tests/test_probe.py b/tests/test_probe.py index ea4ebd3..42c07e4 100644 --- a/tests/test_probe.py +++ b/tests/test_probe.py @@ -68,7 +68,7 @@ def test_nothing_detected(self): class TestDetectRobot(unittest.TestCase): def test_missing_config(self): with mock.patch.object(Path, "exists", return_value=False): - r = detect.detect_robot(robot_env_path="/nonexistent/robot.env") + r = detect.detect_robot(robot_env_path="/nonexistent/.env") self.assertEqual(r.value, "missing-config") def test_reachable(self):