mlx-meralion 0.1.0__tar.gz

mlx_meralion-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Yingxu He
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

mlx_meralion-0.1.0/PKG-INFO

@@ -0,0 +1,131 @@
+Metadata-Version: 2.4
+Name: mlx-meralion
+Version: 0.1.0
+Summary: MLX-native inference for MERaLiON AudioLLM on Apple Silicon
+Author: Yingxu He
+License: MIT
+Project-URL: Homepage, https://github.com/YingxuH/mlx-audiollm
+Project-URL: Repository, https://github.com/YingxuH/mlx-audiollm
+Project-URL: Issues, https://github.com/YingxuH/mlx-audiollm/issues
+Keywords: mlx,apple-silicon,speech,asr,audio,llm,meralion
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: MacOS
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Multimedia :: Sound/Audio :: Speech
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: mlx>=0.18.0
+Requires-Dist: mlx-lm>=0.20.0
+Requires-Dist: numpy>=1.24.0
+Requires-Dist: librosa>=0.10.0
+Requires-Dist: soundfile>=0.12.0
+Requires-Dist: safetensors>=0.4.0
+Requires-Dist: huggingface-hub>=0.20.0
+Requires-Dist: tokenizers>=0.15.0
+Requires-Dist: transformers>=4.46.0
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0.0; extra == "dev"
+Requires-Dist: ruff>=0.4.0; extra == "dev"
+Dynamic: license-file
+
+# mlx-meralion
+
+MLX-native inference for [MERaLiON AudioLLM](https://huggingface.co/MERaLiON) on Apple Silicon.
+
+MERaLiON is A*STAR's multimodal audio-language model for speech transcription, translation, spoken question answering, and more.
+
+## Installation
+
+```bash
+pip install mlx-meralion
+```
+
+Requires macOS on Apple Silicon (M1/M2/M3/M4) and Python 3.10+.
+
+## Quick Start
+
+### Python API
+
+```python
+from mlx_meralion import load_model, transcribe
+
+# Load model (auto-downloads from HuggingFace on first use)
+model = load_model("MERaLiON/MERaLiON-2-10B-MLX")  # 10B 8-bit, recommended
+# model = load_model("MERaLiON/MERaLiON-2-3B-MLX")   # 3B fp16, smaller
+
+# Transcribe speech
+text = transcribe(model, "audio.wav")
+print(text)
+
+# Translate to Chinese
+text = transcribe(model, "audio.wav", task="translate_zh")
+
+# Spoken question answering
+text = transcribe(model, "audio.wav", task="sqa", question="What is the speaker talking about?")
+
+# Summarize dialogue
+text = transcribe(model, "audio.wav", task="summarize")
+```
+
+### CLI
+
+```bash
+# ASR (default task)
+mlx-meralion --model MERaLiON/MERaLiON-2-10B-MLX --audio audio.wav --task asr
+
+# Translation
+mlx-meralion --model MERaLiON/MERaLiON-2-10B-MLX --audio audio.wav --task translate_zh
+
+# Custom instruction
+mlx-meralion --model MERaLiON/MERaLiON-2-10B-MLX --audio audio.wav --instruction "Summarize this in one sentence."
+```
+
+## Supported Tasks
+
+| Task | Description |
+|------|-------------|
+| `asr` | Speech-to-text transcription |
+| `translate_zh` | Translate to Chinese |
+| `translate_id` | Translate to Indonesian |
+| `translate_ms` | Translate to Malay |
+| `translate_ta` | Translate to Tamil |
+| `sqa` | Spoken question answering (requires `question=`) |
+| `summarize` | Dialogue summarization |
+| `paralinguistics` | Speaker characteristic analysis |
+
+## Available Models
+
+| Model | Size | RAM | Quality | HuggingFace |
+|-------|------|-----|---------|-------------|
+| MERaLiON-2-10B-MLX | ~10 GB | 16+ GB | Best | [MERaLiON/MERaLiON-2-10B-MLX](https://huggingface.co/MERaLiON/MERaLiON-2-10B-MLX) |
+| MERaLiON-2-3B-MLX | ~6 GB | 8+ GB | Good | [MERaLiON/MERaLiON-2-3B-MLX](https://huggingface.co/MERaLiON/MERaLiON-2-3B-MLX) |
+
+## Features
+
+- **Apple Silicon native**: Runs entirely on MLX with GPU acceleration
+- **N-gram blocking**: Automatically prevents repetitive output (matching HuggingFace quality)
+- **Smart chunking**: Long audio split at 30s boundaries; short tails merged to prevent hallucination
+- **Auto-download**: HuggingFace models are downloaded and cached automatically
+- **Multiple tasks**: ASR, translation, QA, summarization, and more
+
+## Architecture
+
+```
+Audio (WAV/MP3/FLAC)
+  -> Whisper Encoder (1280-d)
+    -> LayerNorm + MLP Adaptor
+      -> Speech embeddings merged into text sequence
+        -> Gemma2 Decoder -> text output
+```
+
+## License
+
+MIT

mlx_meralion-0.1.0/README.md

@@ -0,0 +1,93 @@
+# mlx-meralion
+
+MLX-native inference for [MERaLiON AudioLLM](https://huggingface.co/MERaLiON) on Apple Silicon.
+
+MERaLiON is A*STAR's multimodal audio-language model for speech transcription, translation, spoken question answering, and more.
+
+## Installation
+
+```bash
+pip install mlx-meralion
+```
+
+Requires macOS on Apple Silicon (M1/M2/M3/M4) and Python 3.10+.
+
+## Quick Start
+
+### Python API
+
+```python
+from mlx_meralion import load_model, transcribe
+
+# Load model (auto-downloads from HuggingFace on first use)
+model = load_model("MERaLiON/MERaLiON-2-10B-MLX")  # 10B 8-bit, recommended
+# model = load_model("MERaLiON/MERaLiON-2-3B-MLX")   # 3B fp16, smaller
+
+# Transcribe speech
+text = transcribe(model, "audio.wav")
+print(text)
+
+# Translate to Chinese
+text = transcribe(model, "audio.wav", task="translate_zh")
+
+# Spoken question answering
+text = transcribe(model, "audio.wav", task="sqa", question="What is the speaker talking about?")
+
+# Summarize dialogue
+text = transcribe(model, "audio.wav", task="summarize")
+```
+
+### CLI
+
+```bash
+# ASR (default task)
+mlx-meralion --model MERaLiON/MERaLiON-2-10B-MLX --audio audio.wav --task asr
+
+# Translation
+mlx-meralion --model MERaLiON/MERaLiON-2-10B-MLX --audio audio.wav --task translate_zh
+
+# Custom instruction
+mlx-meralion --model MERaLiON/MERaLiON-2-10B-MLX --audio audio.wav --instruction "Summarize this in one sentence."
+```
+
+## Supported Tasks
+
+| Task | Description |
+|------|-------------|
+| `asr` | Speech-to-text transcription |
+| `translate_zh` | Translate to Chinese |
+| `translate_id` | Translate to Indonesian |
+| `translate_ms` | Translate to Malay |
+| `translate_ta` | Translate to Tamil |
+| `sqa` | Spoken question answering (requires `question=`) |
+| `summarize` | Dialogue summarization |
+| `paralinguistics` | Speaker characteristic analysis |
+
+## Available Models
+
+| Model | Size | RAM | Quality | HuggingFace |
+|-------|------|-----|---------|-------------|
+| MERaLiON-2-10B-MLX | ~10 GB | 16+ GB | Best | [MERaLiON/MERaLiON-2-10B-MLX](https://huggingface.co/MERaLiON/MERaLiON-2-10B-MLX) |
+| MERaLiON-2-3B-MLX | ~6 GB | 8+ GB | Good | [MERaLiON/MERaLiON-2-3B-MLX](https://huggingface.co/MERaLiON/MERaLiON-2-3B-MLX) |
+
+## Features
+
+- **Apple Silicon native**: Runs entirely on MLX with GPU acceleration
+- **N-gram blocking**: Automatically prevents repetitive output (matching HuggingFace quality)
+- **Smart chunking**: Long audio split at 30s boundaries; short tails merged to prevent hallucination
+- **Auto-download**: HuggingFace models are downloaded and cached automatically
+- **Multiple tasks**: ASR, translation, QA, summarization, and more
+
+## Architecture
+
+```
+Audio (WAV/MP3/FLAC)
+  -> Whisper Encoder (1280-d)
+    -> LayerNorm + MLP Adaptor
+      -> Speech embeddings merged into text sequence
+        -> Gemma2 Decoder -> text output
+```
+
+## License
+
+MIT

mlx_meralion-0.1.0/mlx_meralion/__init__.py

@@ -0,0 +1,13 @@
+"""MLX-native inference for MERaLiON AudioLLM on Apple Silicon.
+
+Quick start:
+    from mlx_meralion import load_model, transcribe
+
+    model = load_model("MERaLiON/MERaLiON-2-10B-MLX")
+    text = transcribe(model, "audio.wav")
+"""
+
+__version__ = "0.1.0"
+
+from .inference import load_model as load_model, transcribe as transcribe, LoadedModel as LoadedModel
+from .processor import get_task_prompt as get_task_prompt

mlx_meralion-0.1.0/mlx_meralion/adaptor.py

@@ -0,0 +1,115 @@
+"""MLP adaptor module for MERaLiON.
+
+Bridges the Whisper encoder output (1280-d, 1500 timesteps) to the
+SEA-LION/Gemma2 decoder input (hidden_size-d, 100 timesteps).
+
+Two variants:
+    - MERaLiONSpeechAudioAdapter: v1 simple MLP (modeling_meralion.py)
+    - MERaLiONSpeechAudioAdapterLarge: v2 gated MLP (modeling_meralion2.py)
+"""
+
+import mlx.core as mx
+import mlx.nn as nn
+
+
+class MERaLiONSpeechAudioAdapter(nn.Module):
+    """Audio-to-text adapter with timestep compression (v1).
+
+    Matches MERaLiONSpeechAudioAdaper from modeling_meralion.py.
+    Architecture: reshape -> Linear+SiLU -> Linear+SiLU+Linear -> output
+    """
+
+    def __init__(
+        self,
+        speech_hidden_size: int = 1280,
+        text_hidden_size: int = 3584,
+        scale_factor: int = 15,
+    ):
+        super().__init__()
+        self.scale_factor = scale_factor
+        mlp_input_dim = speech_hidden_size * scale_factor
+
+        self.mlp_adapter = nn.Sequential(
+            nn.Linear(mlp_input_dim, speech_hidden_size),
+            nn.SiLU(),
+        )
+
+        self.speech_llm_proj = nn.Sequential(
+            nn.Linear(speech_hidden_size, speech_hidden_size * 4),
+            nn.SiLU(),
+            nn.Linear(speech_hidden_size * 4, text_hidden_size),
+        )
+
+    def __call__(self, speech_embeds: mx.array) -> mx.array:
+        B, T, D = speech_embeds.shape
+        new_T = (T // self.scale_factor) * self.scale_factor
+        speech_embeds = speech_embeds[:, :new_T, :]
+        speech_embeds = speech_embeds.reshape(
+            B, new_T // self.scale_factor, D * self.scale_factor
+        )
+        speech_embeds = self.mlp_adapter(speech_embeds)
+        speech_embeds = self.speech_llm_proj(speech_embeds)
+        return speech_embeds
+
+
+class MERaLiONSpeechAudioAdapterLarge(nn.Module):
+    """Audio-to-text adapter with gated linear unit (v2 / MERaLiON-2).
+
+    Matches MERaLiON2SpeechAudioAdaperLarge from modeling_meralion2.py.
+    Architecture: reshape -> Linear+SiLU -> GLU(gate_proj, pool_proj) -> out_proj
+    """
+
+    def __init__(
+        self,
+        speech_hidden_size: int = 1280,
+        text_hidden_size: int = 3584,
+        scale_factor: int = 15,
+    ):
+        super().__init__()
+        self.scale_factor = scale_factor
+        mlp_input_dim = speech_hidden_size * scale_factor
+        intermediate_dim = speech_hidden_size * 5
+
+        self.mlp_adapter = nn.Sequential(
+            nn.Linear(mlp_input_dim, intermediate_dim),
+            nn.SiLU(),
+        )
+
+        self.gate_proj = nn.Linear(intermediate_dim, intermediate_dim)
+        self.pool_proj = nn.Linear(intermediate_dim, intermediate_dim)
+        self.act_fn = nn.SiLU()
+
+        self.out_proj = nn.Linear(intermediate_dim, text_hidden_size)
+
+    def __call__(self, speech_embeds: mx.array) -> mx.array:
+        B, T, D = speech_embeds.shape
+        new_T = (T // self.scale_factor) * self.scale_factor
+        speech_embeds = speech_embeds[:, :new_T, :]
+        speech_embeds = speech_embeds.reshape(
+            B, new_T // self.scale_factor, D * self.scale_factor
+        )
+
+        speech_embeds = self.mlp_adapter(speech_embeds)
+        speech_embeds = self.act_fn(self.gate_proj(speech_embeds)) * self.pool_proj(speech_embeds)
+        speech_embeds = self.out_proj(speech_embeds)
+        return speech_embeds
+
+
+def create_adapter(
+    variant: str,
+    speech_hidden_size: int = 1280,
+    text_hidden_size: int = 3584,
+    scale_factor: int = 15,
+) -> nn.Module:
+    """Create the appropriate adapter variant.
+
+    Args:
+        variant: "v1" for simple MLP, "v2" or "large" for gated MLP
+    """
+    if variant in ("v2", "large", "meralion2"):
+        return MERaLiONSpeechAudioAdapterLarge(
+            speech_hidden_size, text_hidden_size, scale_factor
+        )
+    return MERaLiONSpeechAudioAdapter(
+        speech_hidden_size, text_hidden_size, scale_factor
+    )