processing_midashenglm_spatial.py

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, "<placeholder>", 1)
                while "<placeholder>" in sample:
                    sample = sample.replace("<placeholder>", 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"])
        )