processing_moss_transcribe_diarize.py

"""Processor for MOSS-Transcribe-Diarize audio-text inference."""

from __future__ import annotations

from typing import Optional, Union

import numpy as np
import torch
from transformers.feature_extraction_utils import BatchFeature
from transformers.processing_utils import ProcessingKwargs, ProcessorMixin, Unpack


AUDIO_PAD_TOKEN = "<|audio_pad|>"
AUDIO_START_TOKEN = "<|audio_start|>"
AUDIO_END_TOKEN = "<|audio_end|>"
WHISPER_ENCODER_STRIDE = 2


class MossTranscribeDiarizeProcessorKwargs(ProcessingKwargs, total=False):
    _defaults = {
        "common_kwargs": {
            "return_tensors": "pt",
        },
        "audio_kwargs": {},
    }


def _audio_to_numpy(audio: Union[np.ndarray, torch.Tensor]) -> np.ndarray:
    if torch.is_tensor(audio):
        audio = audio.detach().cpu().numpy()
    audio = np.asarray(audio, dtype=np.float32)
    if audio.ndim > 1:
        audio = np.squeeze(audio)
    if audio.ndim == 0:
        audio = audio.reshape(1)
    if audio.ndim != 1:
        raise ValueError(f"Expected mono audio with shape (num_samples,), got shape {audio.shape}.")
    if audio.shape[0] == 0:
        raise ValueError("Audio must contain at least one sample.")
    return audio


def _pad_or_trim_audio(audio: np.ndarray, length: int) -> np.ndarray:
    if audio.shape[0] > length:
        audio = audio[:length]
    elif audio.shape[0] < length:
        audio = np.pad(audio, (0, length - audio.shape[0]))
    return audio.astype(np.float32, copy=False)


def _compute_audio_token_length(num_samples: int, feature_extractor, audio_merge_size: int) -> int:
    stride = int(feature_extractor.hop_length) * WHISPER_ENCODER_STRIDE * int(audio_merge_size)
    return (int(num_samples) - 1) // stride + 1


def _chunk_audio(
    feature_extractor,
    audio: Union[np.ndarray, torch.Tensor],
    audio_merge_size: int,
) -> tuple[np.ndarray, list[int]]:
    audio = _audio_to_numpy(audio)
    n_samples = int(feature_extractor.n_samples)

    chunks, token_lengths = [], []
    for start in range(0, audio.shape[0], n_samples):
        chunk = audio[start : start + n_samples]
        token_lengths.append(_compute_audio_token_length(chunk.shape[0], feature_extractor, audio_merge_size))
        chunks.append(_pad_or_trim_audio(chunk, n_samples))
    return np.stack(chunks), token_lengths


def _audios_to_input_features(
    feature_extractor,
    audios: list[Union[np.ndarray, torch.Tensor]],
    *,
    audio_merge_size: int,
    feature_extractor_kwargs: Optional[dict] = None,
) -> tuple[torch.Tensor, torch.LongTensor, torch.LongTensor]:
    feature_batches, feature_lengths, chunk_mapping = [], [], []
    feature_extractor_kwargs = dict(feature_extractor_kwargs or {})
    feature_extractor_kwargs.update(
        {
            "sampling_rate": int(feature_extractor.sampling_rate),
            "padding": "max_length",
            "return_tensors": "pt",
        }
    )

    for audio_idx, audio in enumerate(audios):
        chunks, token_lengths = _chunk_audio(feature_extractor, audio, audio_merge_size)
        features = feature_extractor(
            list(chunks),
            **feature_extractor_kwargs,
        )["input_features"]
        feature_batches.append(features)
        feature_lengths.extend(token_lengths)
        chunk_mapping.extend([audio_idx] * len(token_lengths))

    if feature_batches:
        input_features = torch.cat(feature_batches, dim=0)
    else:
        input_features = torch.empty(
            (0, int(feature_extractor.feature_size), int(feature_extractor.nb_max_frames)),
        )

    length_device = input_features.device
    return (
        input_features,
        torch.tensor(feature_lengths, dtype=torch.long, device=length_device),
        torch.tensor(chunk_mapping, dtype=torch.long, device=length_device),
    )


class MossTranscribeDiarizeProcessor(ProcessorMixin):
    """Build MOSS-Transcribe-Diarize model inputs from text prompts and raw waveforms.

    The model consumes log-mel ``input_features``. This processor owns the raw
    waveform preprocessing, audio placeholder expansion, and optional numeric
    time anchors inside the audio span.
    """

    attributes = ["feature_extractor", "tokenizer"]
    feature_extractor_class = "AutoFeatureExtractor"
    tokenizer_class = "AutoTokenizer"

    model_input_names = [
        "input_ids",
        "attention_mask",
        "input_features",
        "audio_feature_lengths",
        "audio_chunk_mapping",
    ]

    def __init__(
        self,
        feature_extractor=None,
        tokenizer=None,
        audio_tokens_per_second: float = 12.5,
        audio_merge_size: int = 4,
        time_marker_every_seconds: int = 2,
        enable_time_marker: bool = True,
        chat_template: Optional[str] = None,
    ):
        if feature_extractor is None:
            raise ValueError("MossTranscribeDiarizeProcessor requires a feature_extractor.")
        if tokenizer is None:
            raise ValueError("MossTranscribeDiarizeProcessor requires a tokenizer.")
        super().__init__(feature_extractor, tokenizer, chat_template=chat_template)
        self.audio_tokens_per_second = audio_tokens_per_second
        self.audio_merge_size = int(audio_merge_size)
        self.time_marker_every_seconds = time_marker_every_seconds
        self.enable_time_marker = enable_time_marker
        self.audio_token = AUDIO_PAD_TOKEN if not hasattr(tokenizer, "audio_token") else tokenizer.audio_token
        self.audio_start_token = (
            AUDIO_START_TOKEN if not hasattr(tokenizer, "audio_start_token") else tokenizer.audio_start_token
        )
        self.audio_end_token = AUDIO_END_TOKEN if not hasattr(tokenizer, "audio_end_token") else tokenizer.audio_end_token
        resolved_audio_token_id = tokenizer.convert_tokens_to_ids(self.audio_token)
        if resolved_audio_token_id is None:
            raise ValueError(f"Tokenizer is missing required audio placeholder token {self.audio_token!r}.")
        self.audio_token_id = int(resolved_audio_token_id)
        self.digit_token_ids = self._get_digit_token_ids()

    def _get_digit_token_ids(self) -> dict[str, int]:
        digit_token_ids = {}
        for digit in "0123456789":
            ids = self.tokenizer.encode(digit, add_special_tokens=False)
            if len(ids) != 1:
                raise ValueError(f"Digit {digit!r} is not a single token: {ids}")
            digit_token_ids[digit] = int(ids[0])
        return digit_token_ids

    def _audio_span_ids(self, audio_seq_len: int) -> list[int]:
        audio_seq_len = int(audio_seq_len)
        if not self.enable_time_marker or audio_seq_len <= 0 or self.time_marker_every_seconds <= 0:
            return [self.audio_token_id] * max(audio_seq_len, 0)

        tokens_per_marker = int(self.audio_tokens_per_second * self.time_marker_every_seconds)
        if tokens_per_marker <= 0:
            return [self.audio_token_id] * audio_seq_len

        duration = audio_seq_len / float(self.audio_tokens_per_second)
        output, consumed = [], 0
        for sec in range(self.time_marker_every_seconds, int(duration) + 1, self.time_marker_every_seconds):
            pos = (sec // self.time_marker_every_seconds) * tokens_per_marker
            segment_len = pos - consumed
            if segment_len > 0:
                output.extend([self.audio_token_id] * segment_len)
                consumed += segment_len
            marker_ids = [self.digit_token_ids[digit] for digit in str(sec)]
            output.extend(marker_ids)

        remainder = audio_seq_len - consumed
        if remainder > 0:
            output.extend([self.audio_token_id] * remainder)
        return output

    def _expand_audio_token(self, text: str, num_audio_tokens: int, max_length: int) -> list[int]:
        audio_ids = self._audio_span_ids(num_audio_tokens)
        audio_token_count = text.count(self.audio_token)
        if audio_token_count != 1:
            raise ValueError(
                f"Expected exactly one {self.audio_token!r} token per text sample, got {audio_token_count}."
            )
        before_audio, after_audio = text.split(self.audio_token, maxsplit=1)
        before_ids = self.tokenizer.encode(before_audio, add_special_tokens=False)
        after_ids = self.tokenizer.encode(after_audio, add_special_tokens=False)

        input_ids = before_ids + audio_ids + after_ids

        if len(input_ids) > max_length:
            raise ValueError(f"Prompt/audio sequence exceeds max_length={max_length}")
        return input_ids

    def __call__(
        self,
        text: Union[str, list[str]],
        audio,
        *,
        max_length: int = 131072,
        **kwargs: Unpack[MossTranscribeDiarizeProcessorKwargs],
    ) -> BatchFeature:
        return_tensors = kwargs.pop("return_tensors", "pt")
        output_kwargs = self._merge_kwargs(
            MossTranscribeDiarizeProcessorKwargs,
            tokenizer_init_kwargs=self.tokenizer.init_kwargs,
            **kwargs,
        )
        if isinstance(text, str):
            texts = [text]
        else:
            texts = list(text)
        audios = audio if isinstance(audio, list) else [audio]
        if len(texts) != len(audios):
            raise ValueError(f"Expected one audio per text prompt, got {len(audios)} audios and {len(texts)} prompts.")

        input_features, audio_feature_lengths, audio_chunk_mapping = _audios_to_input_features(
            self.feature_extractor,
            audios,
            audio_merge_size=self.audio_merge_size,
            feature_extractor_kwargs=output_kwargs["audio_kwargs"],
        )
        audio_token_counts = torch.zeros(len(audios), dtype=torch.long, device=audio_feature_lengths.device)
        audio_token_counts.scatter_add_(0, audio_chunk_mapping, audio_feature_lengths)

        encoded = [
            self._expand_audio_token(prompt, int(num_audio_tokens.item()), max_length)
            for prompt, num_audio_tokens in zip(texts, audio_token_counts)
        ]
        max_seq_len = max(len(ids) for ids in encoded)
        pad_token_id = self.tokenizer.pad_token_id
        if pad_token_id is None:
            pad_token_id = self.tokenizer.eos_token_id or 0

        input_ids, attention_mask = [], []
        for ids in encoded:
            pad_len = max_seq_len - len(ids)
            input_ids.append(ids + [pad_token_id] * pad_len)
            attention_mask.append([1] * len(ids) + [0] * pad_len)

        target_device = input_features.device
        data = {
            "input_ids": torch.tensor(input_ids, dtype=torch.long, device=target_device),
            "attention_mask": torch.tensor(attention_mask, dtype=torch.long, device=target_device),
            "input_features": input_features,
            "audio_feature_lengths": audio_feature_lengths,
            "audio_chunk_mapping": audio_chunk_mapping,
        }
        return BatchFeature(data=data, tensor_type=return_tensors)


__all__ = ["MossTranscribeDiarizeProcessor"]