scripts/moss_audio_mel_mlx.py

"""Pure-MLX mel-spectrogram + input_ids builder for MOSS-Audio.

Replaces the torch-dependent MossAudioProcessor at inference time. Matches the
upstream processor's output byte-for-byte on our test clips (see parity check
in test_moss_audio_mel_parity.py).

What's ported:
  - Log-mel spectrogram (n_mels=128, n_fft=400, hop=160, sr=16000) via mx.fft.rfft,
    reusing mlx-examples/whisper's mel_filters.npz for the filterbank.
  - Whisper-style fbank normalization: log10 → clip to (max - 8.0) → (+4)/4.
  - input_ids construction: audio-span expansion with 2-second time markers,
    chat-template wrapping (<|im_start|>system/user/assistant).

What still uses a non-torch dep:
  - Tokenization: HuggingFace AutoTokenizer (pure Python, no torch at runtime).
    Install path: `pip install transformers` gets the slow BPE tokenizer; no
    torch required at import.
"""
from __future__ import annotations

from functools import lru_cache
from pathlib import Path
from typing import Sequence

import mlx.core as mx
import numpy as np

SAMPLE_RATE = 16_000
N_FFT = 400
HOP_LENGTH = 160
MEL_DIM = 128
AUDIO_TOKENS_PER_SECOND = 12.5

AUDIO_TOKEN_ID = 151654
AUDIO_START_ID = 151669
AUDIO_END_ID = 151670

DIGIT_TOKEN_IDS = {str(i): 15 + i for i in range(10)}

_ASSETS = Path(__file__).resolve().parent / "assets"


@lru_cache(maxsize=None)
def _mel_filters() -> mx.array:
    return mx.load(str(_ASSETS / "mel_filters.npz"))[f"mel_{MEL_DIM}"]


@lru_cache(maxsize=None)
def _hann_window(size: int) -> mx.array:
    return mx.array(np.hanning(size + 1)[:-1])


def _stft_rfft(x: mx.array) -> mx.array:
    """STFT with reflect padding, returning rfft bins. Matches librosa defaults."""
    padding = N_FFT // 2
    prefix = x[1 : padding + 1][::-1]
    suffix = x[-(padding + 1) : -1][::-1]
    x = mx.concatenate([prefix, x, suffix])

    noverlap = HOP_LENGTH
    t = (x.size - N_FFT + noverlap) // noverlap
    strides = [noverlap, 1]
    x = mx.as_strided(x, shape=[t, N_FFT], strides=strides)
    return mx.fft.rfft(x * _hann_window(N_FFT))


def log_mel_spectrogram_mlx(audio: np.ndarray) -> mx.array:
    """Compute (MEL_DIM, n_frames) log-mel spectrogram in pure MLX.

    Matches upstream's WhisperFeatureExtractor._np_extract_fbank_features with
    n_mels=128. Used by the MOSS-Audio encoder, which crops/pads along the
    frame axis.
    """
    if not isinstance(audio, mx.array):
        audio = mx.array(audio.astype(np.float32))

    freqs = _stft_rfft(audio)
    magnitudes = freqs[:-1, :].abs().square()

    filters = _mel_filters()
    mel_spec = magnitudes @ filters.T
    log_spec = mx.maximum(mel_spec, 1e-10).log10()
    log_spec = mx.maximum(log_spec, log_spec.max() - 8.0)
    log_spec = (log_spec + 4.0) / 4.0
    return log_spec.T  # (MEL_DIM, n_frames)


def _conv3_downsample_len(raw_mel_len: int) -> int:
    """MOSS-Audio encoder stem: three conv/stride-2 layers → audio token count."""
    n = int(raw_mel_len)
    for _ in range(3):
        n = (n - 1) // 2 + 1
    return n


def _digit_token_ids(second: int) -> list[int]:
    return [DIGIT_TOKEN_IDS[d] for d in str(second)]


def _build_audio_placeholder_ids(num_audio_tokens: int, *, enable_time_marker: bool) -> list[int]:
    if not enable_time_marker:
        return [AUDIO_TOKEN_ID] * num_audio_tokens

    tokens_per_marker = int(AUDIO_TOKENS_PER_SECOND * 2)  # every 2 seconds
    total_seconds = num_audio_tokens / AUDIO_TOKENS_PER_SECOND
    num_full_seconds = int(total_seconds)

    out: list[int] = []
    consumed = 0
    for second in range(2, num_full_seconds + 1, 2):
        marker_pos = (second // 2) * tokens_per_marker
        segment_len = marker_pos - consumed
        if segment_len > 0:
            out.extend([AUDIO_TOKEN_ID] * segment_len)
            consumed += segment_len
        out.extend(_digit_token_ids(second))

    remaining = num_audio_tokens - consumed
    if remaining > 0:
        out.extend([AUDIO_TOKEN_ID] * remaining)
    return out


def _default_prompt(text: str) -> str:
    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"
    )


def build_input_ids(tokenizer, prompt_text: str, *, audio_token_count: int,
                    enable_time_marker: bool = True) -> list[int]:
    """Expand the <|audio_bos|><|AUDIO|><|audio_eos|> marker into audio_token_count tokens.

    Mirrors MossAudioProcessor._build_input_from_prompt but without torch tensors.
    """
    import re
    audio_span_re = re.compile(r"<\|audio_bos\|>(?:<\|AUDIO\|>)+<\|audio_eos\|>")

    if audio_span_re.search(prompt_text) is None:
        prompt_text = _default_prompt(prompt_text)

    spans = list(audio_span_re.finditer(prompt_text))
    if len(spans) != 1:
        raise ValueError(f"Expected exactly 1 audio span, got {len(spans)}")

    match = spans[0]
    prefix = prompt_text[: match.start()]
    suffix = prompt_text[match.end():]

    ids: list[int] = []
    if prefix:
        ids.extend(tokenizer.encode(prefix, add_special_tokens=False))
    ids.append(AUDIO_START_ID)
    ids.extend(_build_audio_placeholder_ids(audio_token_count, enable_time_marker=enable_time_marker))
    ids.append(AUDIO_END_ID)
    if suffix:
        ids.extend(tokenizer.encode(suffix, add_special_tokens=False))
    return ids


def build_mel_and_input_ids(
    audio: np.ndarray,
    tokenizer,
    *,
    prompt: str,
    enable_time_marker: bool = True,
) -> tuple[mx.array, mx.array, mx.array, int]:
    """Pure-MLX equivalent of MossAudioProcessor(text=prompt, audios=[audio]).

    Returns (mel, lens, input_ids, audio_token_id) — same shapes and semantics
    as the torch bridge's build_mel_spectrogram().
    """
    audio_f32 = audio.astype(np.float32)
    mel = log_mel_spectrogram_mlx(audio_f32)          # (MEL_DIM, T)
    raw_mel_len = int(mel.shape[-1])
    audio_token_count = _conv3_downsample_len(raw_mel_len)

    mel = mel[None, ...]                              # (1, MEL_DIM, T)
    lens = mx.array(np.array([raw_mel_len], dtype=np.int32))

    input_ids = build_input_ids(
        tokenizer, prompt, audio_token_count=audio_token_count,
        enable_time_marker=enable_time_marker,
    )
    input_ids_mx = mx.array(np.array([input_ids], dtype=np.int64))

    return mel, lens, input_ids_mx, AUDIO_TOKEN_ID