"""MiMo-V2-ASR multimodal processor. Audio preprocessing is delegated to :class:`MiMoAudioPipeline`; this processor only handles the special-token contract and content interleaving. """ import asyncio import re from dataclasses import dataclass from typing import List, Literal, Union import numpy as np import torch from sglang.srt.managers.schedule_batch import ( Modality, MultimodalDataItem, MultimodalProcessorOutput, ) from sglang.srt.models.mimo_v2_asr import MiMoV2ASRForCausalLM from sglang.srt.multimodal.processors.base_processor import ( BaseMultimodalProcessor, MultimodalSpecialTokens, ) from sglang.srt.multimodal.processors.mimo_audio import ( AudioInput, MiMoAudioPipeline, ) from sglang.utils import logger TextInput = str | list[int] @dataclass class _Content: type: Literal["text", "audio"] content: TextInput | AudioInput class MiMoV2ASRProcessor(BaseMultimodalProcessor): """ASR-only MiMo processor. Wires three special tokens into the input id stream around each audio span: ``<|sosp|> <|empty|>* <|eosp|>``. The actual mel/codec preparation is owned by :class:`MiMoAudioPipeline`, which is shared with the multimodal MiMo-V2 processor. """ models = [MiMoV2ASRForCausalLM] AUDIO_PAD_TOKEN = "<|empty|>" AUDIO_START_TOKEN = "<|sosp|>" AUDIO_END_TOKEN = "<|eosp|>" AUDIO_REGEX = re.compile(r"<\|sosp\|>(?:<\|empty\|>)+<\|eosp\|>") def __init__(self, hf_config, server_args, _processor, *args, **kwargs): super().__init__(hf_config, server_args, _processor, *args, **kwargs) self.tokenizer = _processor self.audio_pipeline = MiMoAudioPipeline( audio_token_id=self._resolve_special_token_id(self.AUDIO_PAD_TOKEN), audio_start_token_id=self._resolve_special_token_id(self.AUDIO_START_TOKEN), audio_end_token_id=self._resolve_special_token_id(self.AUDIO_END_TOKEN), audio_sampling_rate=24000, ) self.mm_tokens = MultimodalSpecialTokens( audio_token=f"{self.AUDIO_START_TOKEN}{self.AUDIO_PAD_TOKEN}{self.AUDIO_END_TOKEN}", audio_token_id=self.audio_token_id, audio_token_regex=self.AUDIO_REGEX, ).build(_processor) def __getattr__(self, name): # Delegate audio_pipeline fields so callers can use self.audio_token_id # etc. directly. Only triggers when normal attribute lookup fails; # __dict__.get avoids recursion before audio_pipeline is assigned. pipeline = self.__dict__.get("audio_pipeline") if pipeline is not None and hasattr(pipeline, name): return getattr(pipeline, name) raise AttributeError(name) def _resolve_special_token_id(self, name: str) -> int: tid = self.tokenizer.convert_tokens_to_ids(name) if tid is None or tid == self.tokenizer.unk_token_id: raise ValueError( f"tokenizer missing required special token {name!r}; " "checkpoint vocab does not match MiMo-V2-ASR" ) return int(tid) def _process_contents(self, contents: List[_Content]): """Run pipeline + tokenizer over an interleaved content list. Returns ``(input_ids: Tensor[L], audio_inputs: list[Tensor], position_ids: Tensor[3,L], rope_deltas: Tensor[1,1])``. """ input_ids: List[int] = [] audio_inputs: List[torch.Tensor] = [] for content in contents: if content.type == "text": if isinstance(content.content, str): input_ids.extend(self.tokenizer.encode(content.content)) else: input_ids.extend(content.content) elif content.type == "audio": result = self.audio_pipeline.process_audio_input(content.content) audio_inputs.append(result["audio_input"]) input_ids.extend(result["input_ids"]) ids = torch.as_tensor(input_ids) position_ids = torch.arange(ids.shape[0]).expand(3, -1) rope_deltas = torch.zeros((1, 1), dtype=torch.int32) return ids, audio_inputs, position_ids, rope_deltas def process_mm_data( self, input_text, images=None, videos=None, audios=None, **kwargs ) -> dict: if audios and not self.AUDIO_REGEX.search(input_text or ""): input_text = f"{self.mm_tokens.audio_token}{input_text or ''}" processed_audios: List[Union[tuple, torch.Tensor]] = [] if audios: for audio in audios: if isinstance(audio, np.ndarray): audio_tensor = torch.from_numpy(audio).float() elif isinstance(audio, torch.Tensor): audio_tensor = audio.float() else: processed_audios.append(audio) continue if audio_tensor.ndim == 1: processed_audios.append( (audio_tensor.cpu().contiguous(), self.audio_sampling_rate) ) else: processed_audios.append(audio_tensor.cpu().contiguous()) contents: List[_Content] = [] if input_text and processed_audios: multimodal_tokens_pattern = self.mm_tokens.get_combined_regex() text_parts = re.split(multimodal_tokens_pattern, input_text) audio_iter = iter(processed_audios) for text_part in text_parts: if multimodal_tokens_pattern.match(text_part): modality = self.mm_tokens.get_modality_of_token(text_part) if modality == Modality.AUDIO: try: audio = next(audio_iter) contents.append( _Content(type="audio", content=AudioInput(audio=audio)) ) except StopIteration: pass else: if text_part: contents.append(_Content(type="text", content=text_part)) else: contents.extend( _Content(type="audio", content=AudioInput(audio=audio)) for audio in processed_audios ) if not contents: ids = self.tokenizer( input_text or "", return_tensors="pt", add_special_tokens=True, ).input_ids return {"input_ids": ids} input_ids, audio_inputs, position_ids, rope_deltas = self._process_contents( contents ) ret: dict = { "input_ids": input_ids, "mrope_positions": position_ids, "mrope_position_delta": rope_deltas, } if audio_inputs: ret["audio_features"] = audio_inputs return ret async def process_mm_data_async( self, image_data: List[Union[str, bytes]], audio_data: List[Union[str, bytes]], input_text, request_obj, *args, **kwargs, ): if audio_data is None: audio_data = getattr(request_obj, "audio_data", []) if not audio_data: return None if not self.AUDIO_REGEX.search(input_text): input_text = f"{self.mm_tokens.audio_token}{input_text}" base_output = await self.load_mm_data( prompt=input_text, image_data=[], video_data=[], audio_data=audio_data, multimodal_tokens=self.mm_tokens, audio_sample_rate=self.audio_sampling_rate, ) multimodal_tokens_pattern = self.mm_tokens.get_combined_regex() raw_audio_data = audio_data or [] loaded_audio_iter = iter(base_output.audios) raw_audio_iter = iter(raw_audio_data) text_parts = re.split(multimodal_tokens_pattern, base_output.input_text) contents: List[_Content] = [] for text_part in text_parts: if multimodal_tokens_pattern.match(text_part): modality = self.mm_tokens.get_modality_of_token(text_part) assert modality is not None if modality == Modality.AUDIO: loaded_audio = next(loaded_audio_iter) raw_audio_item = next(raw_audio_iter) if isinstance(loaded_audio, np.ndarray): audio_source = loaded_audio elif isinstance(raw_audio_item, dict): audio_source = raw_audio_item.get("url", loaded_audio) elif isinstance(raw_audio_item, (str, bytes, torch.Tensor)): audio_source = raw_audio_item else: raise ValueError( f"unsupported audio item: loaded={type(loaded_audio).__name__}, " f"raw={type(raw_audio_item).__name__}" ) contents.append( _Content( type="audio", content=AudioInput(audio=audio_source), ) ) else: if text_part: contents.append(_Content(type="text", content=text_part)) loop = asyncio.get_running_loop() try: input_ids, audio_inputs, position_ids, rope_deltas = ( await loop.run_in_executor( self.io_executor, lambda: self._process_contents(contents), ) ) except RuntimeError as e: logger.error(f"MiMo ASR processor failed in process_mm_data_async: {e}") raise ValueError(f"Multimodal data is corrupted or cannot be decoded: {e}") input_ids_flat = input_ids.flatten() if audio_inputs: mm_items = [ MultimodalDataItem( modality=Modality.AUDIO, feature=audio_inputs, offsets=self.get_mm_items_offset( input_ids=input_ids_flat, mm_token_id=self.audio_token_id, ), ) ] else: mm_items = [] return MultimodalProcessorOutput( mm_items=mm_items, input_ids=input_ids_flat.tolist(), audio_token_id=self.audio_token_id, audio_start_id=self.audio_start_token_id, audio_end_id=self.audio_end_token_id, mrope_positions=position_ids, mrope_position_delta=rope_deltas, )