-
-
Notifications
You must be signed in to change notification settings - Fork 722
Add MiMo-V2.5-ASR STT support #719
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Open
ailuntx
wants to merge
8
commits into
Blaizzy:main
Choose a base branch
from
ailuntx:feat/mimo-v25-asr
base: main
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
Open
Changes from 7 commits
Commits
Show all changes
8 commits
Select commit
Hold shift + click to select a range
6e7feb1
Add MiMo-V2.5-ASR STT support
ailuntz 5da8bc2
Document MiMo-V2.5-ASR in README
ailuntz 6241c57
Address MiMo review feedback
ailuntz 173badc
Merge branch 'main' into feat/mimo-v25-asr
Blaizzy 70355ea
Restore load_config import
Blaizzy f36d9b6
Format MiMo STT files
ailuntz ff0197c
Use standard MiMo weight loading path
ailuntz d9ceb4b
Merge branch 'main' into feat/mimo-v25-asr
lucasnewman File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -5,6 +5,7 @@ | |
| granite_speech, | ||
| granite_speech_nar, | ||
| lasr_ctc, | ||
| mimo_v2_asr, | ||
| parakeet, | ||
| qwen3_asr, | ||
| sensevoice, | ||
|
|
||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,42 @@ | ||
| # MiMo-V2.5-ASR | ||
|
|
||
| MLX support for Xiaomi's `MiMo-V2.5-ASR`. | ||
|
|
||
| ## Available Model | ||
|
|
||
| - [MiMo-V2.5-ASR](https://huggingface.co/mlx-community/MiMo-V2.5-ASR-MLX) | ||
|
|
||
| The model repo resolves its audio tokenizer from `mlx-community/MiMo-Audio-Tokenizer` | ||
| via `mlx_manifest.json`, so the default Hugging Face path works without extra | ||
| arguments. | ||
|
|
||
| ## Python Usage | ||
|
|
||
| ```python | ||
| from mlx_audio.stt import load | ||
|
|
||
| model = load("mlx-community/MiMo-V2.5-ASR-MLX") | ||
| result = model.generate("audio.wav", language="en") | ||
| print(result.text) | ||
| ``` | ||
|
|
||
| ## Local Usage | ||
|
|
||
| ```python | ||
| from mlx_audio.stt import load | ||
|
|
||
| model = load("/path/to/MiMo-V2.5-ASR-MLX") | ||
| result = model.generate("audio.wav") | ||
| print(result.text) | ||
| ``` | ||
|
|
||
| If you want to override the tokenizer location for a local checkout: | ||
|
|
||
| ```python | ||
| from mlx_audio.stt import load | ||
|
|
||
| model = load( | ||
| "/path/to/MiMo-V2.5-ASR-MLX", | ||
| audio_tokenizer_dir="/path/to/MiMo-Audio-Tokenizer", | ||
| ) | ||
| ``` | ||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,4 @@ | ||
| from .asr import MiMoASR, Model | ||
| from .config import MiMoAudioConfig as ModelConfig | ||
|
|
||
| __all__ = ["MiMoASR", "Model", "ModelConfig"] |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,255 @@ | ||
| import json | ||
| import time | ||
| from pathlib import Path | ||
| from typing import Dict, Optional | ||
|
|
||
| import mlx.core as mx | ||
| import mlx.nn as nn | ||
|
|
||
| from mlx_audio.stt.models.base import STTOutput | ||
| from mlx_audio.utils import get_model_path, load_weights | ||
|
|
||
| from .audio_encoder import AudioEncoder, AudioEncoderConfig | ||
| from .config import MiMoAudioConfig | ||
| from .mel import SAMPLE_RATE as AUDIO_SAMPLE_RATE | ||
| from .mel import log_mel_spectrogram | ||
| from .model import MiMoAudioMLX, MiMoSampler | ||
| from .prompt import build_asr_prompt | ||
|
|
||
|
|
||
| class Model(nn.Module): | ||
| """End-to-end MiMo-V2.5-ASR pipeline for MLX-Audio.""" | ||
|
|
||
| def __init__(self, config: MiMoAudioConfig): | ||
| super().__init__() | ||
| self.config = config | ||
| self.audio_encoder = AudioEncoder(AudioEncoderConfig()) | ||
| self.model = MiMoAudioMLX(config) | ||
|
|
||
| self._tokenizer = None | ||
| self._eos_token_id = None | ||
| self._pad_token_id = None | ||
|
|
||
| @staticmethod | ||
| def sanitize(weights: Dict[str, mx.array]) -> Dict[str, mx.array]: | ||
| return {f"model.{key}": value for key, value in weights.items()} | ||
|
|
||
| def model_quant_predicate(self, path: str, module: nn.Module) -> bool: | ||
| return path.startswith("model.") | ||
|
|
||
| @classmethod | ||
| def post_load_hook(cls, model: "Model", model_path: Path) -> "Model": | ||
| from transformers import AutoTokenizer | ||
|
|
||
| model._tokenizer = AutoTokenizer.from_pretrained( | ||
| str(model_path), | ||
| trust_remote_code=True, | ||
| ) | ||
| model._eos_token_id = model._tokenizer.eos_token_id | ||
| model._pad_token_id = model._tokenizer.pad_token_id or model._eos_token_id | ||
|
|
||
| audio_tokenizer_dir = cls._resolve_audio_tokenizer_dir(model_path) | ||
| with open(audio_tokenizer_dir / "config.json") as f: | ||
| tokenizer_config = json.load(f) | ||
|
|
||
| model.audio_encoder = AudioEncoder( | ||
| AudioEncoderConfig.from_dict(tokenizer_config) | ||
| ) | ||
| audio_encoder_weights = model._sanitize_audio_encoder_weights( | ||
| load_weights(audio_tokenizer_dir) | ||
| ) | ||
| model.audio_encoder.load_weights( | ||
| list(audio_encoder_weights.items()), | ||
| strict=False, | ||
| ) | ||
| mx.eval(model.audio_encoder.parameters()) | ||
| return model | ||
|
|
||
| @staticmethod | ||
| def _resolve_audio_tokenizer_dir(model_path: Path) -> Path: | ||
| def _has_weights(path: Path) -> bool: | ||
| return any(path.glob("*.safetensors")) | ||
|
|
||
| manifest_path = model_path / "mlx_manifest.json" | ||
| manifest = None | ||
| if manifest_path.exists(): | ||
| with open(manifest_path) as f: | ||
| manifest = json.load(f) | ||
|
|
||
| if manifest: | ||
| rel_dir = manifest.get("audio_tokenizer_dir") | ||
| if rel_dir: | ||
| candidate = (model_path / rel_dir).resolve() | ||
| if candidate.exists(): | ||
| return candidate | ||
|
|
||
| repo = manifest.get("audio_tokenizer_repo") | ||
| if repo: | ||
| repo_path = get_model_path(repo) | ||
| if not _has_weights(repo_path): | ||
| repo_path = get_model_path(repo, force_download=True) | ||
| return repo_path | ||
|
|
||
| sibling = (model_path.parent / "MiMo-Audio-Tokenizer").resolve() | ||
| if sibling.exists(): | ||
| return sibling | ||
|
|
||
| raise FileNotFoundError( | ||
| "Unable to resolve MiMo audio tokenizer directory from " | ||
| "mlx_manifest.json or a sibling MiMo-Audio-Tokenizer directory." | ||
| ) | ||
|
|
||
| @staticmethod | ||
| def _sanitize_audio_encoder_weights( | ||
| hf_weights: Dict[str, mx.array], | ||
| ) -> Dict[str, mx.array]: | ||
| out = {} | ||
|
|
||
| for key, tensor in hf_weights.items(): | ||
| new_key = key | ||
| new_tensor = tensor | ||
|
|
||
| if key.startswith("decoder."): | ||
| continue | ||
|
|
||
| if new_key.startswith("encoder."): | ||
| new_key = new_key[len("encoder.") :] | ||
|
|
||
| if new_key.endswith(".weight") and tensor.ndim == 3: | ||
| if tensor.shape[-1] <= 8 and tensor.shape[1] > tensor.shape[-1]: | ||
| new_tensor = tensor.transpose(0, 2, 1) | ||
|
|
||
| if "down_sample_layer.0" in new_key: | ||
| new_key = new_key.replace("down_sample_layer.0", "down_sample") | ||
|
|
||
| if "_codebook." in new_key: | ||
| new_key = new_key.replace("_codebook.", "codebook.") | ||
|
|
||
| if any(x in new_key for x in (".cluster_size", ".embed_avg", ".inited")): | ||
| continue | ||
|
|
||
| out[new_key] = new_tensor | ||
|
|
||
| for key in list(out.keys()): | ||
| if ".self_attn.k_proj.weight" in key: | ||
| bias_key = key.replace(".weight", ".bias") | ||
| if bias_key not in out: | ||
| out[bias_key] = mx.zeros((out[key].shape[0],)) | ||
|
|
||
| return out | ||
|
|
||
| def _normalize_language(self, language: Optional[str]) -> str: | ||
| if language is None: | ||
| return "auto" | ||
| lowered = language.lower() | ||
| if lowered in {"zh", "zh-cn", "chinese", "mandarin"}: | ||
| return "zh" | ||
| if lowered in {"en", "en-us", "english"}: | ||
| return "en" | ||
| return "auto" | ||
|
|
||
| def _clean_text(self, text: str) -> str: | ||
| return ( | ||
| text.replace("<|empty|>", "") | ||
| .replace("<|eot|>", "") | ||
| .replace("<|eostm|>", "") | ||
| .replace("<chinese>", "") | ||
| .replace("<english>", "") | ||
| .strip() | ||
| ) | ||
|
|
||
| def _encode_audio_codes(self, audio) -> mx.array: | ||
| mel = log_mel_spectrogram(audio) | ||
| codes = self.audio_encoder.encode(mel, n_q=self.config.audio_channels) | ||
| mx.eval(codes) | ||
|
|
||
| audio_codes = codes.transpose(1, 0).reshape(-1) | ||
| total_needed = self.config.group_size * self.config.audio_channels | ||
| remainder = audio_codes.shape[0] % total_needed | ||
| if remainder != 0: | ||
| pad_len = total_needed - remainder | ||
| last_frame = audio_codes[-self.config.audio_channels :] | ||
| padding = mx.tile(last_frame, (pad_len // self.config.audio_channels,)) | ||
| audio_codes = mx.concatenate([audio_codes, padding[:pad_len]]) | ||
| return audio_codes | ||
|
|
||
| def generate( | ||
| self, | ||
| audio, | ||
| *, | ||
| max_tokens: int = 256, | ||
| temperature: float = 0.0, | ||
| top_p: float = 0.95, | ||
| top_k: int = 0, | ||
| language: Optional[str] = None, | ||
| verbose: bool = False, | ||
| **kwargs, | ||
| ) -> STTOutput: | ||
| del kwargs | ||
| if self._tokenizer is None: | ||
| raise RuntimeError("Tokenizer not initialized. Call post_load_hook first.") | ||
|
|
||
| if isinstance(audio, str): | ||
| from mlx_audio.stt.utils import load_audio | ||
|
|
||
| audio = load_audio(audio, sr=AUDIO_SAMPLE_RATE) | ||
| elif not isinstance(audio, mx.array): | ||
| audio = mx.array(audio) | ||
|
|
||
| start_time = time.time() | ||
| resolved_language = self._normalize_language(language) | ||
| audio_codes = self._encode_audio_codes(audio) | ||
| prompt = build_asr_prompt( | ||
| audio_codes=audio_codes, | ||
| text_token_ids=[], | ||
| config=self.config, | ||
| tokenizer=self._tokenizer, | ||
| language=resolved_language, | ||
| ) | ||
|
|
||
| global_sampler = MiMoSampler( | ||
| do_sample=temperature > 0, | ||
| temperature=max(temperature, 1e-5) if temperature > 0 else 1.0, | ||
| top_k=top_k, | ||
| top_p=top_p, | ||
| ) | ||
| local_sampler = MiMoSampler(do_sample=False) | ||
|
|
||
| generated = self.model.generate( | ||
| prompt[None], | ||
| max_new_tokens=max_tokens, | ||
| global_sampler=global_sampler, | ||
| local_sampler=local_sampler, | ||
| stop_tokens=[self._eos_token_id, self.config.eot_idx], | ||
| )[0] | ||
|
|
||
| prompt_len = prompt.shape[1] | ||
| text_tokens = generated[0, prompt_len :: self.config.group_size] | ||
| token_list = text_tokens.tolist() | ||
| while token_list and token_list[-1] in { | ||
| self._eos_token_id, | ||
| self.config.eot_idx, | ||
| self.config.eostm_idx, | ||
| }: | ||
| token_list.pop() | ||
|
|
||
| text = self._clean_text( | ||
| self._tokenizer.decode(token_list, skip_special_tokens=False) | ||
| ) | ||
| total_time = time.time() - start_time | ||
| audio_duration = float(audio.shape[0]) / float(AUDIO_SAMPLE_RATE) | ||
|
|
||
| return STTOutput( | ||
| text=text, | ||
| segments=[{"text": text, "start": 0.0, "end": audio_duration}], | ||
| language=resolved_language if resolved_language != "auto" else None, | ||
| prompt_tokens=int(prompt.shape[1]), | ||
| generation_tokens=len(token_list), | ||
| total_tokens=int(prompt.shape[1]) + len(token_list), | ||
| total_time=total_time, | ||
| prompt_tps=(float(prompt.shape[1]) / total_time) if total_time > 0 else 0.0, | ||
| generation_tps=(len(token_list) / total_time) if total_time > 0 else 0.0, | ||
| ) | ||
|
|
||
|
|
||
| MiMoASR = Model |
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The
mlx_manifest.jsoncomment here isn't accurate, can we remove that?