# Copyright (c) 2025, Prince Canuma and contributors (https://github.com/Blaizzy/mlx-audio) import re import time import warnings from pathlib import Path from typing import Any, Callable, Dict, Generator, List, Optional, Tuple import mlx.core as mx import mlx.nn as nn import numpy as np from tqdm import tqdm from ..base import STTOutput from .audio_encoder import AcousticTokenizerEncoder, SemanticTokenizerEncoder from .config import ModelConfig class SpeechConnector(nn.Module): """ MLP connector to project speech features to LM hidden dimension. Structure: Linear -> RMSNorm -> Linear """ def __init__(self, input_dim: int, output_dim: int, eps: float = 1e-6): super().__init__() self.fc1 = nn.Linear(input_dim, output_dim) self.norm = nn.RMSNorm(output_dim, eps=eps) self.fc2 = nn.Linear(output_dim, output_dim) def __call__(self, x: mx.array) -> mx.array: x = self.fc1(x) x = self.norm(x) x = self.fc2(x) return x class LanguageModel(nn.Module): """ Qwen2-based language model wrapper using mlx_lm. This wraps the LlamaModel from mlx_lm (which supports Qwen2 architecture). """ def __init__(self, config): super().__init__() self.config = config self.model_type = config.model_type # Import and use mlx_lm's Qwen2Model try: from mlx_lm.models.qwen2 import Qwen2Model self.model = Qwen2Model(config) except ImportError: # Fallback to llama if qwen2 not available from mlx_lm.models.llama import LlamaModel self.model = LlamaModel(config) if not config.tie_word_embeddings: self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) def __call__( self, inputs: Optional[mx.array] = None, cache: Optional[List[Any]] = None, input_embeddings: Optional[mx.array] = None, ): out = self.model(inputs, cache=cache, input_embeddings=input_embeddings) if self.config.tie_word_embeddings: out = self.model.embed_tokens.as_linear(out) else: out = self.lm_head(out) return out @property def layers(self): return self.model.layers @property def embed_tokens(self): return self.model.embed_tokens class Model(nn.Module): """ VibeVoice-ASR model for speech-to-text transcription. Architecture: - Acoustic tokenizer encoder: Encodes raw audio to acoustic features - Semantic tokenizer encoder: Encodes raw audio to semantic features - Acoustic connector: Projects acoustic features to LM dimension - Semantic connector: Projects semantic features to LM dimension - Language model (Qwen2): Generates text from combined features The model supports: - Long-form audio transcription (up to 60 minutes) - Speaker diarization - Timestamp generation """ def __init__(self, config: ModelConfig): super().__init__() self.config = config # Audio tokenizer encoders self.acoustic_tokenizer = AcousticTokenizerEncoder( config.acoustic_tokenizer_config ) self.semantic_tokenizer = SemanticTokenizerEncoder( config.semantic_tokenizer_config ) # Speech connectors self.acoustic_connector = SpeechConnector( config.acoustic_vae_dim, config.decoder_config.hidden_size ) self.semantic_connector = SpeechConnector( config.semantic_vae_dim, config.decoder_config.hidden_size ) # Language model (includes lm_head for logit computation) self.language_model = LanguageModel(config.decoder_config) def get_input_embeddings(self) -> nn.Embedding: """Get the input embeddings from the language model.""" return self.language_model.embed_tokens def model_quant_predicate(self, p, m): """Only quantize language model layers.""" return p.startswith("language_model") def encode_speech( self, speech_tensors: mx.array, verbose: bool = False ) -> mx.array: """ Encode speech input to features for the language model. Args: speech_tensors: Audio waveform [B, T] at 24kHz verbose: Show progress bar Returns: Combined acoustic + semantic features [B, T', hidden_size] """ # Ensure speech has batch and channel dimensions if speech_tensors.ndim == 1: speech_tensors = speech_tensors[None, :] if speech_tensors.ndim == 2: speech_tensors = speech_tensors[:, None, :] # [B, 1, T] pbar = tqdm(total=4, disable=not verbose, desc="Encoding audio") # Encode through acoustic tokenizer acoustic_tokens = self.acoustic_tokenizer.encode(speech_tensors) mx.eval(acoustic_tokens) pbar.update(1) # Project acoustic features acoustic_features = self.acoustic_connector(acoustic_tokens) mx.eval(acoustic_features) pbar.update(1) # Encode through semantic tokenizer semantic_tokens = self.semantic_tokenizer.encode(speech_tensors) mx.eval(semantic_tokens) pbar.update(1) # Project semantic features semantic_features = self.semantic_connector(semantic_tokens) mx.eval(semantic_features) pbar.update(1) pbar.close() # Combine features combined_features = acoustic_features + semantic_features return combined_features def _merge_speech_text_embeddings( self, input_ids: mx.array, speech_features: Optional[mx.array] = None, acoustic_input_mask: Optional[mx.array] = None, cache: Optional[List[Any]] = None, ) -> mx.array: """ Merge speech features into text embeddings at masked positions. Args: input_ids: Token IDs [B, L] speech_features: Pre-computed speech features [B, T', hidden_size] acoustic_input_mask: Boolean mask indicating where to insert speech [B, L] cache: KV cache (if populated, skip merging as it's done) Returns: Combined embeddings [B, L, hidden_size] """ # Get text embeddings text_embeds = self.get_input_embeddings()(input_ids) # Skip if no speech features or cache already populated if speech_features is None or (cache is not None and cache[0].offset > 0): return text_embeds # Insert speech features at masked positions if acoustic_input_mask is not None: # acoustic_input_mask is [B, L] boolean # speech_features is [B, T', hidden_size] # We use mx.where with 3 args to conditionally select values batch_size, seq_len, hidden_size = text_embeds.shape num_speech_tokens = speech_features.shape[1] # Expand mask to match embedding shape [B, L, hidden_size] mask_expanded = acoustic_input_mask[:, :, None] mask_expanded = mx.broadcast_to(mask_expanded, text_embeds.shape) # Create padded speech features matching text_embeds shape # We need to place speech_features at the mask positions # First, find the indices where mask is True for each batch for b in range(batch_size): # Get flat indices where mask is true mask_flat = acoustic_input_mask[b] # Count cumulative True values to map to speech feature indices cumsum = mx.cumsum(mask_flat.astype(mx.int32)) # speech_idx maps each position to its speech feature index # (or 0 if not a speech position, but we'll use the mask anyway) speech_idx = cumsum - 1 speech_idx = mx.clip(speech_idx, 0, num_speech_tokens - 1) # Gather speech features for this batch # speech_features[b] is [T', hidden_size] # We need to expand it to [L, hidden_size] using speech_idx expanded_speech = speech_features[b][speech_idx] # [L, hidden_size] # Use where to select between speech and text embeddings text_embeds = mx.where( mask_expanded[b : b + 1], expanded_speech[None], text_embeds[b : b + 1], ) return text_embeds def __call__( self, input_ids: mx.array, speech_tensors: Optional[mx.array] = None, acoustic_input_mask: Optional[mx.array] = None, speech_features: Optional[mx.array] = None, cache: Optional[List[Any]] = None, ) -> mx.array: """ Forward pass. Args: input_ids: Token IDs [B, L] speech_tensors: Raw audio [B, T] at 24kHz acoustic_input_mask: Mask for where to insert speech in text sequence speech_features: Pre-computed speech features (optional) cache: KV cache for generation Returns: Logits [B, L, vocab_size] """ # Encode speech if raw audio provided if speech_tensors is not None and speech_features is None: speech_features = self.encode_speech(speech_tensors) # Merge speech and text embeddings input_embeds = self._merge_speech_text_embeddings( input_ids=input_ids, speech_features=speech_features, acoustic_input_mask=acoustic_input_mask, cache=cache, ) # Forward through language model (returns logits including lm_head) logits = self.language_model( inputs=None, cache=cache, input_embeddings=input_embeds ) return logits @staticmethod def sanitize(weights: Dict[str, mx.array]) -> Dict[str, mx.array]: """ Sanitize weights from PyTorch checkpoint to MLX model structure. This handles: - Key name remapping (model.* prefix removal) - Conv weight transposition (PyTorch [O,I,K] -> MLX [O,K,I]) - Tokenizer encoder path remapping PyTorch tokenizer encoder key patterns: - downsample_layers.N.0.conv.conv.X -> downsample_layers.N.conv.X - head.conv.conv.X -> head.conv.X - stages.N.M.mixer.conv.conv.conv.X -> stages.N.M.mixer.conv.conv.X - stages.N.M.{ffn,norm,ffn_norm,gamma,ffn_gamma}.* -> unchanged """ sanitized = {} already_converted = not any(k.startswith("model.") for k in weights.keys()) for key, value in weights.items(): new_key = key # Remove 'model.' prefix if present if new_key.startswith("model."): new_key = new_key[6:] # Skip decoder weights (we only use encoders for ASR) if "acoustic_tokenizer.decoder" in new_key: continue # Handle tokenizer encoder weights (acoustic and semantic) if ( "acoustic_tokenizer.encoder." in new_key or "semantic_tokenizer.encoder." in new_key ): # PyTorch path already matches our wrapper structure: # Model.acoustic_tokenizer (AcousticTokenizerEncoder) # .encoder (TokenizerEncoder) # So acoustic_tokenizer.encoder.* maps directly. # Fix downsample_layers: remove .0. wrapper and one .conv. level # PyTorch: downsample_layers.N.0.conv.conv.X -> MLX: downsample_layers.N.conv.X if ".downsample_layers." in new_key: new_key = re.sub( r"\.downsample_layers\.(\d+)\.0\.conv\.conv\.", r".downsample_layers.\1.conv.", new_key, ) # Fix head: remove one .conv. level # PyTorch: head.conv.conv.X -> MLX: head.conv.X elif ".head.conv.conv." in new_key: new_key = new_key.replace(".head.conv.conv.", ".head.conv.") # Fix mixer in stages: remove one .conv. level # PyTorch: mixer.conv.conv.conv.X -> MLX: mixer.conv.conv.X elif ".mixer.conv.conv.conv." in new_key: new_key = new_key.replace( ".mixer.conv.conv.conv.", ".mixer.conv.conv." ) # Other stage keys (ffn, norm, gamma) map directly # Handle language model weights # PyTorch: language_model.{layers,embed_tokens,norm}.* -> language_model.model.* if new_key.startswith("language_model.layers."): new_key = "language_model.model." + new_key[len("language_model.") :] elif new_key.startswith("language_model.embed_tokens"): new_key = ( "language_model.model.embed_tokens" + new_key[len("language_model.embed_tokens") :] ) elif new_key.startswith("language_model.norm"): new_key = ( "language_model.model.norm" + new_key[len("language_model.norm") :] ) # Map lm_head to language_model.lm_head if new_key.startswith("lm_head."): new_key = "language_model." + new_key # Handle Conv1d weight transposition # PyTorch: [out_channels, in_channels/groups, kernel_size] # MLX: [out_channels, kernel_size, in_channels/groups] if ( not already_converted and "conv" in new_key.lower() and "weight" in new_key ): if value.ndim == 3: value = value.transpose(0, 2, 1) # Skip position_ids as MLX doesn't use them if "position_ids" in new_key: continue # Skip fix_std buffer - we handle it in config if "fix_std" in new_key: continue sanitized[new_key] = value return sanitized @classmethod def post_load_hook(cls, model: "Model", model_path: Path) -> "Model": """ Hook called after model weights are loaded. Initializes the tokenizer for text input/output. """ from transformers import AutoTokenizer # Load tokenizer (Qwen2.5-7B base, which has the special tokens we need) try: tokenizer = AutoTokenizer.from_pretrained( str(model_path), trust_remote_code=True ) except Exception: tokenizer = AutoTokenizer.from_pretrained( "Qwen/Qwen2.5-7B", trust_remote_code=True ) # Set the chat template used by VibeVoice ASR tokenizer.chat_template = ( "{% for message in messages %}" "{{'<|im_start|>' + message['role'] + '\\n' + message['content'] + '<|im_end|>' + '\\n'}}" "{% endfor %}" "{% if add_generation_prompt %}{{ '<|im_start|>assistant\\n' }}{% endif %}" ) model.tokenizer = tokenizer # VibeVoice ASR repurposes existing Qwen2.5 special tokens for speech: # <|object_ref_start|> = speech_start # <|object_ref_end|> = speech_end # <|box_start|> = speech_pad (placeholder for audio features) model._speech_start_id = tokenizer.convert_tokens_to_ids("<|object_ref_start|>") model._speech_end_id = tokenizer.convert_tokens_to_ids("<|object_ref_end|>") model._speech_pad_id = tokenizer.convert_tokens_to_ids("<|box_start|>") return model @classmethod def from_pretrained(cls, model_path: str, **kwargs) -> "Model": """ Load model from pretrained weights. .. deprecated:: Use `mlx_audio.stt.load()` instead. """ warnings.warn( "Model.from_pretrained() is deprecated. Use mlx_audio.stt.load() instead.", DeprecationWarning, stacklevel=2, ) from mlx_audio.stt.utils import load return load(model_path) @staticmethod def _normalize_audio( audio: np.ndarray, target_dB_FS: float = -25.0, eps: float = 1e-6, ) -> np.ndarray: """ Normalize audio to target dB FS level and avoid clipping. This matches the official VibeVoice AudioNormalizer which adjusts audio loudness to a consistent level before encoding. Args: audio: Input audio waveform as numpy array. target_dB_FS: Target loudness in dB FS (default: -25). eps: Small value to avoid division by zero. Returns: Normalized audio as numpy array. """ rms = np.sqrt(np.mean(audio**2)) scalar = 10 ** (target_dB_FS / 20) / (rms + eps) audio = audio * scalar max_val = np.max(np.abs(audio)) if max_val > 1.0: audio = audio / (max_val + eps) return audio def _preprocess_audio( self, audio, *, sampling_rate: Optional[int] = None ) -> mx.array: """ Preprocess audio for the model. Handles loading, resampling to 24 kHz, normalizing loudness to -25 dB FS, and trimming to the maximum supported duration. Args: audio: Audio path (str), waveform (np.ndarray/mx.array) sampling_rate: Sample rate of the input waveform. Required when *audio* is an array that is **not** already at 24 kHz so that the audio can be resampled correctly. Returns: Audio tensor ready for encoding [B, T] """ from mlx_audio.stt.utils import load_audio, resample_audio SAMPLE_RATE = 24000 MAX_DURATION_SECONDS = 59 * 60 # 59 minutes max if isinstance(audio, str): audio = load_audio(audio, sr=SAMPLE_RATE) else: # Convert to numpy for resampling / normalization if isinstance(audio, mx.array): audio_np = np.array(audio, copy=False) else: audio_np = np.asarray(audio, dtype=np.float32) # Flatten to 1-D for resampling if audio_np.ndim > 1: audio_np = audio_np.squeeze() # Resample to 24 kHz when the caller specifies a different rate if sampling_rate is not None and sampling_rate != SAMPLE_RATE: audio_np = resample_audio(audio_np, sampling_rate, SAMPLE_RATE) # Normalize loudness to -25 dB FS (matches official pipeline) audio_np = self._normalize_audio(audio_np) audio = mx.array(audio_np) # Ensure 1D or 2D if audio.ndim == 3: audio = audio.squeeze() if audio.ndim == 1: audio = audio[None, :] # Add batch dim # Check duration and trim if necessary max_samples = MAX_DURATION_SECONDS * SAMPLE_RATE if audio.shape[-1] > max_samples: duration_minutes = audio.shape[-1] / SAMPLE_RATE / 60 print( f"\033[93m[WARNING]\033[0m Audio duration ({duration_minutes:.1f} min) exceeds " f"maximum supported duration (59 min). Trimming to 59 minutes. " f"For longer audio, consider splitting into smaller segments." ) audio = audio[..., :max_samples] return audio def stream_generate( self, input_ids: Optional[mx.array] = None, *, speech_features: Optional[mx.array] = None, acoustic_input_mask: Optional[mx.array] = None, max_tokens: int = 8192, sampler: Optional[Callable[[mx.array], mx.array]] = None, logits_processors: Optional[List[Callable]] = None, prefill_step_size: int = 2048, generation_stream: bool = False, verbose: bool = False, ) -> Generator[Tuple[mx.array, mx.array], None, None]: """ Stream generate tokens. Args: input_ids: Input token IDs speech_features: Pre-computed speech features acoustic_input_mask: Mask for speech token positions max_tokens: Maximum tokens to generate sampler: Sampling function logits_processors: List of logits processors (e.g., repetition penalty) prefill_step_size: Chunk size for prompt prefill (reduces peak memory) generation_stream: Enable streaming generation verbose: Print progress Yields: Tuple of (token, logprobs) """ from mlx_lm.generate import generate_step # Get input embeddings with speech merged in input_embeddings = self._merge_speech_text_embeddings( input_ids=input_ids, speech_features=speech_features, acoustic_input_mask=acoustic_input_mask, )[ 0 ] # Remove batch dim prompt = input_ids[0] if input_ids.ndim > 1 else input_ids # Build EOS token list: include both <|endoftext|> and <|im_end|> # The model uses <|im_end|> (151645) to end assistant turns in chat format eos_token_ids = [151643, 151645] # <|endoftext|>, <|im_end|> if hasattr(self.tokenizer, "eos_token_id") and self.tokenizer.eos_token_id: if self.tokenizer.eos_token_id not in eos_token_ids: eos_token_ids.append(self.tokenizer.eos_token_id) # Create prefill progress bar prefill_pbar = None gen_pbar = None if verbose: prefill_pbar = tqdm(total=1, desc="Prefilling", unit="tok") def prefill_progress(processed: int, total: int): nonlocal gen_pbar if prefill_pbar is not None: if prefill_pbar.total != total: prefill_pbar.total = total prefill_pbar.refresh() prefill_pbar.n = processed prefill_pbar.refresh() if processed >= total: prefill_pbar.close() # Start generating progress bar after prefill completes if gen_pbar is None and verbose: gen_pbar = tqdm(total=max_tokens, desc="Generating", unit="tok") token_count = 0 for token, logprobs in generate_step( prompt=prompt, input_embeddings=input_embeddings, model=self.language_model, max_tokens=max_tokens, sampler=sampler, logits_processors=logits_processors, prefill_step_size=prefill_step_size, prompt_progress_callback=prefill_progress if verbose else None, ): token_count += 1 if gen_pbar is not None: gen_pbar.update(1) # Check for EOS tokens if token in eos_token_ids: break yield token, logprobs if gen_pbar is not None: gen_pbar.close() def generate( self, audio, *, context: Optional[str] = None, sampling_rate: Optional[int] = None, max_tokens: int = 8192, temperature: float = 0.0, top_p: float = 1.0, top_k: int = 0, min_p: float = 0.0, min_tokens_to_keep: int = 1, repetition_penalty: Optional[float] = 1.0, repetition_context_size: int = 100, prefill_step_size: int = 2048, generation_stream: bool = False, verbose: bool = False, **kwargs, ) -> STTOutput: """ Generate transcription from audio. Args: audio: Audio path (str) or waveform (mx.array/np.array) context: Optional context string (hotwords, metadata) sampling_rate: Sample rate of the input waveform. When *audio* is an array not at 24 kHz, provide its actual sample rate so that it is resampled correctly. max_tokens: Maximum tokens to generate temperature: Sampling temperature (0 = greedy) top_p: Top-p sampling (1.0 = no filtering) top_k: Top-k sampling (0 = disabled) min_p: Min-p sampling (0.0 = disabled) min_tokens_to_keep: Min tokens for sampling repetition_penalty: Penalty for repeated tokens (1.0 = no penalty) repetition_context_size: Number of recent tokens to check for repetition prefill_step_size: Chunk size for prompt prefill (reduces peak memory) generation_stream: Enable streaming verbose: Print progress Returns: STTOutput with transcription text and segments """ from mlx_lm.sample_utils import make_logits_processors, make_sampler start_time = time.time() # Preprocess audio audio_tensor = self._preprocess_audio(audio, sampling_rate=sampling_rate) # Encode speech speech_features = self.encode_speech(audio_tensor, verbose=verbose) # Build prompt audio_duration = audio_tensor.shape[1] / 24000 input_ids, acoustic_input_mask = self._build_prompt_tokens( speech_features, audio_duration, context ) # Create sampler sampler = make_sampler( temperature, top_p, min_p, min_tokens_to_keep=min_tokens_to_keep, top_k=top_k, ) # Create logits processors with repetition penalty logits_processors = make_logits_processors( repetition_penalty=repetition_penalty, repetition_context_size=repetition_context_size, ) # Generate tokens generated_tokens = [] for token, _ in self.stream_generate( input_ids=input_ids, speech_features=speech_features, acoustic_input_mask=acoustic_input_mask, max_tokens=max_tokens, sampler=sampler, logits_processors=logits_processors, prefill_step_size=prefill_step_size, generation_stream=generation_stream, verbose=verbose, ): generated_tokens.append(token) end_time = time.time() if verbose: print() # Clear cache mx.clear_cache() # Decode output text = self.tokenizer.decode(generated_tokens, skip_special_tokens=True) # Parse structured output segments = self.parse_transcription(text) return STTOutput( text=text.strip(), segments=segments, prompt_tokens=input_ids.shape[1], generation_tokens=len(generated_tokens), total_tokens=input_ids.shape[1] + len(generated_tokens), total_time=end_time - start_time, prompt_tps=input_ids.shape[1] / (end_time - start_time), generation_tps=len(generated_tokens) / (end_time - start_time), ) def stream_transcribe( self, audio, *, context: Optional[str] = None, sampling_rate: Optional[int] = None, max_tokens: int = 8192, temperature: float = 0.0, top_p: float = 1.0, top_k: int = 0, min_p: float = 0.0, min_tokens_to_keep: int = 1, repetition_penalty: Optional[float] = 1.0, repetition_context_size: int = 100, prefill_step_size: int = 2048, verbose: bool = False, ) -> Generator[str, None, None]: """ Stream transcription token-by-token from audio. Args: audio: Audio path (str) or waveform (mx.array/np.array) context: Optional context string (hotwords, metadata) sampling_rate: Sample rate of the input waveform. When *audio* is an array not at 24 kHz, provide its actual sample rate so that it is resampled correctly. max_tokens: Maximum tokens to generate temperature: Sampling temperature (0 = greedy) top_p: Top-p sampling (1.0 = no filtering) top_k: Top-k sampling (0 = disabled) min_p: Min-p sampling min_tokens_to_keep: Min tokens for sampling repetition_penalty: Penalty for repeated tokens (1.0 = no penalty) repetition_context_size: Number of recent tokens to check for repetition prefill_step_size: Chunk size for prompt prefill (reduces peak memory) verbose: Print progress Yields: Decoded text chunks as they are generated. """ from mlx_lm.sample_utils import make_logits_processors, make_sampler # Preprocess audio audio_tensor = self._preprocess_audio(audio, sampling_rate=sampling_rate) # Encode speech speech_features = self.encode_speech(audio_tensor, verbose=verbose) # Build prompt (same as generate) audio_duration = audio_tensor.shape[1] / 24000 input_ids, acoustic_input_mask = self._build_prompt_tokens( speech_features, audio_duration, context ) # Create sampler sampler = make_sampler( temperature, top_p, min_p, min_tokens_to_keep=min_tokens_to_keep, top_k=top_k, ) # Create logits processors with repetition penalty logits_processors = make_logits_processors( repetition_penalty=repetition_penalty, repetition_context_size=repetition_context_size, ) # Stream tokens for token, _ in self.stream_generate( input_ids=input_ids, speech_features=speech_features, acoustic_input_mask=acoustic_input_mask, max_tokens=max_tokens, sampler=sampler, logits_processors=logits_processors, prefill_step_size=prefill_step_size, verbose=verbose, ): text = self.tokenizer.decode([token]) yield text mx.clear_cache() def _build_prompt_tokens( self, speech_features: mx.array, audio_duration: float, context: Optional[str] = None, ) -> Tuple[mx.array, mx.array]: """ Build input_ids and acoustic_input_mask for generation. Returns: Tuple of (input_ids [1, L], acoustic_input_mask [1, L]) """ system_prompt = ( "You are a helpful assistant that transcribes audio input into text " "output in JSON format." ) vae_tok_len = speech_features.shape[1] # VibeVoice ASR uses repurposed Qwen2.5 special tokens: speech_start_token = "<|object_ref_start|>" speech_pad_token = "<|box_start|>" speech_end_token = "<|object_ref_end|>" show_keys = ["Start time", "End time", "Speaker ID", "Content"] if context and context.strip(): user_suffix = ( f"This is a {audio_duration:.2f} seconds audio, " f"with extra info: {context.strip()}\n\n" f"Please transcribe it with these keys: " + ", ".join(show_keys) ) else: user_suffix = ( f"This is a {audio_duration:.2f} seconds audio, " f"please transcribe it with these keys: " + ", ".join(show_keys) ) messages = [ {"role": "system", "content": system_prompt}, { "role": "user", "content": speech_start_token + speech_pad_token * vae_tok_len + speech_end_token + "\n" + user_suffix, }, ] prompt_text = self.tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=True ) tokens = self.tokenizer.encode(prompt_text) input_ids = mx.array([tokens]) acoustic_input_mask = mx.array( [[t == self._speech_pad_id for t in tokens]] ).astype(mx.bool_) return input_ids, acoustic_input_mask def parse_transcription(self, text: str) -> List[Dict[str, Any]]: """Parse structured JSON output from model.""" import json try: # Try to extract JSON from text if "```json" in text: json_start = text.find("```json") + 7 json_end = text.find("```", json_start) json_str = text[json_start:json_end].strip() else: # Find JSON array or object json_start = text.find("[") if json_start == -1: json_start = text.find("{") if json_start != -1: bracket_count = 0 json_end = json_start for i in range(json_start, len(text)): if text[i] in "[{": bracket_count += 1 elif text[i] in "]}": bracket_count -= 1 if bracket_count == 0: json_end = i + 1 break json_str = text[json_start:json_end] else: json_str = text result = json.loads(json_str) if isinstance(result, dict): result = [result] # Clean up keys segments = [] key_mapping = { "Start time": "start", "Start": "start", "End time": "end", "End": "end", "Speaker ID": "speaker_id", "Speaker": "speaker_id", "Content": "text", } for item in result: if isinstance(item, dict): segment = {} for old_key, new_key in key_mapping.items(): if old_key in item: segment[new_key] = item[old_key] if segment: segments.append(segment) return segments except (json.JSONDecodeError, Exception): return []