""" MLX whisper AlignAtt streaming decoder """ import logging from time import time from typing import Any, List, Optional, Tuple import mlx.core as mx import numpy as np from mlx_whisper.audio import log_mel_spectrogram as mlx_log_mel_spectrogram from mlx_whisper.transcribe import pad_or_trim as mlx_pad_or_trim from whisperlivekit.timed_objects import ASRToken from whisperlivekit.whisper import DecodingOptions, tokenizer from whisperlivekit.whisper.audio import N_FRAMES, N_SAMPLES, TOKENS_PER_SECOND from ..config import AlignAttConfig from .decoder_state import MLXDecoderState from .decoders import MLXBeamSearchDecoder, MLXGreedyDecoder, MLXInference DEC_PAD = 50257 logger = logging.getLogger(__name__) class MLXTokenBuffer: #should try to make it heritate from classic simul whisper class """Token buffer for MLX-based decoding.""" def __init__(self, text="", tokenizer=None, prefix_token_ids=None): self.text = text self.prefix_token_ids = prefix_token_ids or [] self.tokenizer = tokenizer self.pending_token_ids = [] def as_token_ids(self, tokenizer=None): if tokenizer is None: tokenizer = self.tokenizer if tokenizer is None: raise ValueError("Tokenizer is not set.") return self.prefix_token_ids + tokenizer.encode(self.text) def as_mlx_array(self) -> mx.array: """Return tokens as MLX array.""" tok_ids = self.as_token_ids() return mx.array([tok_ids], dtype=mx.int32) def as_mlx_array_beam(self, beam: int) -> mx.array: """Return tokens as MLX array repeated for beam search.""" t = self.as_mlx_array() return mx.repeat(t, beam, axis=0) def as_text(self): return self.text @staticmethod def empty(*a, **kw): return MLXTokenBuffer(*a, **kw) @staticmethod def from_text(text, *a, **kw): return MLXTokenBuffer(*a, text=text, **kw) def is_empty(self): return self.text is None or self.text == "" def trim_words(self, num=1, after=0): """Trim words from the beginning of the context.""" tokenizer = self.tokenizer assert tokenizer is not None, "Tokenizer is not set." ids = tokenizer.encode(self.text[after:]) words, wids = self.tokenizer.split_to_word_tokens(ids) if not words: return 0 self.text = self.text[:after] + "".join(words[num:]) return sum(len(wi) for wi in wids[:num]) def append_token_ids(self, token_ids): """Append token IDs to the buffer, handling incomplete UTF-8.""" tokenizer = self.tokenizer assert tokenizer is not None, "Tokenizer is not set." all_tokens = self.pending_token_ids + token_ids decoded = tokenizer.decode(all_tokens) replacement_char = "\ufffd" if replacement_char in decoded: if len(all_tokens) > 1: decoded_partial = tokenizer.decode(all_tokens[:-1]) if replacement_char not in decoded_partial: self.text += decoded_partial self.pending_token_ids = [all_tokens[-1]] else: self.pending_token_ids = all_tokens else: self.pending_token_ids = all_tokens else: self.text += decoded self.pending_token_ids = [] def mlx_median_filter(x: mx.array, filter_width: int) -> mx.array: """ Apply median filter along the last axis. Args: x: Input array of shape (..., T) filter_width: Width of the median filter (should be odd) Returns: Filtered array of same shape """ if filter_width <= 1: return x pad_width = filter_width // 2 shape = x.shape left_pad = mx.repeat(x[..., :1], pad_width, axis=-1) right_pad = mx.repeat(x[..., -1:], pad_width, axis=-1) x_padded = mx.concatenate([left_pad, x, right_pad], axis=-1) result_shape = list(shape) result = [] for i in range(shape[-1]): window = x_padded[..., i:i + filter_width] sorted_window = mx.sort(window, axis=-1) median_val = sorted_window[..., filter_width // 2:filter_width // 2 + 1] result.append(median_val) return mx.concatenate(result, axis=-1) class MLXAlignAtt: """ MLX-native Alignment-based Attention decoder for SimulStreaming. This class runs entirely on MLX, with no PyTorch dependencies for inference. """ @property def speaker(self): return self.state.speaker @speaker.setter def speaker(self, value): self.state.speaker = value @property def global_time_offset(self): return self.state.global_time_offset @global_time_offset.setter def global_time_offset(self, value): self.state.global_time_offset = value def __init__( self, cfg: AlignAttConfig, mlx_model: Any, ) -> None: """ Initialize MLX AlignAtt decoder. Args: cfg: AlignAtt configuration mlx_model: MLX Whisper model (full model, not just encoder) """ self.model = mlx_model self.cfg = cfg logger.info(f"MLX Model dimensions: {self.model.dims}") self.decode_options = DecodingOptions( language=cfg.language, without_timestamps=True, task=cfg.task ) self.tokenizer_is_multilingual = cfg.tokenizer_is_multilingual self.max_text_len = self.model.dims.n_text_ctx self.num_decoder_layers = len(self.model.decoder.blocks) if self.cfg.max_context_tokens is None: self.max_context_tokens = self.max_text_len else: self.max_context_tokens = self.cfg.max_context_tokens # Initialize per-session state self.state = MLXDecoderState() self._init_state(cfg) def _init_state(self, cfg: AlignAttConfig): """Initialize the per-session decoder state.""" self.create_tokenizer(cfg.language if cfg.language != "auto" else None) self.state.tokenizer = self.tokenizer self.state.detected_language = cfg.language if cfg.language != "auto" else None self.state.global_time_offset = 0.0 self.state.last_attend_frame = -cfg.rewind_threshold self.state.speaker = -1 if cfg.cif_ckpt_path is None or not cfg.cif_ckpt_path: if cfg.never_fire: self.state.never_fire = True self.state.always_fire = False else: self.state.always_fire = True self.state.never_fire = False else: logger.warning("CIF checkpoint provided but MLX CIF not implemented. Using always_fire=True") self.state.always_fire = True self.state.never_fire = cfg.never_fire self._build_alignment_source() suppress_tokens = [ self.tokenizer.transcribe, self.tokenizer.translate, self.tokenizer.sot, self.tokenizer.sot_prev, self.tokenizer.sot_lm, self.tokenizer.no_timestamps, ] + list(self.tokenizer.all_language_tokens) if self.tokenizer.no_speech is not None: suppress_tokens.append(self.tokenizer.no_speech) self.state.suppress_tokens = tuple(sorted(set(suppress_tokens))) logger.debug(f"Suppress tokens: {self.state.suppress_tokens}") self.init_tokens() self.init_context() self.state.decoder_type = cfg.decoder_type if cfg.decoder_type == "greedy": logger.info("Using MLX greedy decoder") self.state.token_decoder = MLXGreedyDecoder(0.0, self.tokenizer.eot) elif cfg.decoder_type == "beam": logger.info("Using MLX beam decoder") self.state.inference = MLXInference(self.model, self.state.initial_token_length) self.state.token_decoder = MLXBeamSearchDecoder( inference=self.state.inference, eot=self.tokenizer.eot, beam_size=cfg.beam_size ) def _build_alignment_source(self): """Build alignment source mapping from model's alignment_heads.""" self.state.align_source = {} self.state.num_align_heads = 0 alignment_heads = self.model.alignment_heads if alignment_heads is None: logger.warning("No alignment heads found in model") return if hasattr(alignment_heads, 'tolist'): heads_list = alignment_heads.tolist() else: heads_list = np.array(alignment_heads).tolist() for layer_rank, head_id in heads_list: layer_rank = int(layer_rank) head_id = int(head_id) heads = self.state.align_source.get(layer_rank, []) heads.append((self.state.num_align_heads, head_id)) self.state.align_source[layer_rank] = heads self.state.num_align_heads += 1 def warmup(self, audio: np.ndarray): """Warmup the model with sample audio.""" try: self.insert_audio(audio) self.infer(is_last=True) self.refresh_segment(complete=True) logger.info("MLX model warmed up successfully") except Exception as e: logger.exception(f"MLX model warmup failed: {e}") def create_tokenizer(self, language=None): """Create tokenizer for the given language.""" self.tokenizer = tokenizer.get_tokenizer( multilingual=self.tokenizer_is_multilingual, language=language, num_languages=self.model.num_languages, task=self.decode_options.task ) self.state.tokenizer = self.tokenizer def init_context(self): """Initialize context buffer.""" kw = { 'tokenizer': self.tokenizer, 'prefix_token_ids': [self.tokenizer.sot_prev] } self.state.context = MLXTokenBuffer.empty(**kw) if self.cfg.static_init_prompt is not None: self.state.context = MLXTokenBuffer.from_text(self.cfg.static_init_prompt, **kw) if self.cfg.init_prompt is not None: self.state.context.text += self.cfg.init_prompt def init_tokens(self): """Initialize token sequence.""" logger.debug(f"init tokens, {len(self.state.segments)}") self.state.initial_tokens = mx.array( [self.tokenizer.sot_sequence_including_notimestamps], dtype=mx.int32 ) self.state.initial_token_length = self.state.initial_tokens.shape[1] self.state.sot_index = self.tokenizer.sot_sequence.index(self.tokenizer.sot) logger.debug(f"init tokens after, {len(self.state.segments)}") self.state.tokens = [self.state.initial_tokens] def trim_context(self): """Trim context if too long.""" logger.info("Trimming context") c = len(self.state.context.as_token_ids()) - len(self.state.context.prefix_token_ids) logger.info(f"Context text: {self.state.context.as_text()}") l = sum(t.shape[1] for t in self.state.tokens) + c if self.cfg.static_init_prompt is None: after = 0 else: after = len(self.cfg.static_init_prompt) while c > self.max_context_tokens or l > self.max_text_len - 20: t = self.state.context.trim_words(after=after) l -= t c -= t logger.debug(f"len {l}, c {c}, max_context_tokens {self.max_context_tokens}") if t == 0: break logger.info(f"Context after trim: {self.state.context.text} (len: {l})") def refresh_segment(self, complete=False): """Refresh segment state.""" logger.debug("Refreshing segment:") self.init_tokens() self.state.last_attend_frame = -self.cfg.rewind_threshold self.state.cumulative_time_offset = 0.0 self.init_context() logger.debug(f"Context: {self.state.context}") if not complete and len(self.state.segments) > 2: self.state.segments = self.state.segments[-2:] else: logger.debug("removing all segments.") self.state.segments = [] self.state.log_segments += 1 self.state.pending_incomplete_tokens = [] def fire_at_boundary(self, chunked_encoder_feature: mx.array) -> bool: """Check if we should fire at word boundary (CIF-based).""" if self.state.always_fire: return True if self.state.never_fire: return False return True def _current_tokens(self) -> mx.array: """Get current token sequence for decoding.""" toks = self.state.tokens if toks[0].shape[0] == 1: toks[0] = mx.repeat(toks[0], self.cfg.beam_size, axis=0) if not self.state.context.is_empty(): context_toks = self.state.context.as_mlx_array_beam(self.cfg.beam_size) toks = [context_toks] + toks # Concatenate all tokens if len(toks) > 1: current_tokens = mx.concatenate(toks, axis=1) else: current_tokens = toks[0] logger.debug("debug print current_tokens:") self.debug_print_tokens(current_tokens) return current_tokens def debug_print_tokens(self, tokens: mx.array): """Debug print token sequences.""" tokens_np = np.array(tokens) for i in range(min(self.cfg.beam_size, tokens_np.shape[0])): logger.debug(self.tokenizer.decode_with_timestamps(tokens_np[i].tolist())) def segments_len(self) -> float: """Get total length of audio segments in seconds.""" return sum(s.shape[0] for s in self.state.segments) / 16000 def _apply_minseglen(self) -> bool: """Check if we have enough audio to process.""" segments_len = self.segments_len() if segments_len < self.cfg.audio_min_len: logger.debug("waiting for next segment") return False return True def insert_audio(self, segment: np.ndarray = None): """Insert audio segment into buffer.""" if segment is not None: if hasattr(segment, 'numpy'): segment = segment.numpy() self.state.segments.append(segment) removed_len = 0 segments_len = self.segments_len() while len(self.state.segments) > 1 and segments_len > self.cfg.audio_max_len: removed_len = self.state.segments[0].shape[0] / 16000 segments_len -= removed_len self.state.last_attend_frame -= int(TOKENS_PER_SECOND * removed_len) self.state.cumulative_time_offset += removed_len self.state.segments = self.state.segments[1:] logger.debug(f"remove segments: {len(self.state.segments)} {len(self.state.tokens)}, cumulative offset: {self.state.cumulative_time_offset:.2f}s") if len(self.state.tokens) > 1: # Convert MLX array to list for context token_list = np.array(self.state.tokens[1][0, :]).tolist() self.state.context.append_token_ids(token_list) self.state.tokens = [self.state.initial_tokens] + self.state.tokens[2:] return removed_len def _clean_cache(self): """Clean the kv_cache after each inference step.""" self.state.clean_cache() def _suppress_tokens(self, logits: mx.array) -> mx.array: """Apply token suppression to logits.""" if self.state.suppress_tokens: suppress_indices = mx.array(list(self.state.suppress_tokens), dtype=mx.int32) logits = logits.at[:, suppress_indices].add(-float('inf')) return logits def lang_id(self, encoder_features: mx.array) -> Tuple[mx.array, List[dict]]: """Language detection from encoder features.""" n_audio = encoder_features.shape[0] x = mx.array([[self.tokenizer.sot]] * n_audio, dtype=mx.int32) logits, _, _ = self.model.decoder(x, encoder_features, kv_cache=None) logits = logits[:, 0] mask = mx.ones(logits.shape[-1], dtype=mx.bool_) language_token_indices = mx.array(list(self.tokenizer.all_language_tokens), dtype=mx.int32) mask = mask.at[language_token_indices].add(False) logits = mx.where(mask, mx.array(-float('inf')), logits) language_tokens = mx.argmax(logits, axis=-1) language_token_probs = mx.softmax(logits, axis=-1) probs_np = np.array(language_token_probs) language_probs = [ { c: float(probs_np[i, j]) for j, c in zip(self.tokenizer.all_language_tokens, self.tokenizer.all_language_codes) } for i in range(n_audio) ] self._clean_cache() return language_tokens, language_probs def infer(self, is_last: bool = False) -> List[ASRToken]: """ Main inference method. Args: is_last: Whether this is the final chunk Returns: List of timestamped ASR tokens """ new_segment = True if len(self.state.segments) == 0: logger.debug("No segments, nothing to do") return [] if not self._apply_minseglen(): logger.debug(f"applied minseglen {self.cfg.audio_min_len} > {self.segments_len()}.") return [] if len(self.state.segments) > 1: input_segments = np.concatenate(self.state.segments, axis=0) else: input_segments = self.state.segments[0] beg_encode = time() mlx_mel_padded = mlx_log_mel_spectrogram( audio=input_segments, n_mels=self.model.dims.n_mels, padding=N_SAMPLES ) mlx_mel = mlx_pad_or_trim(mlx_mel_padded, N_FRAMES, axis=-2) encoder_feature = self.model.encoder(mlx_mel[None]) content_mel_len = int((mlx_mel_padded.shape[0] - mlx_mel.shape[0]) / 2) mx.eval(encoder_feature) end_encode = time() logger.debug(f'MLX Encoder duration: {end_encode - beg_encode:.3f}s') if self.cfg.language == "auto" and self.state.detected_language is None and self.state.first_timestamp: seconds_since_start = self.segments_len() - self.state.first_timestamp if seconds_since_start >= 2.0: language_tokens, language_probs = self.lang_id(encoder_feature) top_lan, p = max(language_probs[0].items(), key=lambda x: x[1]) print(f"Detected language: {top_lan} with p={p:.4f}") self.create_tokenizer(top_lan) self.state.last_attend_frame = -self.cfg.rewind_threshold self.state.cumulative_time_offset = 0.0 self.init_tokens() self.init_context() self.state.detected_language = top_lan logger.info(f"Tokenizer language: {self.tokenizer.language}") self.trim_context() current_tokens = self._current_tokens() fire_detected = self.fire_at_boundary(encoder_feature[:, :content_mel_len, :]) sum_logprobs = mx.zeros((self.cfg.beam_size,), dtype=mx.float32) completed = False attn_of_alignment_heads = None most_attended_frame = None token_len_before_decoding = current_tokens.shape[1] l_absolute_timestamps = [] accumulated_cross_attns = [] audio_duration_s = self.segments_len() # ~15 text tokens/s is a generous upper bound for speech; TOKENS_PER_SECOND (50) # is the mel-frame rate and was causing 10-40x over-allocation on repetition loops. max_tokens_per_chunk = max(50, int(audio_duration_s * 15 * 1.5)) tokens_produced_this_chunk = 0 while not completed and current_tokens.shape[1] < self.max_text_len: tokens_produced_this_chunk += 1 if tokens_produced_this_chunk > max_tokens_per_chunk: logger.warning(f"[Loop Detection] Too many tokens ({tokens_produced_this_chunk}) for {audio_duration_s:.2f}s audio. Breaking.") current_tokens = current_tokens[:, :token_len_before_decoding] break if new_segment: tokens_for_logits = current_tokens else: tokens_for_logits = current_tokens[:, -1:] if self.state.decoder_type == "greedy": logits, self.state.kv_cache, cross_qk = self.model.decoder( tokens_for_logits, encoder_feature, kv_cache=self.state.kv_cache ) else: logits, cross_qk = self.state.inference.logits(tokens_for_logits, encoder_feature) mx.eval(logits) accumulated_cross_attns.append(cross_qk) if len(accumulated_cross_attns) > 16: accumulated_cross_attns = accumulated_cross_attns[-16:] if new_segment and self.tokenizer.no_speech is not None: probs_at_sot = mx.softmax(logits[:, self.state.sot_index, :], axis=-1) no_speech_probs = np.array(probs_at_sot[:, self.tokenizer.no_speech]).tolist() if no_speech_probs[0] > self.cfg.nonspeech_prob: logger.info("no speech, stop") break logits = logits[:, -1, :] # Last token logits # Suppress tokens at segment start if new_segment: blank_tokens = self.tokenizer.encode(" ") + [self.tokenizer.eot] logits = logits.at[:, blank_tokens].add(-float('inf')) new_segment = False logits = self._suppress_tokens(logits) current_tokens, completed = self.state.token_decoder.update( current_tokens, logits, sum_logprobs ) mx.eval(current_tokens) logger.debug(f"Decoding completed: {completed}") self.debug_print_tokens(current_tokens) attn_of_alignment_heads = self._process_cross_attention( accumulated_cross_attns, content_mel_len ) most_attended_frames = mx.argmax(attn_of_alignment_heads[:, -1, :], axis=-1) most_attended_frames_np = np.array(most_attended_frames) absolute_timestamps = [ (frame * 0.02 + self.state.cumulative_time_offset) for frame in most_attended_frames_np.tolist() ] logger.debug(str(most_attended_frames_np.tolist()) + " most att frames") logger.debug(f"Absolute timestamps: {absolute_timestamps}") most_attended_frame = int(most_attended_frames_np[0]) l_absolute_timestamps.append(absolute_timestamps[0]) if completed: current_tokens = current_tokens[:, :-1] break if not is_last and self.state.last_attend_frame - most_attended_frame > self.cfg.rewind_threshold: current_tokens_np = np.array(current_tokens) if current_tokens.shape[1] > 1 and current_tokens_np[0, -2] >= DEC_PAD: logger.debug("omit rewinding from special tokens") self.state.last_attend_frame = most_attended_frame else: logger.debug(f"[rewind detected] current: {most_attended_frame}, last: {self.state.last_attend_frame}") self.state.last_attend_frame = -self.cfg.rewind_threshold current_tokens = mx.concatenate(self.state.tokens, axis=1) if len(self.state.tokens) > 0 else self.state.tokens[0] break else: self.state.last_attend_frame = most_attended_frame if content_mel_len - most_attended_frame <= (4 if is_last else self.cfg.frame_threshold): logger.debug(f"attention reaches the end: {most_attended_frame}/{content_mel_len}") current_tokens = current_tokens[:, :-1] break tokens_to_split = np.array(current_tokens[0, token_len_before_decoding:]).tolist() if self.state.pending_incomplete_tokens: logger.debug(f"[UTF-8 Fix] Prepending pending tokens: {self.state.pending_incomplete_tokens}") tokens_to_split = self.state.pending_incomplete_tokens + tokens_to_split if fire_detected or is_last: new_hypothesis = tokens_to_split split_words, split_tokens = self.tokenizer.split_to_word_tokens(new_hypothesis) else: split_words, split_tokens = self.tokenizer.split_to_word_tokens(tokens_to_split) if len(split_words) > 1: new_hypothesis = [i for sublist in split_tokens[:-1] for i in sublist] else: new_hypothesis = [] logger.debug(f"new_hypothesis: {new_hypothesis}") new_tokens = mx.array([new_hypothesis], dtype=mx.int32) new_tokens = mx.repeat(new_tokens, self.cfg.beam_size, axis=0) self.state.tokens.append(new_tokens) logger.info(f"Output: {self.tokenizer.decode(new_hypothesis)}") self._clean_cache() if len(l_absolute_timestamps) >= 2 and self.state.first_timestamp is None: self.state.first_timestamp = l_absolute_timestamps[0] timestamped_words = [] timestamp_idx = 0 replacement_char = "\ufffd" for word, word_tokens in zip(split_words, split_tokens): if replacement_char in word: logger.warning(f"[UTF-8 Filter] Skipping: {repr(word)}") timestamp_idx += len(word_tokens) continue try: current_timestamp = l_absolute_timestamps[timestamp_idx] except IndexError: pass timestamp_idx += len(word_tokens) timestamp_entry = ASRToken( start=round(current_timestamp, 2), end=round(current_timestamp + 0.1, 2), text=word, speaker=self.state.speaker, detected_language=self.state.detected_language ).with_offset(self.state.global_time_offset) timestamped_words.append(timestamp_entry) self.state.pending_incomplete_tokens = [] MAX_PENDING_TOKENS = 10 if split_words and replacement_char in split_words[-1]: if len(split_tokens[-1]) <= MAX_PENDING_TOKENS: self.state.pending_incomplete_tokens = split_tokens[-1] logger.debug(f"[UTF-8 Fix] Holding incomplete tokens") else: logger.warning(f"[UTF-8 Fix] Skipping too many tokens") return timestamped_words def _process_cross_attention( self, cross_attns: List[List[mx.array]], content_mel_len: int ) -> mx.array: """ Process cross-attention weights for alignment. Args: cross_attns: List of cross-attention from each forward pass Each element is a list of mx.arrays per layer content_mel_len: Length of actual audio content Returns: Processed attention tensor, shape (batch, seq_len, content_mel_len) """ attn_of_alignment_heads = [[] for _ in range(self.state.num_align_heads)] num_decoder_layers = self.num_decoder_layers if cross_attns and isinstance(cross_attns[0], list): flattened_attns = [attn for layer_list in cross_attns for attn in layer_list] else: flattened_attns = cross_attns for idx, attn_mat in enumerate(flattened_attns): if attn_mat is None: continue layer_rank = idx % num_decoder_layers align_heads_in_layer = self.state.align_source.get(layer_rank, []) if len(align_heads_in_layer) == 0: continue attn_mat = mx.softmax(attn_mat, axis=-1) for align_head_rank, head_id in align_heads_in_layer: if self.cfg.beam_size == 1: if attn_mat.ndim == 4: a = attn_mat[0, head_id, :, :] else: a = attn_mat[head_id, :, :] a = a[None, :, :] else: a = attn_mat[:, head_id, :, :] attn_of_alignment_heads[align_head_rank].append(a) tmp = [] for mat in attn_of_alignment_heads: if mat: t = mx.concatenate(mat, axis=1) tmp.append(t) if not tmp: return mx.zeros((self.cfg.beam_size, 1, content_mel_len)) attn_of_alignment_heads = mx.stack(tmp, axis=1) std = mx.std(attn_of_alignment_heads, axis=-2, keepdims=True) mean = mx.mean(attn_of_alignment_heads, axis=-2, keepdims=True) attn_of_alignment_heads = (attn_of_alignment_heads - mean) / (std + 1e-8) attn_of_alignment_heads = mlx_median_filter(attn_of_alignment_heads, 7) attn_of_alignment_heads = mx.mean(attn_of_alignment_heads, axis=1) attn_of_alignment_heads = attn_of_alignment_heads[:, :, :content_mel_len] mx.eval(attn_of_alignment_heads) return attn_of_alignment_heads