""" Processor class for VibeVoice ASR models. """ import os import json import math import warnings from typing import List, Optional, Union, Dict, Any, Tuple import numpy as np import torch from transformers.tokenization_utils_base import BatchEncoding from transformers.utils import TensorType, logging from .vibevoice_tokenizer_processor import VibeVoiceTokenizerProcessor, AudioNormalizer try: from .audio_utils import load_audio_use_ffmpeg HAS_FFMPEG_UTILS = True except ImportError: HAS_FFMPEG_UTILS = False warnings.warn("audio_utils not available, will fall back to soundfile for audio loading") logger = logging.get_logger(__name__) SYSTEM_PROMPT = "You are a helpful assistant that transcribes audio input into text output in JSON format." class VibeVoiceASRProcessor: """ Processor for VibeVoice ASR (Automatic Speech Recognition) models. This processor handles audio preprocessing and tokenization for ASR tasks, following the exact format used in training with proper chat templates. Args: tokenizer: The text tokenizer for processing text audio_processor: The audio processor for processing speech speech_tok_compress_ratio (int): Compression ratio for speech tokenization target_sample_rate (int): Target sample rate for audio normalize_audio (bool): Whether to normalize audio input """ def __init__( self, tokenizer=None, audio_processor=None, speech_tok_compress_ratio=320, target_sample_rate=24000, normalize_audio=True, **kwargs ): self.tokenizer = tokenizer self.audio_processor = audio_processor or VibeVoiceTokenizerProcessor( sampling_rate=target_sample_rate, normalize_audio=normalize_audio ) self.speech_tok_compress_ratio = speech_tok_compress_ratio self.target_sample_rate = target_sample_rate self.normalize_audio = normalize_audio if normalize_audio: self.audio_normalizer = AudioNormalizer() else: self.audio_normalizer = None # Cache special token IDs self._cache_special_tokens() def _cache_special_tokens(self): """Cache special token IDs for efficiency.""" # Add safety checks for special tokens if hasattr(self.tokenizer, 'speech_start_id'): self.speech_start_id = self.tokenizer.speech_start_id else: self.speech_start_id = self.tokenizer.convert_tokens_to_ids("<|speech_start|>") if hasattr(self.tokenizer, 'speech_end_id'): self.speech_end_id = self.tokenizer.speech_end_id else: self.speech_end_id = self.tokenizer.convert_tokens_to_ids("<|speech_end|>") if hasattr(self.tokenizer, 'speech_pad_id'): self.speech_pad_id = self.tokenizer.speech_pad_id else: self.speech_pad_id = self.tokenizer.convert_tokens_to_ids("<|speech_pad|>") if hasattr(self.tokenizer, 'pad_id'): self.pad_id = self.tokenizer.pad_id elif hasattr(self.tokenizer, 'pad_token_id'): self.pad_id = self.tokenizer.pad_token_id else: self.pad_id = self.tokenizer.convert_tokens_to_ids("<|endoftext|>") @classmethod def from_pretrained(cls, pretrained_model_name_or_path, **kwargs): """ Load processor from a pretrained model path. Args: pretrained_model_name_or_path: Path to the pretrained model **kwargs: Additional keyword arguments Returns: VibeVoiceASRProcessor: The loaded processor """ import json from transformers.utils import cached_file from vibevoice.modular.modular_vibevoice_text_tokenizer import VibeVoiceASRTextTokenizerFast # Try to load configuration config_path = os.path.join(pretrained_model_name_or_path, "preprocessor_config.json") config = {} if os.path.exists(config_path): with open(config_path, 'r') as f: config = json.load(f) else: try: config_file = cached_file( pretrained_model_name_or_path, "preprocessor_config.json", **kwargs ) with open(config_file, 'r') as f: config = json.load(f) except Exception as e: logger.warning(f"Could not load preprocessor_config.json: {e}") logger.warning("Using default configuration") # Extract parameters speech_tok_compress_ratio = config.get("speech_tok_compress_ratio", 3200) target_sample_rate = config.get("target_sample_rate", 24000) normalize_audio = config.get("normalize_audio", True) # Load tokenizer language_model_pretrained_name = config.get("language_model_pretrained_name", None) or kwargs.pop("language_model_pretrained_name", "Qwen/Qwen2.5-1.5B") logger.info(f"Loading tokenizer from {language_model_pretrained_name}") if 'qwen' in language_model_pretrained_name.lower(): tokenizer = VibeVoiceASRTextTokenizerFast.from_pretrained( language_model_pretrained_name, **kwargs ) else: raise ValueError(f"Unsupported tokenizer type for {language_model_pretrained_name}") # Load audio processor audio_processor = VibeVoiceTokenizerProcessor( sampling_rate=target_sample_rate, normalize_audio=normalize_audio, target_dB_FS=config.get("target_dB_FS", -25), eps=config.get("eps", 1e-6), ) return cls( tokenizer=tokenizer, audio_processor=audio_processor, speech_tok_compress_ratio=speech_tok_compress_ratio, target_sample_rate=target_sample_rate, normalize_audio=normalize_audio, ) def save_pretrained(self, save_directory: Union[str, os.PathLike], **kwargs): """ Save processor configuration to a directory. Args: save_directory: Directory to save the configuration **kwargs: Additional keyword arguments """ import json os.makedirs(save_directory, exist_ok=True) # Save processor configuration processor_config = { "processor_class": "VibeVoiceASRProcessor", "speech_tok_compress_ratio": self.speech_tok_compress_ratio, "target_sample_rate": self.target_sample_rate, "normalize_audio": self.normalize_audio, "target_dB_FS": -25, "eps": 1e-6, } config_path = os.path.join(save_directory, "preprocessor_config.json") with open(config_path, 'w') as f: json.dump(processor_config, f, indent=2) logger.info(f"Processor configuration saved in {config_path}") def __call__( self, audio: Optional[Union[str, np.ndarray, torch.Tensor, List[Union[str, np.ndarray, torch.Tensor]]]] = None, sampling_rate: Optional[int] = None, return_tensors: Optional[Union[str, TensorType]] = None, padding: bool = True, max_length: Optional[int] = None, truncation: bool = False, add_generation_prompt: bool = True, use_streaming: bool = True, context_info: Optional[str] = None, **kwargs ) -> BatchEncoding: """ Process audio input for ASR model. Args: audio: Audio input(s). Can be: - str: Path to audio file - np.ndarray: Audio array - torch.Tensor: Audio tensor - List of the above for batch processing sampling_rate: Sampling rate of input audio return_tensors: Output format ('pt' for PyTorch, 'np' for NumPy) padding: Whether to pad batch inputs max_length: Maximum sequence length truncation: Whether to truncate long sequences add_generation_prompt: Whether to add generation prompt for inference use_streaming: Whether to use streaming mode (True by default, auto False if <60s) context_info: Optional context information (e.g., hotwords, metadata) to help transcription Returns: BatchEncoding with: - input_ids: Token IDs for the model - attention_mask: Attention mask - acoustic_input_mask: Mask indicating speech token positions - speech_tensors: Processed speech features - speech_masks: Valid speech masks - vae_tok_seqlens: Length of each speech segment in tokens """ if audio is None: raise ValueError("Audio input is required for ASR processing") # Handle single vs batch input if isinstance(audio, list): is_batched = True audio_list = audio else: is_batched = False audio_list = [audio] # Process each audio input all_encodings = [] for audio_input in audio_list: encoding = self._process_single_audio( audio_input, sampling_rate=sampling_rate, add_generation_prompt=add_generation_prompt, use_streaming=use_streaming, context_info=context_info, ) all_encodings.append(encoding) # Combine into batch batch_encoding = self._batch_encode( all_encodings, padding=padding, max_length=max_length, truncation=truncation, return_tensors=return_tensors, ) return batch_encoding def _process_single_audio( self, audio: Union[str, np.ndarray, torch.Tensor], sampling_rate: Optional[int] = None, add_generation_prompt: bool = True, use_streaming: bool = True, context_info: Optional[str] = None, ) -> Dict[str, Any]: """ Process a single audio input. Args: audio: Single audio input sampling_rate: Audio sampling rate add_generation_prompt: Whether to add generation prompt context_info: Optional context information (e.g., hotwords, metadata) to help transcription Returns: Dictionary with processed tokens and audio features """ # Process audio through audio processor if isinstance(audio, str): # Load from file using ffmpeg for better format support if HAS_FFMPEG_UTILS: try: audio_array, file_sr = load_audio_use_ffmpeg(audio, resample=False) except Exception as e: # Fall back to soundfile if ffmpeg fails warnings.warn(f"ffmpeg loading failed, falling back to soundfile: {e}") import soundfile as sf audio_array, file_sr = sf.read(audio) if audio_array.ndim > 1: audio_array = audio_array.mean(axis=1) # Convert to mono else: import soundfile as sf audio_array, file_sr = sf.read(audio) if audio_array.ndim > 1: audio_array = audio_array.mean(axis=1) # Convert to mono # Resample if needed if file_sr != self.target_sample_rate: import librosa audio_array = librosa.resample( audio_array, orig_sr=file_sr, target_sr=self.target_sample_rate ) elif isinstance(audio, torch.Tensor): audio_array = audio.cpu().numpy() if audio_array.ndim > 1: audio_array = audio_array.squeeze() else: audio_array = np.array(audio, dtype=np.float32) if audio_array.ndim > 1: audio_array = audio_array.squeeze() # Ensure float32 audio_array = audio_array.astype(np.float32) # Normalize if needed if self.normalize_audio and self.audio_normalizer: audio_array = self.audio_normalizer(audio_array) # Calculate audio duration audio_duration = len(audio_array) / self.target_sample_rate # Auto-disable streaming for short audio (<60s) if use_streaming and audio_duration < 60.0: use_streaming = False # Calculate token length based on streaming mode # Non-streaming: uses ceil (encoder adds extra_padding for stride alignment) # Streaming: uses floor (segments processed independently, no global alignment) # if use_streaming: # vae_tok_len = len(audio_array) // self.speech_tok_compress_ratio # else: vae_tok_len = math.ceil(len(audio_array) / self.speech_tok_compress_ratio) # Build token sequence following training format # 1. System prompt - use apply_chat_template then encode like in training system_prompt_text = self.tokenizer.apply_chat_template( [{"role": "system", "content": SYSTEM_PROMPT}], tokenize=False ) system_tokens = self.tokenizer.encode(system_prompt_text) # 2. User input with speech tokens # Build speech placeholder string sp_start_token = self.tokenizer.convert_ids_to_tokens(self.speech_start_id) sp_pad_token = self.tokenizer.convert_ids_to_tokens(self.speech_pad_id) sp_end_token = self.tokenizer.convert_ids_to_tokens(self.speech_end_id) # User suffix with audio duration info show_keys = ['Start time', 'End time', 'Speaker ID', 'Content'] if context_info and context_info.strip(): user_suffix = f"This is a {audio_duration:.2f} seconds audio, with extra info: {context_info.strip()}\n\nPlease transcribe it with these keys: " + ", ".join(show_keys) else: user_suffix = f"This is a {audio_duration:.2f} seconds audio, please transcribe it with these keys: " + ", ".join(show_keys) user_input_string = ''.join( [sp_start_token] + [sp_pad_token] * vae_tok_len + [sp_end_token] ) + '\n' + user_suffix user_tokens = self.tokenizer.apply_chat_template( [{"role": "user", "content": user_input_string}], tokenize=True ) # Combine tokens full_tokens = system_tokens + user_tokens # Create acoustic input mask acoustic_input_mask = [1 if token == self.speech_pad_id else 0 for token in full_tokens] return { "input_ids": full_tokens, "acoustic_input_mask": acoustic_input_mask, "speech": audio_array, "vae_tok_len": vae_tok_len, } def _batch_encode( self, encodings: List[Dict[str, Any]], padding: bool = True, max_length: Optional[int] = None, truncation: bool = False, return_tensors: Optional[str] = None, ) -> BatchEncoding: """ Combine multiple encodings into a batch. Args: encodings: List of encoded samples padding: Whether to pad sequences max_length: Maximum sequence length truncation: Whether to truncate return_tensors: Output format Returns: BatchEncoding with batched data """ # Extract components input_ids_list = [enc["input_ids"] for enc in encodings] acoustic_masks_list = [enc["acoustic_input_mask"] for enc in encodings] speech_list = [enc["speech"] for enc in encodings] vae_tok_lens = [enc["vae_tok_len"] for enc in encodings] # Determine max length for padding if padding: if max_length is not None: target_length = max_length else: target_length = max(len(ids) for ids in input_ids_list) # Pad sequences padded_input_ids = [] padded_acoustic_masks = [] attention_masks = [] for input_ids, acoustic_mask in zip(input_ids_list, acoustic_masks_list): # Truncate if needed if truncation and len(input_ids) > target_length: input_ids = input_ids[:target_length] acoustic_mask = acoustic_mask[:target_length] # Pad sequences to left (for autoregressive generation) padding_length = target_length - len(input_ids) padded_ids = [self.pad_id] * padding_length + input_ids padded_acoustic = [0] * padding_length + acoustic_mask attention_mask = [0] * padding_length + [1] * len(input_ids) padded_input_ids.append(padded_ids) padded_acoustic_masks.append(padded_acoustic) attention_masks.append(attention_mask) input_ids_list = padded_input_ids acoustic_masks_list = padded_acoustic_masks else: attention_masks = [[1] * len(ids) for ids in input_ids_list] # Process speech tensors - raw audio is 1D, so we keep it as is max_speech_length = max(len(s) for s in speech_list) padded_speeches = np.zeros((len(speech_list), max_speech_length), dtype=np.float32) speech_masks = np.zeros((len(speech_list), max(vae_tok_lens)), dtype=bool) for i, (speech, vae_len) in enumerate(zip(speech_list, vae_tok_lens)): padded_speeches[i, :len(speech)] = speech speech_masks[i, :vae_len] = True # Create batch encoding batch_encoding = BatchEncoding() if return_tensors == "pt": batch_encoding["input_ids"] = torch.tensor(input_ids_list, dtype=torch.long) batch_encoding["attention_mask"] = torch.tensor(attention_masks, dtype=torch.long) batch_encoding["acoustic_input_mask"] = torch.tensor(acoustic_masks_list, dtype=torch.bool) batch_encoding["speech_tensors"] = torch.tensor(padded_speeches, dtype=torch.float32) batch_encoding["speech_masks"] = torch.tensor(speech_masks, dtype=torch.bool) # Note: vae_tok_seqlens and speech_type are not included as they are not model inputs else: batch_encoding["input_ids"] = input_ids_list if len(input_ids_list) > 1 else input_ids_list[0] batch_encoding["attention_mask"] = attention_masks if len(attention_masks) > 1 else attention_masks[0] batch_encoding["acoustic_input_mask"] = acoustic_masks_list if len(acoustic_masks_list) > 1 else acoustic_masks_list[0] batch_encoding["speech_tensors"] = padded_speeches if len(padded_speeches) > 1 else padded_speeches[0] batch_encoding["speech_masks"] = speech_masks if len(speech_masks) > 1 else speech_masks[0] return batch_encoding def batch_decode(self, *args, **kwargs): """ Decode batch of token IDs to text. Forwards to tokenizer's batch_decode method. """ return self.tokenizer.batch_decode(*args, **kwargs) def decode(self, *args, **kwargs): """ Decode token IDs to text. Forwards to tokenizer's decode method. """ return self.tokenizer.decode(*args, **kwargs) def post_process_transcription(self, text: str) -> List[Dict[str, Any]]: """ Post-process the generated transcription text to extract structured data. Args: text: Generated text from the model Returns: List of dictionaries with transcription segments """ try: # Try to parse as JSON if "```json" in text: # Extract JSON from markdown code block json_start = text.find("```json") + 7 json_end = text.find("```", json_start) json_str = text[json_start:json_end].strip() else: # Try to find JSON array or object json_start = text.find("[") if json_start == -1: json_start = text.find("{") if json_start != -1: # Find matching closing bracket 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 # Parse JSON result = json.loads(json_str) # Ensure it's a list if isinstance(result, dict): result = [result] # Validate and clean up the result cleaned_result = [] for item in result: if isinstance(item, dict): cleaned_item = {} # Map keys to expected format key_mapping = { "Start time": "start_time", "Start": "start_time", "End time": "end_time", "End": "end_time", "Speaker ID": "speaker_id", "Speaker": "speaker_id", "Content": "text", } for key, mapped_key in key_mapping.items(): if key in item: cleaned_item[mapped_key] = item[key] if cleaned_item: cleaned_result.append(cleaned_item) return cleaned_result except json.JSONDecodeError as e: logger.warning(f"Failed to parse JSON from transcription: {e}") logger.debug(f"Raw text: {text}") return [] except Exception as e: logger.warning(f"Error post-processing transcription: {e}") return [] @property def model_input_names(self): """Return the list of inputs accepted by the model.""" return ["input_ids", "attention_mask", "acoustic_input_mask", "speech_tensors", "speech_masks"] __all__ = ["VibeVoiceASRProcessor"]