OpenMOSS-Team/MOSS-Music-8B-Instruct

1import importlib.util2import os3import re4import sys5import types6from dataclasses import dataclass7from typing import List, Optional, Sequence, Union8 9import numpy as np10import torch11import torchaudio  # noqa: F40112from transformers import AutoTokenizer, BatchEncoding13 14 15@dataclass16class MelConfig:17    mel_sr: int = 1600018    mel_dim: int = 12819    mel_n_fft: int = 40020    mel_hop_length: int = 16021    mel_dtype: torch.dtype = torch.bfloat1622    use_whisper_feature_extractor: bool = True23 24 25def load_chat_template(template_path: str, mossflux_path: str = None) -> List:26    if mossflux_path is None:27        template_dir = os.path.dirname(os.path.abspath(template_path))28        current = template_dir29        while current and os.path.basename(current) != "mossLite":30            parent = os.path.dirname(current)31            if parent == current:32                break33            current = parent34        if os.path.basename(current) == "mossLite":35            mossflux_path = os.path.join(current, "mossflux")36 37    if mossflux_path and mossflux_path not in sys.path:38        sys.path.insert(0, mossflux_path)39 40    spec = importlib.util.spec_from_file_location("chat_template_module", template_path)41    module = importlib.util.module_from_spec(spec)42    sys.modules["chat_template_module"] = module43    spec.loader.exec_module(module)44    return module.chat_template45 46 47class MossMusicProcessor:48    _AUDIO_SPAN_RE = re.compile(r"<\|audio_bos\|>(?:<\|AUDIO\|>)+<\|audio_eos\|>")49    _auto_class = None50 51    @classmethod52    def register_for_auto_class(cls, auto_class="AutoProcessor"):53        if not isinstance(auto_class, str):54            auto_class = auto_class.__name__55        cls._auto_class = auto_class56 57    def __init__(58        self,59        tokenizer,60        *,61        mel_config: Optional[MelConfig] = None,62        template_path: Optional[str] = None,63        enable_time_marker: bool = True,64        audio_token_id: int = 151654,65        audio_start_id: int = 151669,66        audio_end_id: int = 151670,67    ):68        self._base_tokenizer = tokenizer69        self.tokenizer = tokenizer70        self.audio_token_id = int(audio_token_id)71        self.audio_start_id = int(audio_start_id)72        self.audio_end_id = int(audio_end_id)73        self.chat_template = (74            None if template_path is None else load_chat_template(template_path)75        )76        self.custom_texts = {}77        self.enable_time_marker = bool(enable_time_marker)78        self.config = mel_config or MelConfig()79        self._whisper_feature_extractor = None80 81        alias_map = {82            "<|AUDIO|>": self.audio_token_id,83            "<|audio_bos|>": self.audio_start_id,84            "<|audio_eos|>": self.audio_end_id,85        }86        orig_convert_tokens_to_ids = self.tokenizer.convert_tokens_to_ids87 88        def _patched_convert_tokens_to_ids(tokenizer_self, tokens):89            if isinstance(tokens, (list, tuple)):90                converted = [91                    _patched_convert_tokens_to_ids(tokenizer_self, token)92                    for token in tokens93                ]94                return converted if isinstance(tokens, list) else tuple(converted)95            if isinstance(tokens, str) and tokens in alias_map:96                return alias_map[tokens]97            return orig_convert_tokens_to_ids(tokens)98 99        self.tokenizer.convert_tokens_to_ids = types.MethodType(100            _patched_convert_tokens_to_ids, self.tokenizer101        )102 103        self._digit_token_ids = {104            "0": 15,105            "1": 16,106            "2": 17,107            "3": 18,108            "4": 19,109            "5": 20,110            "6": 21,111            "7": 22,112            "8": 23,113            "9": 24,114        }115        self.audio_tokens_per_second = 12.5116        self.time_marker_every_seconds = 2117        self.time_marker_every_audio_tokens = int(118            self.audio_tokens_per_second * self.time_marker_every_seconds119        )120        self.model_input_names = [121            "input_ids",122            "attention_mask",123            "audio_data",124            "audio_data_seqlens",125        ]126 127    @classmethod128    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):129        tokenizer_kwargs = {}130        for key in ["cache_dir", "revision", "token", "local_files_only"]:131            if key in kwargs:132                tokenizer_kwargs[key] = kwargs[key]133 134        tokenizer = AutoTokenizer.from_pretrained(135            pretrained_model_name_or_path,136            use_fast=False,137            **tokenizer_kwargs,138        )139 140        mel_config = kwargs.pop("mel_config", None)141        template_path = kwargs.pop("template_path", None)142        enable_time_marker = kwargs.pop("enable_time_marker", False)143        audio_token_id = kwargs.pop("audio_token_id", 151654)144        audio_start_id = kwargs.pop("audio_start_id", 151669)145        audio_end_id = kwargs.pop("audio_end_id", 151670)146 147        return cls(148            tokenizer,149            mel_config=mel_config,150            template_path=template_path,151            enable_time_marker=enable_time_marker,152            audio_token_id=audio_token_id,153            audio_start_id=audio_start_id,154            audio_end_id=audio_end_id,155        )156 157    def load_template(self, template_path: str):158        self.chat_template = load_chat_template(template_path)159        return self160 161    def set_custom_text(self, key: str, text: str):162        self.custom_texts[key] = text163        return self164 165    def clear_custom_text(self, key: Optional[str] = None):166        if key is None:167            self.custom_texts.clear()168        else:169            self.custom_texts.pop(key, None)170        return self171 172    def _template_requires_audio(self) -> bool:173        if self.chat_template is None:174            return False175        for segment in self.chat_template:176            if segment.type in {"audio_contiguous", "audio_token"}:177                return True178        return False179 180    @staticmethod181    def _conv3_downsample_len(raw_mel_len: int) -> int:182        def conv_out_len(length: int) -> int:183            return (length - 1) // 2 + 1184 185        length1 = conv_out_len(int(raw_mel_len))186        length2 = conv_out_len(length1)187        length3 = conv_out_len(length2)188        return int(length3)189 190    def _get_whisper_feature_extractor(self):191        if self._whisper_feature_extractor is not None:192            return self._whisper_feature_extractor193 194        from transformers.models.whisper.feature_extraction_whisper import (195            WhisperFeatureExtractor,196        )197 198        self._whisper_feature_extractor = WhisperFeatureExtractor(199            feature_size=int(self.config.mel_dim),200            sampling_rate=int(self.config.mel_sr),201            hop_length=int(self.config.mel_hop_length),202            n_fft=int(self.config.mel_n_fft),203        )204        return self._whisper_feature_extractor205 206    def _extract_mel(self, audio: Union[np.ndarray, torch.Tensor]) -> torch.Tensor:207        if isinstance(audio, np.ndarray):208            wav = torch.from_numpy(audio)209        else:210            wav = audio211        wav = wav.to(dtype=torch.float32)212        if wav.dim() == 1:213            wav = wav.unsqueeze(0)214 215        if bool(getattr(self.config, "use_whisper_feature_extractor", False)):216            fe = self._get_whisper_feature_extractor()217            wav_np = wav.detach().to("cpu", torch.float32).contiguous().numpy()218            if wav_np.ndim == 2:219                wav_np = wav_np[0]220            feats = fe._np_extract_fbank_features(wav_np[None, ...], device="cpu")221            mel = torch.from_numpy(feats[0])222 223        return mel.to(dtype=self.config.mel_dtype)224 225    def _get_time_marker_token_ids(self, second: int) -> List[int]:226        return [self._digit_token_ids[digit] for digit in str(second)]227 228    def _build_audio_tokens_with_time_markers(self, audio_seq_len: int) -> List[int]:229        total_duration_seconds = audio_seq_len / self.audio_tokens_per_second230        num_full_seconds = int(total_duration_seconds)231 232        token_ids: List[int] = []233        audio_tokens_consumed = 0234        for second in range(235            self.time_marker_every_seconds,236            num_full_seconds + 1,237            self.time_marker_every_seconds,238        ):239            marker_pos = (240                second // self.time_marker_every_seconds241            ) * self.time_marker_every_audio_tokens242            audio_segment_len = marker_pos - audio_tokens_consumed243            if audio_segment_len > 0:244                token_ids.extend([self.audio_token_id] * audio_segment_len)245                audio_tokens_consumed += audio_segment_len246            token_ids.extend(self._get_time_marker_token_ids(second))247 248        remaining = audio_seq_len - audio_tokens_consumed249        if remaining > 0:250            token_ids.extend([self.audio_token_id] * remaining)251        return token_ids252 253    def _build_audio_placeholder_ids(self, num_audio_tokens: int) -> List[int]:254        if self.enable_time_marker:255            return self._build_audio_tokens_with_time_markers(num_audio_tokens)256        return [self.audio_token_id] * num_audio_tokens257 258    def _build_input_from_template(259        self, num_audio_tokens: int, include_answer: bool = False260    ) -> List[int]:261        if self.chat_template is None:262            raise ValueError("Chat template not loaded.")263 264        input_ids: List[int] = []265        for segment in self.chat_template:266            seg_type = segment.type267            if seg_type == "constant_text_token":268                input_ids.extend(segment.text_ids.tolist())269            elif seg_type in {"audio_contiguous", "audio_token"}:270                input_ids.extend(self._build_audio_placeholder_ids(num_audio_tokens))271            elif seg_type == "text_token":272                text_token_key = segment.text_token_key273                if "answer" in text_token_key.lower() and not include_answer:274                    break275                if text_token_key not in self.custom_texts:276                    break277                text_ids = self._base_tokenizer.encode(278                    self.custom_texts[text_token_key], add_special_tokens=False279                )280                input_ids.extend(text_ids)281 282        return input_ids283 284    def _build_default_prompt(self, text: str, has_audio: bool) -> str:285        if has_audio:286            return (287                "<|im_start|>system\n"288                "You are a helpful assistant.<|im_end|>\n"289                "<|im_start|>user\n"290                "<|audio_bos|><|AUDIO|><|audio_eos|>\n"291                f"{text}<|im_end|>\n"292                "<|im_start|>assistant\n"293            )294        return (295            "<|im_start|>system\n"296            "You are a helpful assistant.<|im_end|>\n"297            "<|im_start|>user\n"298            f"{text}<|im_end|>\n"299            "<|im_start|>assistant\n"300        )301 302    def _build_input_from_prompt(self, prompt: str, token_lens: List[int]) -> List[int]:303        spans = list(self._AUDIO_SPAN_RE.finditer(prompt))304        if len(spans) != len(token_lens):305            raise ValueError(306                f"Audio placeholder count mismatch: found {len(spans)} spans in text, "307                f"but got {len(token_lens)} audio inputs."308            )309 310        input_ids: List[int] = []311        cursor = 0312        for index, match in enumerate(spans):313            prefix = prompt[cursor : match.start()]314            if prefix:315                input_ids.extend(316                    self._base_tokenizer.encode(prefix, add_special_tokens=False)317                )318 319            input_ids.append(self.audio_start_id)320            input_ids.extend(self._build_audio_placeholder_ids(int(token_lens[index])))321            input_ids.append(self.audio_end_id)322            cursor = match.end()323 324        suffix = prompt[cursor:]325        if suffix:326            input_ids.extend(327                self._base_tokenizer.encode(suffix, add_special_tokens=False)328            )329        return input_ids330 331    def __call__(332        self,333        *,334        text: Union[str, Sequence[str], None] = None,335        audios: Optional[Sequence[Union[np.ndarray, torch.Tensor]]] = None,336        audio: Optional[Sequence[Union[np.ndarray, torch.Tensor]]] = None,337        return_tensors: str = "pt",338        **kwargs,339    ):340        if isinstance(text, (list, tuple)):341            if len(text) != 1:342                raise ValueError(f"Expected text batch size 1, got {len(text)}")343            prompt_text = text[0]344        else:345            prompt_text = text346 347        audio_list = audios if audios is not None else audio348        audio_list = [] if audio_list is None else list(audio_list)349 350        mels: List[torch.Tensor] = []351        raw_lengths: List[int] = []352        token_lens: List[int] = []353        for one_audio in audio_list:354            mel = self._extract_mel(one_audio)355            raw_len = int(mel.shape[-1])356            mels.append(mel)357            raw_lengths.append(raw_len)358            token_lens.append(self._conv3_downsample_len(raw_len))359 360        if mels:361            max_length = max(raw_lengths)362            audio_batch = torch.zeros(363                (len(mels), self.config.mel_dim, max_length),364                dtype=self.config.mel_dtype,365            )366            for index, mel in enumerate(mels):367                audio_batch[index, :, : mel.shape[-1]] = mel368            seqlens_tensor = torch.tensor(raw_lengths, dtype=torch.long)369        else:370            audio_batch = None371            seqlens_tensor = None372 373        if prompt_text is not None:374            if self._AUDIO_SPAN_RE.search(prompt_text) is None and audio_list:375                prompt_text = self._build_default_prompt(prompt_text, has_audio=True)376            elif self._AUDIO_SPAN_RE.search(prompt_text) is None and not audio_list:377                prompt_text = self._build_default_prompt(prompt_text, has_audio=False)378            input_ids_list = self._build_input_from_prompt(prompt_text, token_lens)379        elif self.chat_template is not None:380            input_ids_list = self._build_input_from_template(381                token_lens[0] if token_lens else 0382            )383        else:384            raise ValueError(385                "Either provide text or load a chat_template before calling the processor."386            )387 388        input_ids_tensor = torch.tensor([input_ids_list], dtype=torch.long)389        attention_mask_tensor = torch.ones_like(input_ids_tensor)390 391        data = {392            "input_ids": input_ids_tensor,393            "attention_mask": attention_mask_tensor,394        }395        if audio_batch is not None:396            data["audio_data"] = audio_batch397            data["audio_data_seqlens"] = seqlens_tensor398        return BatchEncoding(data=data, tensor_type=return_tensors)399 400    def batch_decode(self, *args, **kwargs):401        return self._base_tokenizer.batch_decode(*args, **kwargs)402 403    def decode(self, *args, **kwargs):404        return self._base_tokenizer.decode(*args, **kwargs)405 406 407__all__ = ["MelConfig", "MossMusicProcessor"]408