diff --git a/mlx_audio/stt/models/__init__.py b/mlx_audio/stt/models/__init__.py index 04a25f62a..f7986e08b 100644 --- a/mlx_audio/stt/models/__init__.py +++ b/mlx_audio/stt/models/__init__.py @@ -3,6 +3,7 @@ lasr_ctc, parakeet, qwen3_asr, + sensevoice, voxtral, voxtral_realtime, wav2vec, diff --git a/mlx_audio/stt/models/sensevoice/README.md b/mlx_audio/stt/models/sensevoice/README.md new file mode 100644 index 000000000..b97b00ea6 --- /dev/null +++ b/mlx_audio/stt/models/sensevoice/README.md @@ -0,0 +1,39 @@ +# SenseVoice + +MLX implementation of [SenseVoice Small](https://github.com/FunAudioLLM/SenseVoice) from Alibaba DAMO Academy. Non-autoregressive CTC-based model with speech recognition, language identification, emotion recognition, and audio event detection. + +## Features + +- 50+ languages (strong on zh, en, ja, ko, yue) +- emotion detection (happy, sad, angry, neutral, etc.) +- audio event detection (speech, bgm, laughter, applause) +- non-autoregressive inference (very fast, ~70ms for 10s audio) +- ~234M parameters + +## Usage + +```python +from mlx_audio.stt import load + +# downloads automatically from huggingface +model = load("mlx-community/SenseVoiceSmall") + +result = model.generate("audio.wav", language="auto") +print(result.text) +print(result.language) + +# rich info (emotion, event) is in the first segment +seg = result.segments[0] +print(seg["emotion"]) # e.g. "neutral" +print(seg["event"]) # e.g. "Speech" +``` + +### Language options + +`"auto"`, `"zh"`, `"en"`, `"ja"`, `"ko"`, `"yue"`, `"nospeech"` + +### Inverse text normalization + +```python +result = model.generate("audio.wav", language="auto", use_itn=True) +``` diff --git a/mlx_audio/stt/models/sensevoice/__init__.py b/mlx_audio/stt/models/sensevoice/__init__.py new file mode 100644 index 000000000..c28cad6d4 --- /dev/null +++ b/mlx_audio/stt/models/sensevoice/__init__.py @@ -0,0 +1,3 @@ +from .config import ModelConfig +from .sensevoice import SenseVoiceSmall +from .sensevoice import SenseVoiceSmall as Model diff --git a/mlx_audio/stt/models/sensevoice/config.py b/mlx_audio/stt/models/sensevoice/config.py new file mode 100644 index 000000000..181c5ccd0 --- /dev/null +++ b/mlx_audio/stt/models/sensevoice/config.py @@ -0,0 +1,95 @@ +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional + + +@dataclass +class EncoderConfig: + output_size: int = 512 + attention_heads: int = 4 + linear_units: int = 2048 + num_blocks: int = 50 + tp_blocks: int = 20 + dropout_rate: float = 0.1 + positional_dropout_rate: float = 0.1 + attention_dropout_rate: float = 0.1 + kernel_size: int = 11 + sanm_shift: int = 0 + normalize_before: bool = True + + @classmethod + def from_dict(cls, params: Dict[str, Any]) -> "EncoderConfig": + import inspect + + valid = { + k: v for k, v in params.items() if k in inspect.signature(cls).parameters + } + # Handle the typo in upstream config ("sanm_shfit" -> "sanm_shift") + if "sanm_shfit" in params and "sanm_shift" not in params: + valid["sanm_shift"] = params["sanm_shfit"] + return cls(**valid) + + +@dataclass +class FrontendConfig: + fs: int = 16000 + window: str = "hamming" + n_mels: int = 80 + frame_length: int = 25 + frame_shift: int = 10 + lfr_m: int = 7 + lfr_n: int = 6 + + @classmethod + def from_dict(cls, params: Dict[str, Any]) -> "FrontendConfig": + import inspect + + valid = { + k: v for k, v in params.items() if k in inspect.signature(cls).parameters + } + return cls(**valid) + + +@dataclass +class ModelConfig: + model_type: str = "sensevoice" + vocab_size: int = 25055 + input_size: int = 560 + encoder_conf: Optional[EncoderConfig] = None + frontend_conf: Optional[FrontendConfig] = None + # CMVN stats (loaded from am.mvn) + cmvn_means: Optional[List[float]] = None + cmvn_istd: Optional[List[float]] = None + # Path for loading ancillary files + model_path: Optional[str] = None + + def __post_init__(self): + if self.encoder_conf is None: + self.encoder_conf = EncoderConfig() + elif isinstance(self.encoder_conf, dict): + self.encoder_conf = EncoderConfig.from_dict(self.encoder_conf) + if self.frontend_conf is None: + self.frontend_conf = FrontendConfig() + elif isinstance(self.frontend_conf, dict): + self.frontend_conf = FrontendConfig.from_dict(self.frontend_conf) + + @classmethod + def from_dict(cls, params: Dict[str, Any]) -> "ModelConfig": + import inspect + + encoder_conf = params.pop("encoder_conf", None) + frontend_conf = params.pop("frontend_conf", None) + valid = { + k: v for k, v in params.items() if k in inspect.signature(cls).parameters + } + config = cls(**valid) + if encoder_conf is not None: + if isinstance(encoder_conf, dict): + config.encoder_conf = EncoderConfig.from_dict(encoder_conf) + else: + config.encoder_conf = encoder_conf + if frontend_conf is not None: + if isinstance(frontend_conf, dict): + config.frontend_conf = FrontendConfig.from_dict(frontend_conf) + else: + config.frontend_conf = frontend_conf + return config diff --git a/mlx_audio/stt/models/sensevoice/sensevoice.py b/mlx_audio/stt/models/sensevoice/sensevoice.py new file mode 100644 index 000000000..30bb8bc1d --- /dev/null +++ b/mlx_audio/stt/models/sensevoice/sensevoice.py @@ -0,0 +1,598 @@ +import json +import math +import re +from pathlib import Path +from typing import Dict, List, Optional, Tuple, Union + +import mlx.core as mx +import mlx.nn as nn +import numpy as np + +from mlx_audio.stt.models.base import STTOutput +from mlx_audio.stt.utils import load_audio + +from .config import ModelConfig + + +def _compute_fbank( + waveform: mx.array, + sample_rate: int = 16000, + n_mels: int = 80, + frame_length_ms: int = 25, + frame_shift_ms: int = 10, + window: str = "hamming", +) -> mx.array: + from mlx_audio.dsp import compute_fbank_kaldi + + win_len = int(sample_rate * frame_length_ms / 1000) + win_inc = int(sample_rate * frame_shift_ms / 1000) + + waveform = waveform * (1 << 15) + + return compute_fbank_kaldi( + waveform, + sample_rate=sample_rate, + win_len=win_len, + win_inc=win_inc, + num_mels=n_mels, + win_type=window, + preemphasis=0.97, + dither=0.0, + snip_edges=True, + low_freq=20.0, + high_freq=0.0, + ) + + +def _apply_lfr(feats: mx.array, lfr_m: int = 7, lfr_n: int = 6) -> mx.array: + T, D = feats.shape + T_lfr = math.ceil(T / lfr_n) + + # left-pad with copies of the first frame + left_pad = (lfr_m - 1) // 2 + if left_pad > 0: + feats = mx.concatenate([mx.tile(feats[:1], (left_pad, 1)), feats], axis=0) + T_padded = feats.shape[0] + + frames = [] + for i in range(T_lfr): + start = i * lfr_n + end = start + lfr_m + if end <= T_padded: + stacked = feats[start:end].reshape(-1) + else: + # right-pad with copies of the last frame + available = feats[start:T_padded] + pad_count = lfr_m - available.shape[0] + padded = mx.concatenate( + [available, mx.tile(feats[-1:], (pad_count, 1))], axis=0 + ) + stacked = padded.reshape(-1) + frames.append(stacked) + return mx.stack(frames, axis=0) + + +def _apply_cmvn( + feats: mx.array, + means: mx.array, + istd: mx.array, +) -> mx.array: + return (feats + means) * istd + + +def _parse_am_mvn(path: str) -> Tuple[List[float], List[float]]: + with open(path) as f: + text = f.read() + + # Extract AddShift values (negative means) + shift_match = re.search( + r".*?\s+\d+\s+\[(.*?)\]", text, re.DOTALL + ) + if not shift_match: + raise ValueError("Could not parse AddShift from am.mvn") + means = [float(x) for x in shift_match.group(1).split()] + + # Extract Rescale values (inverse std) + rescale_match = re.search( + r".*?\s+\d+\s+\[(.*?)\]", text, re.DOTALL + ) + if not rescale_match: + raise ValueError("Could not parse Rescale from am.mvn") + istd = [float(x) for x in rescale_match.group(1).split()] + + return means, istd + + +class SinusoidalPositionEncoder(nn.Module): + def __call__(self, x: mx.array) -> mx.array: + batch_size, timesteps, input_dim = x.shape + positions = mx.arange(1, timesteps + 1)[None, :] # (1, T) + + # Compute inverse timescales + half_dim = input_dim // 2 + log_timescale_increment = math.log(10000) / (half_dim - 1) + inv_timescales = mx.exp(mx.arange(half_dim) * (-log_timescale_increment)) + inv_timescales = mx.broadcast_to( + inv_timescales[None, :], (batch_size, half_dim) + ) + + scaled_time = positions[:, :, None] * inv_timescales[:, None, :] + encoding = mx.concatenate([mx.sin(scaled_time), mx.cos(scaled_time)], axis=2) + + return x + encoding + + +class PositionwiseFeedForward(nn.Module): + def __init__(self, idim: int, hidden_units: int, dropout_rate: float = 0.0): + super().__init__() + self.w_1 = nn.Linear(idim, hidden_units) + self.w_2 = nn.Linear(hidden_units, idim) + + def __call__(self, x: mx.array) -> mx.array: + return self.w_2(nn.relu(self.w_1(x))) + + +class MultiHeadedAttentionSANM(nn.Module): + def __init__( + self, + n_head: int, + in_feat: int, + n_feat: int, + dropout_rate: float = 0.0, + kernel_size: int = 11, + sanm_shift: int = 0, + ): + super().__init__() + assert n_feat % n_head == 0 + self.d_k = n_feat // n_head + self.h = n_head + self.n_feat = n_feat + + self.linear_out = nn.Linear(n_feat, n_feat) + self.linear_q_k_v = nn.Linear(in_feat, n_feat * 3) + + self.fsmn_block = nn.Conv1d( + n_feat, + n_feat, + kernel_size=kernel_size, + stride=1, + padding=0, + groups=n_feat, + bias=False, + ) + + left_padding = (kernel_size - 1) // 2 + if sanm_shift > 0: + left_padding = left_padding + sanm_shift + self.left_padding = left_padding + self.right_padding = kernel_size - 1 - left_padding + + def _forward_fsmn(self, inputs: mx.array) -> mx.array: + x = mx.pad( + inputs, + pad_width=((0, 0), (self.left_padding, self.right_padding), (0, 0)), + ) + x = self.fsmn_block(x) + x = x + inputs + return x + + def __call__(self, x: mx.array) -> mx.array: + B, T, _ = x.shape + + q_k_v = self.linear_q_k_v(x) + q, k, v = mx.split(q_k_v, 3, axis=-1) + + fsmn_memory = self._forward_fsmn(v) + + q = q.reshape(B, T, self.h, self.d_k).transpose(0, 2, 1, 3) + k = k.reshape(B, T, self.h, self.d_k).transpose(0, 2, 1, 3) + v_h = v.reshape(B, T, self.h, self.d_k).transpose(0, 2, 1, 3) + + scores = (q * (self.d_k**-0.5)) @ k.transpose(0, 1, 3, 2) + attn = mx.softmax(scores, axis=-1) + att_out = attn @ v_h + + att_out = att_out.transpose(0, 2, 1, 3).reshape(B, T, self.n_feat) + att_out = self.linear_out(att_out) + + return att_out + fsmn_memory + + +class EncoderLayerSANM(nn.Module): + def __init__( + self, + in_size: int, + size: int, + self_attn: MultiHeadedAttentionSANM, + feed_forward: PositionwiseFeedForward, + dropout_rate: float = 0.0, + normalize_before: bool = True, + ): + super().__init__() + self.self_attn = self_attn + self.feed_forward = feed_forward + self.norm1 = nn.LayerNorm(in_size) + self.norm2 = nn.LayerNorm(size) + self.in_size = in_size + self.size = size + self.normalize_before = normalize_before + + def __call__(self, x: mx.array) -> mx.array: + residual = x + if self.normalize_before: + x = self.norm1(x) + + attn_out = self.self_attn(x) + + if self.in_size == self.size: + x = residual + attn_out + else: + x = attn_out + + residual = x + if self.normalize_before: + x = self.norm2(x) + + x = residual + self.feed_forward(x) + return x + + +class SenseVoiceEncoder(nn.Module): + def __init__(self, config: ModelConfig): + super().__init__() + enc = config.encoder_conf + input_size = config.input_size + output_size = enc.output_size + + self.embed = SinusoidalPositionEncoder() + self._output_size = output_size + + self.encoders0 = [ + EncoderLayerSANM( + in_size=input_size, + size=output_size, + self_attn=MultiHeadedAttentionSANM( + n_head=enc.attention_heads, + in_feat=input_size, + n_feat=output_size, + dropout_rate=enc.attention_dropout_rate, + kernel_size=enc.kernel_size, + sanm_shift=enc.sanm_shift, + ), + feed_forward=PositionwiseFeedForward( + idim=output_size, + hidden_units=enc.linear_units, + dropout_rate=enc.dropout_rate, + ), + dropout_rate=enc.dropout_rate, + normalize_before=enc.normalize_before, + ) + ] + + self.encoders = [ + EncoderLayerSANM( + in_size=output_size, + size=output_size, + self_attn=MultiHeadedAttentionSANM( + n_head=enc.attention_heads, + in_feat=output_size, + n_feat=output_size, + dropout_rate=enc.attention_dropout_rate, + kernel_size=enc.kernel_size, + sanm_shift=enc.sanm_shift, + ), + feed_forward=PositionwiseFeedForward( + idim=output_size, + hidden_units=enc.linear_units, + dropout_rate=enc.dropout_rate, + ), + dropout_rate=enc.dropout_rate, + normalize_before=enc.normalize_before, + ) + for _ in range(enc.num_blocks - 1) + ] + + self.after_norm = nn.LayerNorm(output_size) + + self.tp_encoders = [ + EncoderLayerSANM( + in_size=output_size, + size=output_size, + self_attn=MultiHeadedAttentionSANM( + n_head=enc.attention_heads, + in_feat=output_size, + n_feat=output_size, + dropout_rate=enc.attention_dropout_rate, + kernel_size=enc.kernel_size, + sanm_shift=enc.sanm_shift, + ), + feed_forward=PositionwiseFeedForward( + idim=output_size, + hidden_units=enc.linear_units, + dropout_rate=enc.dropout_rate, + ), + dropout_rate=enc.dropout_rate, + normalize_before=enc.normalize_before, + ) + for _ in range(enc.tp_blocks) + ] + + self.tp_norm = nn.LayerNorm(output_size) + + def __call__(self, xs_pad: mx.array) -> mx.array: + xs_pad = xs_pad * (self._output_size**0.5) + xs_pad = self.embed(xs_pad) + + for layer in self.encoders0: + xs_pad = layer(xs_pad) + + for layer in self.encoders: + xs_pad = layer(xs_pad) + + xs_pad = self.after_norm(xs_pad) + + for layer in self.tp_encoders: + xs_pad = layer(xs_pad) + + xs_pad = self.tp_norm(xs_pad) + return xs_pad + + +class SenseVoiceSmall(nn.Module): + def __init__(self, config: ModelConfig): + super().__init__() + self.config = config + self.encoder = SenseVoiceEncoder(config) + + self.ctc_lo = nn.Linear(config.encoder_conf.output_size, config.vocab_size) + self.embed = nn.Embedding(16, config.input_size) + + self.lid_dict = { + "auto": 0, + "zh": 3, + "en": 4, + "yue": 7, + "ja": 11, + "ko": 12, + "nospeech": 13, + } + self.textnorm_dict = {"withitn": 14, "woitn": 15} + self.emo_dict = { + "unk": 25009, + "happy": 25001, + "sad": 25002, + "angry": 25003, + "neutral": 25004, + } + + self.blank_id = 0 + + # CMVN stats (loaded in post_load_hook) + self._cmvn_means: Optional[mx.array] = None + self._cmvn_istd: Optional[mx.array] = None + + # Tokenizer (loaded in post_load_hook) + self._tokenizer = None + self._token_list: Optional[List[str]] = None + + def _extract_features(self, audio: mx.array) -> mx.array: + fc = self.config.frontend_conf + + fbank = _compute_fbank( + audio, + sample_rate=fc.fs, + n_mels=fc.n_mels, + frame_length_ms=fc.frame_length, + frame_shift_ms=fc.frame_shift, + window=fc.window, + ) + + feats = _apply_lfr(fbank, lfr_m=fc.lfr_m, lfr_n=fc.lfr_n) + + if self._cmvn_means is not None and self._cmvn_istd is not None: + feats = _apply_cmvn(feats, self._cmvn_means, self._cmvn_istd) + + return feats + + def _build_query( + self, + batch_size: int, + language: str = "auto", + use_itn: bool = False, + ) -> Tuple[mx.array, mx.array]: + lid = self.lid_dict.get(language, 0) + language_query = self.embed(mx.array([[lid]])) + + textnorm = "withitn" if use_itn else "woitn" + textnorm_id = self.textnorm_dict[textnorm] + textnorm_query = self.embed(mx.array([[textnorm_id]])) + + event_emo_query = self.embed(mx.array([[1, 2]])) + + if batch_size > 1: + language_query = mx.broadcast_to( + language_query, (batch_size,) + language_query.shape[1:] + ) + textnorm_query = mx.broadcast_to( + textnorm_query, (batch_size,) + textnorm_query.shape[1:] + ) + event_emo_query = mx.broadcast_to( + event_emo_query, (batch_size,) + event_emo_query.shape[1:] + ) + + input_query = mx.concatenate([language_query, event_emo_query], axis=1) + return textnorm_query, input_query + + def __call__( + self, feats: mx.array, language: str = "auto", use_itn: bool = False + ) -> mx.array: + B = feats.shape[0] + textnorm_query, input_query = self._build_query(B, language, use_itn) + + speech = mx.concatenate([textnorm_query, feats], axis=1) + speech = mx.concatenate([input_query, speech], axis=1) + + encoder_out = self.encoder(speech) + logits = self.ctc_lo(encoder_out) + return nn.log_softmax(logits, axis=-1) + + def _decode_tokens(self, token_ids: List[int]) -> str: + if self._tokenizer is not None: + return self._tokenizer.decode(token_ids) + if self._token_list is not None: + pieces = [ + self._token_list[t] for t in token_ids if 0 <= t < len(self._token_list) + ] + text = "".join(pieces).replace("▁", " ").strip() + return text + return " ".join(str(t) for t in token_ids) + + def _greedy_ctc_decode(self, log_probs: mx.array) -> Tuple[List[int], str]: + pred = mx.argmax(log_probs, axis=-1) + pred_list = pred.tolist() + + deduped = [] + prev = None + for t in pred_list: + if t != prev: + deduped.append(t) + prev = t + + token_ids = [t for t in deduped if t != self.blank_id] + text = self._decode_tokens(token_ids) + return token_ids, text + + def _extract_rich_info(self, log_probs: mx.array) -> Dict[str, str]: + info = {} + + lid_pred = mx.argmax(log_probs[0]).item() + lid_map = { + 24884: "zh", + 24885: "en", + 24888: "yue", + 24892: "ja", + 24896: "ko", + 24992: "nospeech", + } + info["language"] = lid_map.get(lid_pred, "unknown") + + emo_pred = mx.argmax(log_probs[1]).item() + emo_map = { + 25001: "happy", + 25002: "sad", + 25003: "angry", + 25004: "neutral", + 25005: "fearful", + 25006: "disgusted", + 25007: "surprised", + 25008: "other", + 25009: "unk", + } + info["emotion"] = emo_map.get(emo_pred, f"token_{emo_pred}") + + event_pred = mx.argmax(log_probs[2]).item() + event_map = { + 24993: "Speech", + 24995: "BGM", + 24997: "Laughter", + 24999: "Applause", + } + info["event"] = event_map.get(event_pred, f"token_{event_pred}") + + return info + + def generate( + self, + audio: Union[str, Path, mx.array, np.ndarray], + *, + language: str = "auto", + use_itn: bool = False, + verbose: bool = False, + **kwargs, + ) -> STTOutput: + kwargs.pop("max_tokens", None) + kwargs.pop("generation_stream", None) + kwargs.pop("dtype", None) + + if isinstance(audio, (str, Path)): + audio_data = load_audio(str(audio), sr=self.config.frontend_conf.fs) + elif isinstance(audio, np.ndarray): + audio_data = mx.array(audio, dtype=mx.float32) + elif isinstance(audio, mx.array): + audio_data = audio + else: + raise TypeError(f"Unsupported audio type: {type(audio)}") + + feats = self._extract_features(audio_data) + feats = feats[None, :, :] + + log_probs = self(feats, language=language, use_itn=use_itn) + log_probs = log_probs[0] + + rich_info = self._extract_rich_info(log_probs[:4]) + token_ids, text = self._greedy_ctc_decode(log_probs[4:]) + + if verbose: + print(f"Language: {rich_info.get('language', '?')}") + print(f"Emotion: {rich_info.get('emotion', '?')}") + print(f"Event: {rich_info.get('event', '?')}") + print(f"Text: {text}") + + return STTOutput( + text=text, + language=rich_info.get("language"), + segments=[ + { + "text": text, + "language": rich_info.get("language"), + "emotion": rich_info.get("emotion"), + "event": rich_info.get("event"), + } + ], + ) + + @staticmethod + def sanitize(weights: Dict[str, mx.array]) -> Dict[str, mx.array]: + new_weights = {} + for k, v in weights.items(): + new_key = k + new_key = new_key.replace("ctc.ctc_lo.", "ctc_lo.") + + if "fsmn_block.weight" in new_key and v.ndim == 3: + v = v.transpose(0, 2, 1) + + new_weights[new_key] = v + return new_weights + + @staticmethod + def post_load_hook(model: "SenseVoiceSmall", model_path: Path) -> "SenseVoiceSmall": + model_path = Path(model_path) + + mvn_path = model_path / "am.mvn" + if mvn_path.exists(): + means, istd = _parse_am_mvn(str(mvn_path)) + model._cmvn_means = mx.array(means) + model._cmvn_istd = mx.array(istd) + + if model._cmvn_means is None and model.config.cmvn_means is not None: + model._cmvn_means = mx.array(model.config.cmvn_means) + model._cmvn_istd = mx.array(model.config.cmvn_istd) + + bpe_path = model_path / "chn_jpn_yue_eng_ko_spectok.bpe.model" + tokens_path = model_path / "tokens.json" + + if bpe_path.exists(): + try: + import sentencepiece as spm + + sp = spm.SentencePieceProcessor() + sp.Load(str(bpe_path)) + model._tokenizer = sp + except ImportError: + pass + + if model._tokenizer is None and tokens_path.exists(): + with open(tokens_path) as f: + model._token_list = json.load(f) + + return model diff --git a/mlx_audio/stt/models/sensevoice/tests/__init__.py b/mlx_audio/stt/models/sensevoice/tests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/mlx_audio/stt/tests/test_sensevoice.py b/mlx_audio/stt/tests/test_sensevoice.py new file mode 100644 index 000000000..e66da24f4 --- /dev/null +++ b/mlx_audio/stt/tests/test_sensevoice.py @@ -0,0 +1,223 @@ +import math + +import mlx.core as mx +import mlx.nn as nn +import numpy as np +import pytest + +from mlx_audio.stt.models.sensevoice.config import ( + EncoderConfig, + FrontendConfig, + ModelConfig, +) +from mlx_audio.stt.models.sensevoice.sensevoice import ( + EncoderLayerSANM, + MultiHeadedAttentionSANM, + PositionwiseFeedForward, + SenseVoiceEncoder, + SenseVoiceSmall, + SinusoidalPositionEncoder, + _apply_cmvn, + _apply_lfr, +) + + +class TestConfig: + def test_default_config(self): + config = ModelConfig() + assert config.model_type == "sensevoice" + assert config.vocab_size == 25055 + assert config.input_size == 560 + assert config.encoder_conf.output_size == 512 + assert config.encoder_conf.num_blocks == 50 + assert config.encoder_conf.tp_blocks == 20 + + def test_from_dict(self): + d = { + "model_type": "sensevoice", + "vocab_size": 25055, + "input_size": 560, + "encoder_conf": { + "output_size": 512, + "attention_heads": 4, + "linear_units": 2048, + "num_blocks": 50, + "tp_blocks": 20, + "kernel_size": 11, + "sanm_shfit": 0, # upstream typo + }, + "frontend_conf": { + "fs": 16000, + "n_mels": 80, + "lfr_m": 7, + "lfr_n": 6, + }, + } + config = ModelConfig.from_dict(d) + assert config.encoder_conf.sanm_shift == 0 + assert config.frontend_conf.lfr_m == 7 + + +class TestSinusoidalPositionEncoder: + def test_shape(self): + enc = SinusoidalPositionEncoder() + x = mx.zeros((2, 10, 512)) + out = enc(x) + assert out.shape == (2, 10, 512) + + def test_adds_positional_info(self): + enc = SinusoidalPositionEncoder() + x = mx.zeros((1, 5, 64)) + out = enc(x) + # Output should not be all zeros since position encoding is added + assert mx.any(out != 0).item() + + +class TestPositionwiseFeedForward: + def test_shape(self): + ff = PositionwiseFeedForward(512, 2048) + x = mx.zeros((2, 10, 512)) + out = ff(x) + assert out.shape == (2, 10, 512) + + +class TestMultiHeadedAttentionSANM: + def test_shape_same_dim(self): + attn = MultiHeadedAttentionSANM( + n_head=4, in_feat=512, n_feat=512, kernel_size=11 + ) + x = mx.zeros((2, 20, 512)) + out = attn(x) + assert out.shape == (2, 20, 512) + + def test_shape_projection(self): + """Test first layer where in_feat != n_feat (560 -> 512).""" + attn = MultiHeadedAttentionSANM( + n_head=4, in_feat=560, n_feat=512, kernel_size=11 + ) + x = mx.zeros((1, 15, 560)) + out = attn(x) + assert out.shape == (1, 15, 512) + + +class TestEncoderLayerSANM: + def test_shape_same_dim(self): + attn = MultiHeadedAttentionSANM( + n_head=4, in_feat=512, n_feat=512, kernel_size=11 + ) + ff = PositionwiseFeedForward(512, 2048) + layer = EncoderLayerSANM(in_size=512, size=512, self_attn=attn, feed_forward=ff) + x = mx.zeros((2, 10, 512)) + out = layer(x) + assert out.shape == (2, 10, 512) + + def test_shape_projection(self): + attn = MultiHeadedAttentionSANM( + n_head=4, in_feat=560, n_feat=512, kernel_size=11 + ) + ff = PositionwiseFeedForward(512, 2048) + layer = EncoderLayerSANM(in_size=560, size=512, self_attn=attn, feed_forward=ff) + x = mx.zeros((1, 10, 560)) + out = layer(x) + assert out.shape == (1, 10, 512) + + +class TestLFR: + def test_basic(self): + feats = mx.ones((100, 80)) + lfr = _apply_lfr(feats, lfr_m=7, lfr_n=6) + expected_t = math.ceil(100 / 6) + assert lfr.shape == (expected_t, 560) + + def test_short_input(self): + feats = mx.ones((10, 80)) + lfr = _apply_lfr(feats, lfr_m=7, lfr_n=6) + assert lfr.shape[1] == 560 + assert lfr.shape[0] == math.ceil(10 / 6) + + def test_left_padding(self): + feats = mx.arange(20).reshape(20, 1).astype(mx.float32) + lfr = _apply_lfr(feats, lfr_m=7, lfr_n=6) + # first output frame should start with 3 copies of feats[0] then feats[0:4] + first_frame = lfr[0] + assert first_frame[0].item() == 0.0 # left-padded with feats[0] + assert first_frame[1].item() == 0.0 + assert first_frame[2].item() == 0.0 + assert first_frame[3].item() == 0.0 # feats[0] + assert first_frame[4].item() == 1.0 # feats[1] + + +class TestCMVN: + def test_basic(self): + feats = mx.ones((10, 560)) + means = mx.zeros((560,)) + istd = mx.ones((560,)) + out = _apply_cmvn(feats, means, istd) + assert out.shape == (10, 560) + np.testing.assert_allclose(np.array(out), np.ones((10, 560)), atol=1e-5) + + +class TestSenseVoiceEncoder: + @pytest.fixture + def small_config(self): + return ModelConfig( + input_size=560, + encoder_conf=EncoderConfig( + output_size=64, + attention_heads=2, + linear_units=128, + num_blocks=3, + tp_blocks=2, + kernel_size=5, + ), + ) + + def test_shape(self, small_config): + encoder = SenseVoiceEncoder(small_config) + x = mx.zeros((1, 10, 560)) + mx.eval(encoder.parameters()) + out = encoder(x) + assert out.shape == (1, 10, 64) + + +class TestSenseVoiceSmall: + @pytest.fixture + def small_model(self): + config = ModelConfig( + input_size=560, + vocab_size=100, + encoder_conf=EncoderConfig( + output_size=64, + attention_heads=2, + linear_units=128, + num_blocks=3, + tp_blocks=2, + kernel_size=5, + ), + ) + model = SenseVoiceSmall(config) + mx.eval(model.parameters()) + return model + + def test_forward_shape(self, small_model): + feats = mx.zeros((1, 10, 560)) + log_probs = small_model(feats, language="en") + # 4 query tokens + 10 feature frames = 14 + assert log_probs.shape == (1, 14, 100) + + def test_sanitize(self): + weights = { + "ctc.ctc_lo.weight": mx.zeros((100, 64)), + "ctc.ctc_lo.bias": mx.zeros((100,)), + "encoder.encoders.0.self_attn.fsmn_block.weight": mx.zeros((64, 1, 11)), + "embed.weight": mx.zeros((16, 560)), + } + sanitized = SenseVoiceSmall.sanitize(weights) + assert "ctc_lo.weight" in sanitized + assert "ctc.ctc_lo.weight" not in sanitized + # Conv1d weights transposed + assert sanitized["encoder.encoders.0.self_attn.fsmn_block.weight"].shape == ( + 64, + 11, + 1, + ) diff --git a/mlx_audio/stt/utils.py b/mlx_audio/stt/utils.py index afda67f64..307c1a75f 100644 --- a/mlx_audio/stt/utils.py +++ b/mlx_audio/stt/utils.py @@ -11,6 +11,7 @@ MODEL_REMAPPING = { "glm": "glmasr", + "sensevoice": "sensevoice", "voxtral": "voxtral", "voxtral_realtime": "voxtral_realtime", "vibevoice": "vibevoice_asr",