""" Parakeet TDT Speech-to-Text Handler Supports NVIDIA Parakeet TDT model for high-quality multilingual ASR. - On Apple Silicon (MPS): Uses mlx-audio with mlx-community/parakeet-tdt-0.6b-v3 - On CUDA/CPU: Uses nano-parakeet (pure PyTorch) with nvidia/parakeet-tdt-0.6b-v3 Model supports 25 European languages with automatic language detection. """ from __future__ import annotations import logging from contextlib import contextmanager from sys import platform from threading import Lock from time import perf_counter from typing import Any, Iterator, Optional import numpy as np from rich.console import Console from rich.text import Text from speech_to_speech.pipeline.handler_types import STTIn, STTOut from speech_to_speech.pipeline.messages import PartialTranscription, Transcription from speech_to_speech.STT.base_stt_handler import BaseSTTHandler from speech_to_speech.STT.smart_progressive_streaming import PartialTranscription as ProgressiveStreamPartial from speech_to_speech.utils.mlx_lock import MLXLockContext try: from lingua import Language, LanguageDetectorBuilder LINGUA_AVAILABLE = True except ImportError: LINGUA_AVAILABLE = False logger = logging.getLogger(__name__) console = Console() # Parakeet TDT v3 supports 25 European languages SUPPORTED_LANGUAGES = [ "en", "de", "fr", "es", "it", "pt", "nl", "pl", "ru", "uk", "cs", "sk", "hu", "ro", "bg", "hr", "sl", "sr", "da", "no", "sv", "fi", "et", "lv", "lt", ] # Lingua uses "nb" (Bokmål) for Norwegian instead of "no" _LINGUA_CODE_MAP = {"no": "nb"} if LINGUA_AVAILABLE: _lingua_iso_to_code = { lang.iso_code_639_1.name.lower(): lang for lang in Language.all() if lang.iso_code_639_1 is not None } _lingua_languages = [ _lingua_iso_to_code[_LINGUA_CODE_MAP.get(code, code)] for code in SUPPORTED_LANGUAGES if _LINGUA_CODE_MAP.get(code, code) in _lingua_iso_to_code ] def _build_lingua_detector(): # Preloading can take multiple seconds on some hardware, including the # deployed server. Pay that cost at startup instead of on the first user # request, where it would look like slow STT. return LanguageDetectorBuilder.from_languages(*_lingua_languages).with_preloaded_language_models().build() _lingua_detector = _build_lingua_detector() class ParakeetTDTSTTHandler(BaseSTTHandler): """ Handles Speech-to-Text using NVIDIA Parakeet TDT model. On Apple Silicon (MPS): Uses mlx-audio with the MLX-converted model. On CUDA/CPU: Uses nano-parakeet (pure PyTorch) for NeMo-free inference. Parakeet TDT 0.6B v3 is a 600M parameter multilingual ASR model supporting 25 European languages with automatic language detection. """ def setup( self, model_name: Optional[str] = None, device: str = "auto", compute_type: str = "float16", language: Optional[str] = None, gen_kwargs: dict[str, Any] = {}, enable_live_transcription: bool = False, live_transcription_update_interval: float = 0.5, ) -> None: """ Initialize the Parakeet TDT model. Args: model_name: Model identifier. Defaults are: - MPS: "mlx-community/parakeet-tdt-0.6b-v3" - CUDA/CPU: "nvidia/parakeet-tdt-0.6b-v3" device: Device to use ("auto", "cuda", "mps", "cpu") compute_type: Compute precision ("float16", "float32") language: Target language code (optional, model auto-detects) gen_kwargs: Additional generation kwargs """ self.gen_kwargs = gen_kwargs self.start_language = language self.last_language = language if language else "en" self.enable_live_transcription = enable_live_transcription self.live_transcription_update_interval = live_transcription_update_interval self.compute_lock = Lock() self.sample_rate = 16000 # Determine device if device == "auto": if platform == "darwin": self.device = "mps" else: import torch self.device = "cuda" if torch.cuda.is_available() else "cpu" else: self.device = device # Set default model based on device if model_name is None: if self.device == "mps": model_name = "mlx-community/parakeet-tdt-0.6b-v3" else: model_name = "nvidia/parakeet-tdt-0.6b-v3" self.model_name = model_name self.compute_type = compute_type logger.info(f"Loading Parakeet TDT model: {model_name} on {self.device}") if self.device == "mps": self._setup_mlx(model_name) else: self._setup_nano_parakeet(model_name) # Setup streaming handler if live transcription is enabled self.streaming_handler = None if self.enable_live_transcription: from speech_to_speech.STT.smart_progressive_streaming import ( SmartProgressiveStreamingHandler, ) self.streaming_handler = SmartProgressiveStreamingHandler( self.model, emission_interval=self.live_transcription_update_interval, max_window_size=15.0, sentence_buffer=2.0, ) self.processing_final = False # Track if we're processing final audio logger.info(f"Live transcription enabled for Parakeet TDT ({self.backend})") self._live_transcription_active = False self._live_turn_key: tuple[str | None, int | None] | None = None self.warmup() def _setup_mlx(self, model_name: str) -> None: """Setup for Apple Silicon using mlx-audio.""" try: from mlx_audio.stt.generate import load_model self.backend = "mlx" self.model = load_model(model_name) logger.info("MLX Audio Parakeet model loaded successfully") except ImportError as e: raise ImportError( "mlx-audio is required for Parakeet TDT on Apple Silicon. Install with: pip install mlx-audio" ) from e def _setup_nano_parakeet(self, model_name: str) -> None: """Setup for CUDA/CPU using nano-parakeet.""" try: import torch from nano_parakeet import from_pretrained self.backend = "nano_parakeet" if self.device == "cuda" and not torch.cuda.is_available(): logger.warning("CUDA requested but not available. Falling back to CPU for nano-parakeet.") self.device = "cpu" self.model = from_pretrained(model_name=model_name, device=self.device) logger.info(f"nano-parakeet model loaded successfully on {self.device}") except ImportError as e: raise ImportError( "nano-parakeet is required for Parakeet TDT on CUDA/CPU. Install with: pip install nano-parakeet" ) from e def warmup(self) -> None: """Warm up the model with a dummy input.""" logger.info(f"Warming up {self.__class__.__name__}") # Create 1 second of silence at 16kHz dummy_audio = np.zeros(16000, dtype=np.float32) try: if self.backend == "mlx": import mlx.core as mx # Convert to mx.array and call decode_chunk directly audio_mx = mx.array(dummy_audio, dtype=mx.float32) _ = self.model.decode_chunk(audio_mx, verbose=False) elif self.backend == "nano_parakeet": _ = self.model.transcribe(dummy_audio) else: _ = self.model.transcribe([dummy_audio], batch_size=1, verbose=False) logger.info("Model warmed up and ready") except Exception as e: logger.warning(f"Warmup failed: {e}") def process(self, vad_audio: STTIn) -> Iterator[STTOut]: """ Process audio and generate transcription. Yields: :class:`PartialTranscription` or :class:`Transcription` """ process_start_s = perf_counter() is_progressive = vad_audio.mode == "progressive" audio_input = vad_audio.audio # Ensure audio is float32 numpy array if not isinstance(audio_input, np.ndarray): audio_input = np.array(audio_input, dtype=np.float32) else: audio_input = audio_input.astype(np.float32) audio_duration_s = len(audio_input) / getattr(self, "sample_rate", 16000) item_age_s = self._item_age_s(vad_audio) self._prepare_live_transcription_turn(vad_audio.turn_id, vad_audio.turn_revision) # Handle progressive updates: yield tagged partial for TranscriptionNotifier if self.enable_live_transcription and is_progressive: # Ignore progressive updates if we're already processing final audio if self.processing_final: logger.debug("Skipping stale progressive update (final audio already received)") return # Try to acquire lock with short timeout - skip if busy lock_scope_start_s = perf_counter() with self._compute_lock_context(handler_name="ParakeetSTT-Progressive", timeout=0.01) as acquired: if acquired: try: inference_start_s = perf_counter() progressive_text = self._show_progressive_transcription(audio_input) inference_s = perf_counter() - inference_start_s if inference_s >= 0.25: logger.info( "Parakeet progressive STT timing turn=%s rev=%s audio=%.3fs age=%.3fs " "lock_scope=%.3fs inference=%.3fs chars=%d", vad_audio.turn_id, vad_audio.turn_revision, audio_duration_s, item_age_s, perf_counter() - lock_scope_start_s, inference_s, len(progressive_text), ) if progressive_text: yield PartialTranscription( text=progressive_text, turn_id=vad_audio.turn_id, turn_revision=vad_audio.turn_revision, ) return except Exception as e: logger.debug(f"Progressive transcription failed: {e}") else: logger.debug("Skipping progressive update (compute busy)") return # Handle final transcription (send to LLM) logger.info( "Parakeet final STT start turn=%s rev=%s audio=%.3fs age=%.3fs", vad_audio.turn_id, vad_audio.turn_revision, audio_duration_s, item_age_s, ) inference_s = 0.0 lock_scope_s = 0.0 try: if self.enable_live_transcription: # Mark that we're processing final audio (ignore stale progressive updates) self.processing_final = True # Acquire lock with longer timeout for final transcription lock_scope_start_s = perf_counter() with self._compute_lock_context(handler_name="ParakeetSTT-Final", timeout=5.0) as acquired: lock_scope_s = perf_counter() - lock_scope_start_s if not acquired: logger.error("Failed to acquire compute lock for final transcription") pred_text = "" language_code = self.last_language else: inference_start_s = perf_counter() if self.backend == "mlx": pred_text, language_code = self._process_mlx_final(audio_input) else: pred_text, language_code = self._process_nano_parakeet(audio_input) inference_s = perf_counter() - inference_start_s lock_scope_s = perf_counter() - lock_scope_start_s # Validate and update language if language_code and language_code in SUPPORTED_LANGUAGES: self.last_language = language_code else: language_code = self.last_language except Exception as e: logger.error(f"Parakeet TDT inference failed: {e}") pred_text = "" language_code = self.last_language total_s = perf_counter() - process_start_s logger.info( "Parakeet final STT done turn=%s rev=%s total=%.3fs lock_scope=%.3fs inference=%.3fs chars=%d", vad_audio.turn_id, vad_audio.turn_revision, total_s, lock_scope_s, inference_s, len(pred_text), ) logger.debug("Finished Parakeet TDT inference") self._clear_live_transcription_line() if pred_text.strip(): console.print(f"[yellow]USER: {pred_text.strip()}") if language_code: console.print(f"[dim]Language: {language_code}[/dim]") # Reset per-utterance live transcription state only after final STT # completes. The streaming handler carries fixed sentence timing within # an utterance, and stale timing must not leak into the next turn. if self.enable_live_transcription: self.processing_final = False self._reset_live_transcription_state(clear_turn=True) yield Transcription( text=pred_text, language_code=language_code, turn_id=vad_audio.turn_id, turn_revision=vad_audio.turn_revision, speech_stopped_at_s=vad_audio.created_at_s, ) @property def timing_log_level(self) -> int: return logging.INFO def should_log_timing(self, output: STTOut) -> bool: return isinstance(output, Transcription) and self.last_time > self.min_time_to_debug def _detect_language_from_text(self, text: str) -> Optional[str]: """ Detect language from transcribed text using lingua-py. Args: text: Transcribed text string Returns: Detected language code or None if detection fails """ if not LINGUA_AVAILABLE: logger.warning("lingua-py not available, cannot detect language from text") return None # Skip very short utterances where language ID is still too noisy. if not text or len(text.strip()) < 20: return None detected = _lingua_detector.detect_language_of(text) if detected is None: return None code = detected.iso_code_639_1.name.lower() # Map back lingua-specific codes to our supported codes return {v: k for k, v in _LINGUA_CODE_MAP.items()}.get(code, code) @contextmanager def _compute_lock_context(self, handler_name: str, timeout: float) -> Iterator[bool]: if self.backend == "mlx": with MLXLockContext(handler_name=handler_name, timeout=timeout) as acquired: yield acquired return lock_start_s = perf_counter() acquired = self.compute_lock.acquire(timeout=timeout) wait_s = perf_counter() - lock_start_s hold_start_s: float | None = None if acquired: if wait_s >= 0.25: logger.info("%s: compute lock acquired after %.2fs", handler_name, wait_s) else: logger.debug("%s: compute lock acquired after %.3fs", handler_name, wait_s) hold_start_s = perf_counter() else: logger.warning("%s: Failed to acquire compute lock after %.3fs (timeout=%s)", handler_name, wait_s, timeout) try: yield acquired finally: if acquired: assert hold_start_s is not None self.compute_lock.release() hold_s = perf_counter() - hold_start_s if hold_s >= 0.25: logger.info("%s: compute lock released after holding %.2fs", handler_name, hold_s) else: logger.debug("%s: compute lock released after holding %.3fs", handler_name, hold_s) def _show_progressive_transcription(self, audio_input: np.ndarray) -> str: """Run progressive transcription, print to console, and return the text.""" result = self.streaming_handler.transcribe_incremental(audio_input) rich_text = Text() if result.fixed_text: rich_text.append("Live: ", style="dim") rich_text.append(result.fixed_text, style="yellow") if result.active_text: rich_text.append(" ", style="dim") if result.active_text: if not result.fixed_text: rich_text.append("Live: ", style="dim") rich_text.append(result.active_text, style="cyan dim") progressive_text = self._build_progressive_text(result) if progressive_text: self._print_live_transcription(rich_text, progressive_text) return progressive_text def _print_live_transcription(self, rich_text: Text, progressive_text: str) -> None: is_terminal = bool(getattr(console, "is_terminal", False)) if is_terminal: self._write_live_control("\r\x1b[2K") if rich_text: console.print(self._truncate_live_transcription(rich_text), end="") else: fallback = Text("Live: ", style="dim") fallback.append(progressive_text, style="cyan dim") console.print(self._truncate_live_transcription(fallback), end="") self._write_live_control("\r") self._live_transcription_active = True return if rich_text: console.print(rich_text) else: console.print(f"[dim]Live: [/dim]{progressive_text}") def _clear_live_transcription_line(self) -> None: if not getattr(self, "_live_transcription_active", False): return self._write_live_control("\r\x1b[2K") self._live_transcription_active = False def _write_live_control(self, sequence: str) -> None: file = getattr(console, "file", None) if file is None: return file.write(sequence) file.flush() def _truncate_live_transcription(self, text: Text) -> Text: text = text.copy() width = getattr(console, "width", 80) try: max_width = max(1, int(width) - 1) except (TypeError, ValueError): max_width = 79 text.truncate(max_width, overflow="ellipsis") return text def _prepare_live_transcription_turn(self, turn_id: str | None, turn_revision: int | None) -> None: if not self.enable_live_transcription: return turn_key = (turn_id, turn_revision) if getattr(self, "_live_turn_key", None) == turn_key: return self._reset_live_transcription_state(clear_turn=False) self._live_turn_key = turn_key def _reset_live_transcription_state(self, clear_turn: bool) -> None: self._clear_live_transcription_line() streaming_handler = getattr(self, "streaming_handler", None) if streaming_handler is not None: streaming_handler.reset() if clear_turn: self._live_turn_key = None def _build_progressive_text(self, result: ProgressiveStreamPartial) -> str: parts = [] if result.fixed_text: parts.append(result.fixed_text.strip()) if result.active_text: parts.append(result.active_text.strip()) return " ".join(part for part in parts if part).strip() def _process_mlx_final(self, audio_input: np.ndarray) -> tuple[str, str]: """Process final audio using MLX backend with streaming handler.""" # If we have fixed sentences from progressive updates, only transcribe the new part if ( self.streaming_handler is not None and hasattr(self.streaming_handler, "fixed_sentences") and self.streaming_handler.fixed_sentences ): self._clear_live_transcription_line() # Get fixed text from previous progressive updates fixed_text = " ".join(self.streaming_handler.fixed_sentences).strip() # Calculate where fixed part ends in audio fixed_end_time = self.streaming_handler.fixed_end_time sample_rate = 16000 fixed_end_sample = int(fixed_end_time * sample_rate) if fixed_end_sample > len(audio_input): logger.warning( "Ignoring stale progressive fixed text: fixed_end_sample=%d exceeds final audio samples=%d", fixed_end_sample, len(audio_input), ) pred_text, language_code = self._process_mlx(audio_input) return pred_text, language_code # Only transcribe the part after fixed sentences if fixed_end_sample < len(audio_input): remaining_audio = audio_input[fixed_end_sample:] import mlx.core as mx audio_mx = mx.array(remaining_audio, dtype=mx.float32) result = self.model.decode_chunk(audio_mx, verbose=False) if hasattr(result, "text"): new_text = result.text.strip() else: new_text = str(result).strip() # Combine fixed + new pred_text = f"{fixed_text} {new_text}".strip() if new_text else fixed_text else: # All audio already transcribed in progressive updates pred_text = fixed_text # Reset streaming handler for next utterance self.streaming_handler.reset() else: # No progressive updates, transcribe everything pred_text, language_code = self._process_mlx(audio_input) return pred_text, language_code # Determine language if self.start_language and self.start_language != "auto": language_code = self.start_language else: detected_lang = self._detect_language_from_text(pred_text) if detected_lang: language_code = detected_lang else: language_code = self.last_language return pred_text, language_code def _process_mlx(self, audio_input: np.ndarray) -> tuple[str, str]: """Process audio using MLX backend.""" import mlx.core as mx # Convert numpy array to mx.array audio_mx = mx.array(audio_input, dtype=mx.float32) # Call decode_chunk directly with the audio array result = self.model.decode_chunk(audio_mx, verbose=False) # Extract text from result if hasattr(result, "text"): pred_text = result.text.strip() else: pred_text = str(result).strip() # Determine language: # 1. Use fixed language if specified by user # 2. Try to detect from transcribed text using langdetect # 3. Fall back to last known language if self.start_language and self.start_language != "auto": language_code = self.start_language else: # Detect language from transcribed text detected_lang = self._detect_language_from_text(pred_text) if detected_lang: language_code = detected_lang else: language_code = self.last_language return pred_text, language_code def _process_nano_parakeet(self, audio_input: np.ndarray) -> tuple[str, str]: """Process audio using nano-parakeet backend.""" pred_text = self.model.transcribe(audio_input).strip() if self.start_language and self.start_language != "auto": language_code = self.start_language else: detected_lang = self._detect_language_from_text(pred_text) if detected_lang: language_code = detected_lang else: language_code = self.last_language return pred_text, language_code def cleanup(self) -> None: """Clean up model resources.""" logger.info(f"Cleaning up {self.__class__.__name__}") if hasattr(self, "model"): del self.model def on_session_end(self) -> None: super().on_session_end() self.last_language = self.start_language if self.start_language else "en" if self.enable_live_transcription: self.processing_final = False self._reset_live_transcription_state(clear_turn=True) logger.debug("Parakeet TDT session state reset")