import importlib.util import os import re import sys import types from dataclasses import dataclass from typing import List, Optional, Sequence, Union import numpy as np import torch import torchaudio from transformers import AutoTokenizer, BatchEncoding @dataclass class MelConfig: mel_sr: int = 16000 mel_dim: int = 128 mel_n_fft: int = 400 mel_hop_length: int = 160 mel_dtype: torch.dtype = torch.bfloat16 use_whisper_feature_extractor: bool = True def load_chat_template(template_path: str, mossflux_path: str = None) -> List: if mossflux_path is None: template_dir = os.path.dirname(os.path.abspath(template_path)) current = template_dir while current and os.path.basename(current) != "mossLite": parent = os.path.dirname(current) if parent == current: break current = parent if os.path.basename(current) == "mossLite": mossflux_path = os.path.join(current, "mossflux") if mossflux_path and mossflux_path not in sys.path: sys.path.insert(0, mossflux_path) spec = importlib.util.spec_from_file_location("chat_template_module", template_path) module = importlib.util.module_from_spec(spec) sys.modules["chat_template_module"] = module spec.loader.exec_module(module) return module.chat_template class MossAudioProcessor: _AUDIO_SPAN_RE = re.compile(r"<\|audio_bos\|>(?:<\|AUDIO\|>)+<\|audio_eos\|>") _auto_class = None @classmethod def register_for_auto_class(cls, auto_class="AutoProcessor"): if not isinstance(auto_class, str): auto_class = auto_class.__name__ cls._auto_class = auto_class def __init__( self, tokenizer, *, mel_config: Optional[MelConfig] = None, template_path: Optional[str] = None, enable_time_marker: bool = True, audio_token_id: int = 151654, audio_start_id: int = 151669, audio_end_id: int = 151670, ): self._base_tokenizer = tokenizer self.tokenizer = tokenizer self.audio_token_id = int(audio_token_id) self.audio_start_id = int(audio_start_id) self.audio_end_id = int(audio_end_id) self.chat_template = ( None if template_path is None else load_chat_template(template_path) ) self.custom_texts = {} self.enable_time_marker = bool(enable_time_marker) self.config = mel_config or MelConfig() self._whisper_feature_extractor = None alias_map = { "<|AUDIO|>": self.audio_token_id, "<|audio_bos|>": self.audio_start_id, "<|audio_eos|>": self.audio_end_id, } orig_convert_tokens_to_ids = self.tokenizer.convert_tokens_to_ids def _patched_convert_tokens_to_ids(tokenizer_self, tokens): if isinstance(tokens, (list, tuple)): converted = [ _patched_convert_tokens_to_ids(tokenizer_self, token) for token in tokens ] return converted if isinstance(tokens, list) else tuple(converted) if isinstance(tokens, str) and tokens in alias_map: return alias_map[tokens] return orig_convert_tokens_to_ids(tokens) self.tokenizer.convert_tokens_to_ids = types.MethodType( _patched_convert_tokens_to_ids, self.tokenizer ) self._digit_token_ids = { "0": 15, "1": 16, "2": 17, "3": 18, "4": 19, "5": 20, "6": 21, "7": 22, "8": 23, "9": 24, } self.audio_tokens_per_second = 12.5 self.time_marker_every_seconds = 2 self.time_marker_every_audio_tokens = int( self.audio_tokens_per_second * self.time_marker_every_seconds ) self.model_input_names = [ "input_ids", "attention_mask", "audio_data", "audio_data_seqlens", ] @classmethod def from_pretrained(cls, pretrained_model_name_or_path, **kwargs): tokenizer_kwargs = {} for key in ["cache_dir", "revision", "token", "local_files_only"]: if key in kwargs: tokenizer_kwargs[key] = kwargs[key] tokenizer = AutoTokenizer.from_pretrained( pretrained_model_name_or_path, use_fast=False, **tokenizer_kwargs, ) mel_config = kwargs.pop("mel_config", None) template_path = kwargs.pop("template_path", None) enable_time_marker = kwargs.pop("enable_time_marker", False) audio_token_id = kwargs.pop("audio_token_id", 151654) audio_start_id = kwargs.pop("audio_start_id", 151669) audio_end_id = kwargs.pop("audio_end_id", 151670) return cls( tokenizer, mel_config=mel_config, template_path=template_path, enable_time_marker=enable_time_marker, audio_token_id=audio_token_id, audio_start_id=audio_start_id, audio_end_id=audio_end_id, ) def load_template(self, template_path: str): self.chat_template = load_chat_template(template_path) return self def set_custom_text(self, key: str, text: str): self.custom_texts[key] = text return self def clear_custom_text(self, key: Optional[str] = None): if key is None: self.custom_texts.clear() else: self.custom_texts.pop(key, None) return self def _template_requires_audio(self) -> bool: if self.chat_template is None: return False for segment in self.chat_template: if segment.type in {"audio_contiguous", "audio_token"}: return True return False @staticmethod def _conv3_downsample_len(raw_mel_len: int) -> int: def conv_out_len(length: int) -> int: return (length - 1) // 2 + 1 length1 = conv_out_len(int(raw_mel_len)) length2 = conv_out_len(length1) length3 = conv_out_len(length2) return int(length3) def _get_whisper_feature_extractor(self): if self._whisper_feature_extractor is not None: return self._whisper_feature_extractor from transformers.models.whisper.feature_extraction_whisper import ( WhisperFeatureExtractor, ) self._whisper_feature_extractor = WhisperFeatureExtractor( feature_size=int(self.config.mel_dim), sampling_rate=int(self.config.mel_sr), hop_length=int(self.config.mel_hop_length), n_fft=int(self.config.mel_n_fft), ) return self._whisper_feature_extractor def _extract_mel(self, audio: Union[np.ndarray, torch.Tensor]) -> torch.Tensor: if isinstance(audio, np.ndarray): wav = torch.from_numpy(audio) else: wav = audio wav = wav.to(dtype=torch.float32) if wav.dim() == 1: wav = wav.unsqueeze(0) if bool(getattr(self.config, "use_whisper_feature_extractor", False)): fe = self._get_whisper_feature_extractor() wav_np = wav.detach().to("cpu", torch.float32).contiguous().numpy() if wav_np.ndim == 2: wav_np = wav_np[0] feats = fe._np_extract_fbank_features(wav_np[None, ...], device="cpu") mel = torch.from_numpy(feats[0]) return mel.to(dtype=self.config.mel_dtype) def _get_time_marker_token_ids(self, second: int) -> List[int]: return [self._digit_token_ids[digit] for digit in str(second)] def _build_audio_tokens_with_time_markers(self, audio_seq_len: int) -> List[int]: total_duration_seconds = audio_seq_len / self.audio_tokens_per_second num_full_seconds = int(total_duration_seconds) token_ids: List[int] = [] audio_tokens_consumed = 0 for second in range( self.time_marker_every_seconds, num_full_seconds + 1, self.time_marker_every_seconds, ): marker_pos = ( second // self.time_marker_every_seconds ) * self.time_marker_every_audio_tokens audio_segment_len = marker_pos - audio_tokens_consumed if audio_segment_len > 0: token_ids.extend([self.audio_token_id] * audio_segment_len) audio_tokens_consumed += audio_segment_len token_ids.extend(self._get_time_marker_token_ids(second)) remaining = audio_seq_len - audio_tokens_consumed if remaining > 0: token_ids.extend([self.audio_token_id] * remaining) return token_ids def _build_audio_placeholder_ids(self, num_audio_tokens: int) -> List[int]: if self.enable_time_marker: return self._build_audio_tokens_with_time_markers(num_audio_tokens) return [self.audio_token_id] * num_audio_tokens def _build_input_from_template( self, num_audio_tokens: int, include_answer: bool = False ) -> List[int]: if self.chat_template is None: raise ValueError("Chat template not loaded.") input_ids: List[int] = [] for segment in self.chat_template: seg_type = segment.type if seg_type == "constant_text_token": input_ids.extend(segment.text_ids.tolist()) elif seg_type in {"audio_contiguous", "audio_token"}: input_ids.extend(self._build_audio_placeholder_ids(num_audio_tokens)) elif seg_type == "text_token": text_token_key = segment.text_token_key if "answer" in text_token_key.lower() and not include_answer: break if text_token_key not in self.custom_texts: break text_ids = self._base_tokenizer.encode( self.custom_texts[text_token_key], add_special_tokens=False ) input_ids.extend(text_ids) return input_ids def _build_default_prompt(self, text: str, has_audio: bool) -> str: if has_audio: return ( "<|im_start|>system\n" "You are a helpful assistant.<|im_end|>\n" "<|im_start|>user\n" "<|audio_bos|><|AUDIO|><|audio_eos|>\n" f"{text}<|im_end|>\n" "<|im_start|>assistant\n" ) return ( "<|im_start|>system\n" "You are a helpful assistant.<|im_end|>\n" "<|im_start|>user\n" f"{text}<|im_end|>\n" "<|im_start|>assistant\n" ) def _build_input_from_prompt(self, prompt: str, token_lens: List[int]) -> List[int]: spans = list(self._AUDIO_SPAN_RE.finditer(prompt)) if len(spans) != len(token_lens): raise ValueError( f"Audio placeholder count mismatch: found {len(spans)} spans in text, " f"but got {len(token_lens)} audio inputs." ) input_ids: List[int] = [] cursor = 0 for index, match in enumerate(spans): prefix = prompt[cursor : match.start()] if prefix: input_ids.extend( self._base_tokenizer.encode(prefix, add_special_tokens=False) ) input_ids.append(self.audio_start_id) input_ids.extend(self._build_audio_placeholder_ids(int(token_lens[index]))) input_ids.append(self.audio_end_id) cursor = match.end() suffix = prompt[cursor:] if suffix: input_ids.extend( self._base_tokenizer.encode(suffix, add_special_tokens=False) ) return input_ids def __call__( self, *, text: Union[str, Sequence[str], None] = None, audios: Optional[Sequence[Union[np.ndarray, torch.Tensor]]] = None, audio: Optional[Sequence[Union[np.ndarray, torch.Tensor]]] = None, return_tensors: str = "pt", **kwargs, ): if isinstance(text, (list, tuple)): if len(text) != 1: raise ValueError(f"Expected text batch size 1, got {len(text)}") prompt_text = text[0] else: prompt_text = text audio_list = audios if audios is not None else audio audio_list = [] if audio_list is None else list(audio_list) mels: List[torch.Tensor] = [] raw_lengths: List[int] = [] token_lens: List[int] = [] for one_audio in audio_list: mel = self._extract_mel(one_audio) raw_len = int(mel.shape[-1]) mels.append(mel) raw_lengths.append(raw_len) token_lens.append(self._conv3_downsample_len(raw_len)) if mels: max_length = max(raw_lengths) audio_batch = torch.zeros( (len(mels), self.config.mel_dim, max_length), dtype=self.config.mel_dtype, ) for index, mel in enumerate(mels): audio_batch[index, :, : mel.shape[-1]] = mel seqlens_tensor = torch.tensor(raw_lengths, dtype=torch.long) else: audio_batch = None seqlens_tensor = None if prompt_text is not None: if self._AUDIO_SPAN_RE.search(prompt_text) is None and audio_list: prompt_text = self._build_default_prompt(prompt_text, has_audio=True) elif self._AUDIO_SPAN_RE.search(prompt_text) is None and not audio_list: prompt_text = self._build_default_prompt(prompt_text, has_audio=False) input_ids_list = self._build_input_from_prompt(prompt_text, token_lens) elif self.chat_template is not None: input_ids_list = self._build_input_from_template( token_lens[0] if token_lens else 0 ) else: raise ValueError( "Either provide text or load a chat_template before calling the processor." ) input_ids_tensor = torch.tensor([input_ids_list], dtype=torch.long) attention_mask_tensor = torch.ones_like(input_ids_tensor) data = { "input_ids": input_ids_tensor, "attention_mask": attention_mask_tensor, } if audio_batch is not None: data["audio_data"] = audio_batch data["audio_data_seqlens"] = seqlens_tensor return BatchEncoding(data=data, tensor_type=return_tensors) def batch_decode(self, *args, **kwargs): return self._base_tokenizer.batch_decode(*args, **kwargs) def decode(self, *args, **kwargs): return self._base_tokenizer.decode(*args, **kwargs) __all__ = ["MelConfig", "MossAudioProcessor"]