"""AlignAtt streaming transcription for Whisper. This module implements low-latency streaming transcription using the AlignAtt algorithm, which monitors cross-attention weights to determine when decoded tokens are stable enough to emit. Reference: https://arxiv.org/abs/2211.00895 (SimulMT with AlignAtt) """ from dataclasses import dataclass from typing import List import mlx.core as mx @dataclass class StreamingConfig: """Configuration for AlignAtt streaming transcription. Attributes: frame_threshold: Number of frames from audio end before emitting tokens. Lower = faster but may cut words. Default 25 (~0.5s at 50 tokens/sec). min_chunk_duration: Minimum audio duration (seconds) before processing. emit_partial: Whether to emit partial (uncommitted) results. """ frame_threshold: int = 25 min_chunk_duration: float = 0.5 emit_partial: bool = True @dataclass class StreamingResult: """Result from streaming transcription. Attributes: text: Decoded text for this emission. tokens: Token IDs that were decoded. is_final: True if this is a final (committed) result, False if partial. start_time: Start timestamp in seconds. end_time: End timestamp in seconds. progress: Progress from 0.0 to 1.0 (percentage of audio processed). audio_position: Current position in audio (seconds). audio_duration: Total audio duration (seconds). """ text: str tokens: List[int] is_final: bool start_time: float end_time: float progress: float = 0.0 audio_position: float = 0.0 audio_duration: float = 0.0 def get_most_attended_frame(cross_qk: List[mx.array], alignment_heads: mx.array) -> int: """Extract the most attended audio frame from cross-attention weights. Uses the same alignment heads as word-level timestamps (timing.py). Args: cross_qk: List of cross-attention weights per layer. Each element shape: [batch, n_heads, seq_len, audio_frames] alignment_heads: Array of [layer, head] pairs to use for alignment. Returns: Frame index that receives highest average attention for the last token. """ weights = mx.stack( [cross_qk[layer][0, head, -1, :] for layer, head in alignment_heads.tolist()] ) avg_attention = weights.mean(axis=0) most_attended = int(mx.argmax(avg_attention).item()) return most_attended def should_emit( most_attended_frame: int, content_frames: int, config: StreamingConfig ) -> bool: """Determine if tokens should be emitted based on AlignAtt algorithm. The AlignAtt insight: when cross-attention focuses near the end of available audio, the model has "seen enough" to emit stable tokens. Args: most_attended_frame: Frame index with highest attention for last token. content_frames: Total number of audio frames available. config: Streaming configuration with frame_threshold. Returns: True if tokens should be emitted (attention near audio end). """ distance_to_end = content_frames - most_attended_frame return distance_to_end <= config.frame_threshold class StreamingDecoder: """AlignAtt-based streaming decoder for Whisper. This decoder accumulates audio over time and monitors cross-attention weights during token generation to determine when tokens are stable enough to emit, enabling low-latency streaming transcription. The key insight is that we accumulate mel spectrograms across chunks rather than processing each chunk independently. This allows the model to maintain context as more audio arrives. """ def __init__( self, model, config: StreamingConfig = None, language: str = None, task: str = "transcribe", ): """Initialize streaming decoder. Args: model: Loaded Whisper model instance. config: Streaming configuration. Uses defaults if None. language: Language code (e.g., 'en'). Auto-detected if None. task: "transcribe" or "translate" (to English). """ from mlx_audio.stt.models.whisper.decoding import ( Inference, SuppressBlank, SuppressTokens, get_suppress_tokens, ) self.model = model self.config = config or StreamingConfig() self.inference = Inference(model) self.tokenizer = model.get_tokenizer(language=language or "en", task=task) self._emitted_tokens = [] self._pending_tokens = [] self._accumulated_mel = None # Use sot_sequence with notimestamps for text-only output self._sot_sequence = self.tokenizer.sot_sequence_including_notimestamps # Use shared helper for suppress tokens suppress_tokens = set(get_suppress_tokens(self.tokenizer)) # Don't suppress notimestamps since we need it in sot sequence suppress_tokens.discard(self.tokenizer.no_timestamps) # Initialize logit filters self._sample_begin = len(self._sot_sequence) self._logit_filters = [ SuppressBlank(self.tokenizer, self._sample_begin, model.dims.n_vocab), SuppressTokens(list(suppress_tokens), model.dims.n_vocab), ] def reset(self): """Reset decoder state for new transcription session.""" self.inference.reset() self._emitted_tokens = [] self._pending_tokens = [] self._accumulated_mel = None def decode_chunk(self, mel: mx.array, is_last: bool = False) -> StreamingResult: """Decode an audio chunk using AlignAtt streaming. Audio is accumulated across chunks to maintain context. The AlignAtt algorithm determines when tokens are stable enough to emit based on cross-attention weights. Args: mel: Mel spectrogram of audio chunk, shape [n_frames, n_mels]. is_last: True if this is the final chunk (emit all remaining). Returns: StreamingResult with decoded text and metadata. """ from mlx_audio.stt.models.whisper.audio import ( N_FRAMES, TOKENS_PER_SECOND, pad_or_trim, ) # Accumulate mel spectrogram across chunks if self._accumulated_mel is None: self._accumulated_mel = mel else: self._accumulated_mel = mx.concatenate([self._accumulated_mel, mel], axis=0) # Trim to max 30 seconds if accumulated too much (sliding window) if self._accumulated_mel.shape[0] > N_FRAMES: self._accumulated_mel = self._accumulated_mel[-N_FRAMES:] # Reset KV cache for fresh decoding of accumulated audio self.inference.reset() # Pad mel to required size for encoder mel_padded = pad_or_trim(self._accumulated_mel, N_FRAMES, axis=-2) # Add batch dimension if needed if mel_padded.ndim == 2: mel_padded = mel_padded[None, :] # Encode audio audio_features = self.model.encoder(mel_padded.astype(self.model.dtype)) content_frames = self._accumulated_mel.shape[0] // 2 # Encoder has stride 2 # Initialize tokens with SOT sequence (including notimestamps) tokens = mx.array([list(self._sot_sequence)]) # Decode loop with AlignAtt # First iteration: pass all initial tokens; subsequent: only last token first_iteration = True for _ in range(self.model.dims.n_text_ctx // 2): inputs = tokens if first_iteration else tokens[:, -1:] first_iteration = False logits, cross_qk = self.inference.logits_with_cross_qk( inputs, audio_features, ) # Apply logit filters to suppress non-speech and special tokens filtered_logits = logits[:, -1] for logit_filter in self._logit_filters: filtered_logits = logit_filter.apply(filtered_logits[None, :], tokens)[ 0 ] # Get next token from filtered logits next_token = int(mx.argmax(filtered_logits, axis=-1).item()) if next_token == self.tokenizer.eot: break tokens = mx.concatenate([tokens, mx.array([[next_token]])], axis=-1) # Check AlignAtt condition (only if we have alignment heads) if ( hasattr(self.model, "alignment_heads") and self.model.alignment_heads is not None ): most_attended = get_most_attended_frame( cross_qk, self.model.alignment_heads ) threshold = 4 if is_last else self.config.frame_threshold if should_emit( most_attended, content_frames, StreamingConfig(frame_threshold=threshold), ): break # Extract text tokens (excluding special tokens) text_tokens = [ t for t in tokens[0].tolist() if t < self.tokenizer.eot and t not in self.tokenizer.sot_sequence ] # Calculate what's new since last emission new_tokens = text_tokens[len(self._emitted_tokens) :] self._emitted_tokens = text_tokens emit_start_time = ( len(self._emitted_tokens) - len(new_tokens) ) / TOKENS_PER_SECOND emit_end_time = len(self._emitted_tokens) / TOKENS_PER_SECOND return StreamingResult( text=self.tokenizer.decode(new_tokens), tokens=new_tokens, is_final=is_last, start_time=emit_start_time, end_time=emit_end_time, )