import logging from collections.abc import Mapping from typing import Dict, List, Optional, Union, cast import numpy as np import torch from transformers import Qwen2Tokenizer, Qwen2TokenizerFast, Wav2Vec2FeatureExtractor from transformers.feature_extraction_utils import BatchFeature from transformers.processing_utils import ( AllKwargsForChatTemplate, ProcessingKwargs, ProcessorMixin, ) from typing_extensions import Unpack DEFAULT_SYSTEM_PROMPT = ( "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, " "capable of perceiving auditory and visual inputs, as well as generating text and speech." ) class MiDashengLMSpatialProcessorKwargs(ProcessingKwargs): _defaults = { # type: ignore "text_kwargs": { "padding": True, "padding_side": "left", }, "audio_kwargs": {}, } def calculate_mel_frames_dasheng( audio_length_samples: int, n_fft: int = 512, hop_size: int = 160, dasheng_subsampling: int = 4, center=True, model_subsampling: int = 5, ) -> int: if center: audio_length_samples = audio_length_samples + n_fft return ( int(1 + ((audio_length_samples - n_fft) / hop_size)) // dasheng_subsampling // model_subsampling ) class MiDashengLMSpatialProcessor(ProcessorMixin): attributes = ["feature_extractor", "tokenizer"] valid_kwargs = [ "chat_template", "audio_token", "audio_bos_token", "audio_eos_token", "audio_channels", "sem_n_fft", "min_audio_duration", "model_subsampling", ] feature_extractor_class = "Wav2Vec2FeatureExtractor" tokenizer_class = ("Qwen2Tokenizer", "Qwen2TokenizerFast") def __init__( self, feature_extractor: Wav2Vec2FeatureExtractor, tokenizer: Union[Qwen2Tokenizer, Qwen2TokenizerFast], model_subsampling: int = 5, audio_channels: int = 2, sem_n_fft: int = 512, min_audio_duration: float = 0.2, chat_template: Optional[Union[str, Dict[str, str]]] = None, audio_token: Optional[str] = None, audio_bos_token: Optional[str] = None, audio_eos_token: Optional[str] = None, ): assert audio_token is not None or hasattr(tokenizer, "audio_token"), ( "Either `audio_token` must be provided or tokenizer must have `audio_token` attribute." ) assert audio_bos_token is not None or hasattr(tokenizer, "audio_bos_token"), ( "Either `audio_bos_token` must be provided or tokenizer must have `audio_bos_token` attribute." ) assert audio_eos_token is not None or hasattr(tokenizer, "audio_eos_token"), ( "Either `audio_eos_token` must be provided or tokenizer must have `audio_eos_token` attribute." ) if chat_template is None: chat_template = tokenizer.chat_template def get_token(token_name: str) -> str: if not hasattr(tokenizer, token_name): raise ValueError( f"Tokenizer does not have attribute `{token_name}`. " "Please provide it as an argument to the processor." ) token = getattr(tokenizer, token_name) if not isinstance(token, str): raise TypeError( f"Expected token {token_name} to be a string, but got {type(token)}." ) return token self.audio_token = audio_token or get_token("audio_token") self.audio_bos_token = audio_bos_token or get_token("audio_bos_token") self.audio_eos_token = audio_eos_token or get_token("audio_eos_token") self.audio_token_id = cast(int, tokenizer.convert_tokens_to_ids(self.audio_token)) self.model_subsampling = model_subsampling self.audio_channels = audio_channels self.sem_n_fft = sem_n_fft self.min_audio_duration = min_audio_duration self.sampling_rate = feature_extractor.sampling_rate super().__init__(feature_extractor, tokenizer, chat_template=chat_template) self.feature_extractor: Wav2Vec2FeatureExtractor self.tokenizer: Union[Qwen2Tokenizer, Qwen2TokenizerFast] self.chat_template: Optional[Union[str, Dict[str, str]]] def _process_messages_for_chat_template( self, conversation, batch_images, batch_videos, batch_video_metadata, **mm_load_kwargs, ): if (sr := mm_load_kwargs.get("sampling_rate", None)) is not None: if sr != self.sampling_rate: raise ValueError( f"This model is trained with a sampling rate of {self.sampling_rate}, " f"but the sampling rate {sr} is used to load audio." ) return super()._process_messages_for_chat_template( conversation, batch_images, batch_videos, batch_video_metadata, **mm_load_kwargs, ) def _coerce_binaural(self, sample: Union[np.ndarray, torch.Tensor]) -> np.ndarray: if isinstance(sample, torch.Tensor): arr = sample.detach().cpu().numpy() elif isinstance(sample, np.ndarray): arr = sample elif isinstance(sample, str): raise TypeError( "Expected audio to be a numpy array or torch tensor, but got a string. " "If you passed audios through `apply_chat_template`, make sure the audio " "paths are valid URLs (http/https) or existing local files." ) else: raise TypeError( f"Expected audio to be a numpy array, torch tensor, or string, but got {type(sample)}." ) arr = np.asarray(arr, dtype=np.float32) C = self.audio_channels if arr.ndim == 1: arr = arr[None, :] # (1, T) elif arr.ndim == 2: if arr.shape[0] <= C: pass elif arr.shape[1] <= C: arr = arr.T elif arr.shape[0] > arr.shape[1]: arr = arr.T else: raise ValueError(f"Audio array must be 1D or 2D, got shape {arr.shape}.") if arr.shape[0] != C: arr = arr[:1] if arr.shape[0] == 1 and C > 1: arr = np.repeat(arr, C, axis=0) # up-mix mono -> C channels return arr def _pad_audio_batch(self, audios: List[np.ndarray]): raw_lengths = [a.shape[-1] for a in audios] min_t = max(1, int(np.ceil(self.sampling_rate * self.min_audio_duration))) lengths = [max(length, min_t) for length in raw_lengths] max_t = max(lengths) C = audios[0].shape[0] batch = torch.zeros(len(audios), C, max_t, dtype=torch.float32) for i, a in enumerate(audios): batch[i, :, : a.shape[-1]] = torch.from_numpy(a) return batch, torch.tensor(lengths, dtype=torch.long) def _load_audio_faithfully(self, src: str) -> np.ndarray: import torchaudio wav, sr = torchaudio.load(src) # (C, T), float32 if sr != self.sampling_rate: wav = torchaudio.functional.resample(wav, sr, self.sampling_rate) return wav.numpy() def _preload_audio_in_conversation(self, conversation): """Replace `{"type":"audio","path"/"url": ...}` with an in-memory binaural array (torchaudio) so the base `apply_chat_template` never routes the file through the mono-forcing `load_audio`. Array inputs are left untouched. Returns a new conversation; the caller's object is not mutated.""" def fix_message(msg): content = msg.get("content") if not isinstance(content, list): return msg new_content = [] for part in content: if ( isinstance(part, dict) and part.get("type") == "audio" and "audio" not in part and ("path" in part or "url" in part) ): src = part.get("path") or part.get("url") arr = self._load_audio_faithfully(src) new_part = {k: v for k, v in part.items() if k not in ("path", "url")} new_part["audio"] = arr new_content.append(new_part) else: new_content.append(part) return {**msg, "content": new_content} if conversation and isinstance(conversation[0], Mapping): return [fix_message(m) for m in conversation] # single conversation return [[fix_message(m) for m in conv] for conv in conversation] # batch def apply_chat_template( self, conversation: Union[List[dict], List[List[dict]]], chat_template: Optional[str] = None, **kwargs: Unpack[AllKwargsForChatTemplate], ) -> str: if conversation: conversation = self._preload_audio_in_conversation(conversation) if conversation: first_msgs = ( [conversation[0]] if isinstance(conversation[0], Mapping) else [conv[0] for conv in conversation if conv] ) for first_msg in first_msgs: if first_msg["role"] != "system": continue system_prompt = None if isinstance(first_msg["content"], str): system_prompt = first_msg["content"] elif isinstance(first_msg["content"], list): for part in first_msg["content"]: if isinstance(part, dict) and "text" in part: system_prompt = part["text"] break if system_prompt is not None and system_prompt != DEFAULT_SYSTEM_PROMPT: logging.warning( "The system prompt has been modified, which may reduce model " "performance. Prefer the default system prompt by omitting the " "system role from the input." ) return super().apply_chat_template(conversation, chat_template, **kwargs) def __call__( self, text: Optional[Union[str, List[str]]] = None, audio: Optional[Union[List[np.ndarray], List[torch.Tensor]]] = None, **kwargs: Unpack[MiDashengLMSpatialProcessorKwargs], ) -> BatchFeature: if text is None: raise ValueError("You need to specify `text` input to process.") elif isinstance(text, str): text = [text] elif not isinstance(text, list) or not isinstance(text[0], str): raise ValueError("Invalid input text. Please provide a string, or a list of strings") if kwargs.get("images", None) is not None or kwargs.get("videos", None) is not None: raise ValueError("This model does not support images or videos.") output_kwargs = self._merge_kwargs( MiDashengLMSpatialProcessorKwargs, tokenizer_init_kwargs=self.tokenizer.init_kwargs, **kwargs, ) audio_inputs: Dict[str, torch.Tensor] = {} if audio is not None: binaural = [self._coerce_binaural(sample) for sample in audio] num_audio_tokens = sum(sample.count(self.audio_token) for sample in text) num_audios = len(binaural) if num_audio_tokens != num_audios: raise ValueError( f"Found {num_audio_tokens} {self.audio_token} token(s) in provided " f"text but received {num_audios} audio(s)" ) audio_tensor, audio_length = self._pad_audio_batch(binaural) audio_inputs["input_values"] = audio_tensor # (B, C, T) audio_inputs["audio_length"] = audio_length # (B,) audio_lengths = audio_length.tolist() expanded_text = [] for sample in text: replace_str = [] while self.audio_token in sample: length = audio_lengths.pop(0) n_tok = calculate_mel_frames_dasheng( length, n_fft=self.sem_n_fft, model_subsampling=self.model_subsampling ) expanded_audio_token = self.audio_token * n_tok start = sample.find(self.audio_token) end = start + len(self.audio_token) has_bos = sample[start - len(self.audio_bos_token) : start] == self.audio_bos_token has_eos = sample[end : end + len(self.audio_eos_token)] == self.audio_eos_token if not has_bos and not has_eos: expanded_audio_token = ( self.audio_bos_token + expanded_audio_token + self.audio_eos_token ) replace_str.append(expanded_audio_token) sample = sample.replace(self.audio_token, "", 1) while "" in sample: sample = sample.replace("", replace_str.pop(0), 1) expanded_text.append(sample) text = expanded_text return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", "pt") inputs = self.tokenizer(text, **output_kwargs["text_kwargs"]) self._check_special_mm_tokens(text, BatchFeature(inputs), modalities=["audio"]) if audio is not None: inputs.update(audio_inputs) return BatchFeature(data={**inputs}, tensor_type=return_tensors) @property def model_input_names(self): tokenizer_input_names = self.tokenizer.model_input_names return list( dict.fromkeys(tokenizer_input_names + ["input_values", "audio_length"]) )