diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 1cac60aa1..3c6b932d8 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -48,6 +48,24 @@ jobs: - name: Run core tests run: pytest -s mlx_audio/tests/ + # Differential tests: vendored mlx_audio.lm vs upstream mlx-lm. + parity: + runs-on: macos-14 + needs: style + steps: + - uses: actions/checkout@v6 + - uses: actions/setup-python@v6 + with: + python-version: '3.10' + + - name: Install with parity extra + run: | + python -m pip install --upgrade pip + pip install -e ".[all,dev,parity]" + + - name: Run parity tests + run: pytest -s tests/vendor_parity/ + # Modular installation tests - validates issue #287 # Only verifies imports work with minimal deps installed. # Full tests run separately with all deps (test files import models that need extra deps). diff --git a/docs/contributing/adding-a-model.md b/docs/contributing/adding-a-model.md index 9dcafccb7..22c3e7c0b 100644 --- a/docs/contributing/adding-a-model.md +++ b/docs/contributing/adding-a-model.md @@ -46,6 +46,26 @@ TTS models use the base classes from `mlx_audio/tts/models/base.py`: - **`GenerationResult`** -- Dataclass returned by `generate()`. Contains `audio`, `sample_rate`, `token_count`, timing information, and streaming flags. - **`BatchGenerationResult`** -- Dataclass for batch generation results. +### Transformer Components + +Many speech models wrap a language-model backbone. Import that machinery from +`mlx_audio.lm`, **not** from `mlx-lm` -- mlx-audio vendors it so the package has +no runtime dependency on mlx-lm, and a test fails the build if `mlx_lm` is +imported anywhere outside `sts/voice_pipeline.py`. + +| Need | Import from | +|---|---| +| KV caches (`KVCache`, `RotatingKVCache`, `BatchKVCache`, `make_prompt_cache`) | `mlx_audio.lm.models.cache` | +| Attention masks, `scaled_dot_product_attention` | `mlx_audio.lm.models.base` | +| Samplers, logits processors | `mlx_audio.lm.sample_utils` | +| `generate_step`, `stream_generate` | `mlx_audio.lm.generate` | +| Backbones (llama, qwen2, qwen3, gpt2, granite, lfm2, bailing_moe, gemma3) | `mlx_audio.lm.models.` | +| Quantize / save helpers | `mlx_audio.lm.convert` | + +If your model needs a backbone that is not vendored yet, copy it from mlx-lm +into `mlx_audio/lm/models/` verbatim and add the provenance header used by the +other files there (upstream path, version, commit). + ### Model Configuration Create a dataclass for your model's config that extends `BaseModelArgs`: diff --git a/mlx_audio/codec/models/mimi/modules/__init__.py b/mlx_audio/codec/models/mimi/modules/__init__.py index 63cc74c8e..1c9ed3441 100644 --- a/mlx_audio/codec/models/mimi/modules/__init__.py +++ b/mlx_audio/codec/models/mimi/modules/__init__.py @@ -4,7 +4,7 @@ # flake8: noqa """Modules used for building the models.""" -from mlx_lm.models.cache import KVCache, RotatingKVCache +from mlx_audio.lm.models.cache import KVCache, RotatingKVCache from .conv import ( Conv1d, diff --git a/mlx_audio/codec/models/mimi/modules/transformer.py b/mlx_audio/codec/models/mimi/modules/transformer.py index ba96f4485..3e4a853d7 100644 --- a/mlx_audio/codec/models/mimi/modules/transformer.py +++ b/mlx_audio/codec/models/mimi/modules/transformer.py @@ -8,7 +8,8 @@ import mlx.core as mx import mlx.nn as nn -from mlx_lm.models.cache import KVCache, RotatingKVCache + +from mlx_audio.lm.models.cache import KVCache, RotatingKVCache @dataclass diff --git a/mlx_audio/convert.py b/mlx_audio/convert.py index 293471be2..02df88af3 100644 --- a/mlx_audio/convert.py +++ b/mlx_audio/convert.py @@ -486,7 +486,7 @@ def base_requirements(path: str, module) -> bool: if not quant_predicate_name: return base_requirements - from mlx_lm.convert import mixed_quant_predicate_builder + from mlx_audio.lm.convert import mixed_quant_predicate_builder mixed_predicate = mixed_quant_predicate_builder(quant_predicate_name, model) return lambda p, m: base_requirements(p, m) and mixed_predicate(p, m) @@ -590,7 +590,12 @@ def convert( q_mode: Quantization mode (affine, mxfp4, nvfp4, mxfp8). model_domain: Force model domain ("tts", "stt", or "sts"). Auto-detected if None. """ - from mlx_lm.utils import dequantize_model, quantize_model, save_config, save_model + from mlx_audio.lm.convert import ( + dequantize_model, + quantize_model, + save_config, + save_model, + ) if quantize and dequantize: raise ValueError("Choose either quantize or dequantize, not both.") diff --git a/mlx_audio/lm/__init__.py b/mlx_audio/lm/__init__.py new file mode 100644 index 000000000..35f79ffea --- /dev/null +++ b/mlx_audio/lm/__init__.py @@ -0,0 +1 @@ +"""Transformer machinery vendored from mlx-lm. Import submodules directly.""" diff --git a/mlx_audio/lm/convert.py b/mlx_audio/lm/convert.py new file mode 100644 index 000000000..7915a61cb --- /dev/null +++ b/mlx_audio/lm/convert.py @@ -0,0 +1,184 @@ +# Copyright © 2023-2024 Apple Inc. +# Derived from mlx-lm v0.31.3 (ed1fca4cef15a824c5f1702c80f70b4cffc8e4dd). + +import copy +import json +from pathlib import Path +from typing import Callable, Optional, Union + +import mlx.core as mx +import mlx.nn as nn +from mlx.utils import tree_flatten, tree_map, tree_unflatten + +MAX_FILE_SIZE_GB = 5 + + +def mixed_quant_predicate_builder(recipe: str, model: nn.Module, group_size: int = 64): + recipes = { + "mixed_2_6": (2, 6), + "mixed_3_4": (3, 4), + "mixed_3_6": (3, 6), + "mixed_4_6": (4, 6), + } + if recipe not in recipes: + raise ValueError(f"Invalid quant recipe {recipe}") + low_bits, high_bits = recipes[recipe] + down_keys = [name for name, _ in model.named_modules() if "down_proj" in name] + if not down_keys: + raise ValueError("Model does not have expected keys for mixed quant.") + layer_location = next( + index for index, key in enumerate(down_keys[0].split(".")) if key.isdigit() + ) + num_layers = len(model.layers) + + def predicate(path: str, module: nn.Module) -> Union[bool, dict]: + del module + index = ( + int(path.split(".")[layer_location]) + if len(path.split(".")) > layer_location + else 0 + ) + high_precision = ( + index < num_layers // 8 + or index >= 7 * num_layers // 8 + or (index - num_layers // 8) % 3 == 2 + ) + wide = ( + "v_proj" in path + or "v_a_proj" in path + or "v_b_proj" in path + or "down_proj" in path + ) + # lm_head takes high bits regardless of depth, as upstream does. + bits = high_bits if (wide and high_precision) or "lm_head" in path else low_bits + return {"group_size": group_size, "bits": bits, "mode": "affine"} + + return predicate + + +def quantize_model( + model: nn.Module, + config: dict, + group_size: Optional[int], + bits: Optional[int], + mode: str = "affine", + quant_predicate: Optional[Callable] = None, +): + defaults = {"affine": (64, 4), "mxfp4": (32, 4), "nvfp4": (16, 4), "mxfp8": (32, 8)} + group_size, bits = group_size or defaults[mode][0], bits or defaults[mode][1] + config = copy.deepcopy(config) + quant_predicate = quant_predicate or getattr(model, "quant_predicate", None) + params = {"group_size": group_size, "bits": bits, "mode": mode} + + # An existing "quantization" key means the model is already partially + # quantized, so record parameters per layer rather than globally. + fine_grained = "quantization" in config + if not fine_grained: + config["quantization"] = params + + def predicate(path, module): + if not hasattr(module, "to_quantized") or module.weight.shape[-1] % group_size: + return False + result = quant_predicate(path, module) if quant_predicate else True + if isinstance(result, dict): + config["quantization"][path] = result + elif fine_grained and result: + config["quantization"][path] = params + return result + + nn.quantize(model, group_size, bits, mode=mode, class_predicate=predicate) + config["quantization_config"] = config["quantization"] + return model, config + + +def dequantize_model(model: nn.Module) -> nn.Module: + replacements = [] + for name, module in model.named_modules(): + if isinstance(module, nn.QuantizedLinear): + layer = nn.Linear(*module.weight.shape[::-1], bias="bias" in module) + elif isinstance(module, nn.QuantizedEmbedding): + layer = nn.Embedding(*module.weight.shape) + else: + continue + layer.weight = mx.dequantize( + module.weight, + module.scales, + module.biases, + module.group_size, + module.bits, + module.mode, + ) + if "bias" in module: + layer.bias = module.bias + replacements.append((name, layer)) + if replacements: + model.update_modules(tree_unflatten(replacements)) + return model + + +def save_config(config: dict, config_path: Union[str, Path]) -> None: + config = copy.deepcopy(config) + config.pop("_name_or_path", None) + config.pop("vision_config", None) + if "quantization" in config: + config["quantization_config"] = config["quantization"] + with open(config_path, "w") as handle: + json.dump(dict(sorted(config.items())), handle, indent=4) + + +def make_shards(weights: dict, max_file_size_gb: int = MAX_FILE_SIZE_GB) -> list: + max_file_size_bytes = max_file_size_gb << 30 + shards = [] + shard, shard_size = {}, 0 + for name, weight in weights.items(): + if shard_size + weight.nbytes > max_file_size_bytes: + shards.append(shard) + shard, shard_size = {}, 0 + shard[name] = weight + shard_size += weight.nbytes + shards.append(shard) + return shards + + +def save_model( + save_path: Union[str, Path], model: nn.Module, *, donate_model: bool = False +) -> None: + save_path = Path(save_path) + save_path.mkdir(parents=True, exist_ok=True) + + weights = dict(tree_flatten(model.parameters())) + total_size = sum(value.nbytes for value in weights.values()) + shards = make_shards(weights) + name_format = ( + "model-{:05d}-of-{:05d}.safetensors" if len(shards) > 1 else "model.safetensors" + ) + weight_map = { + name: name_format.format(index + 1, len(shards)) + for index, shard in enumerate(shards) + for name in shard + } + + # Release the model's references before serializing so each shard can be + # freed as it is written, rather than holding every weight twice. + if donate_model: + model.update(tree_map(lambda _: mx.array([]), model.parameters())) + weights.clear() + + for index, shard in enumerate(shards): + shards[index] = None + mx.save_safetensors( + str(save_path / name_format.format(index + 1, len(shards))), + shard, + metadata={"format": "mlx"}, + ) + del shard + + with open(save_path / "model.safetensors.index.json", "w") as handle: + json.dump( + { + "metadata": {"total_size": total_size}, + "weight_map": {k: weight_map[k] for k in sorted(weight_map)}, + }, + handle, + indent=4, + ) diff --git a/mlx_audio/lm/generate.py b/mlx_audio/lm/generate.py new file mode 100644 index 000000000..5161f6bef --- /dev/null +++ b/mlx_audio/lm/generate.py @@ -0,0 +1,277 @@ +# Copyright © 2023-2024 Apple Inc. +# Derived from mlx-lm v0.31.3 (ed1fca4cef15a824c5f1702c80f70b4cffc8e4dd), +# mlx_lm/generate.py. Trimmed to mlx-audio's single-stream generation APIs. + +import contextlib +import inspect +import time +import warnings +from dataclasses import dataclass +from typing import Any, Callable, Generator, List, Optional, Tuple, Union + +import mlx.core as mx +import mlx.nn as nn +from mlx.utils import tree_reduce + +from .models.cache import make_prompt_cache + +generation_stream = mx.new_thread_local_stream(mx.default_device()) + + +@contextlib.contextmanager +def wired_limit(model: nn.Module, streams: Optional[List[mx.Stream]] = None): + if not mx.metal.is_available(): + yield + return + + model_bytes = tree_reduce( + lambda total, value: ( + total + value.nbytes if isinstance(value, mx.array) else total + ), + model, + 0, + ) + max_rec_size = mx.device_info()["max_recommended_working_set_size"] + if model_bytes > 0.9 * max_rec_size: + warnings.warn( + "Generating with a model that requires " + f"{model_bytes / 1 << 30:.2f} GB, close to the maximum recommended " + f"size of {max_rec_size / 1 << 30:.2f} GB. This can be slow; see the " + "MLX documentation on tuning the wired limit.", + stacklevel=2, + ) + old_limit = mx.set_wired_limit(max_rec_size) + try: + yield + finally: + if streams is None: + mx.synchronize() + else: + for stream in streams: + mx.synchronize(stream) + mx.set_wired_limit(old_limit) + + +@dataclass +class GenerationResponse: + text: str + token: int + logprobs: mx.array + from_draft: bool + prompt_tokens: int + prompt_tps: float + generation_tokens: int + generation_tps: float + peak_memory: float + finish_reason: Optional[str] = None + + +class _NaiveDetokenizer: + def __init__(self, tokenizer): + self.tokenizer = tokenizer + self.reset() + + def reset(self): + self.tokens = [] + self.text = "" + self.offset = 0 + + def add_token(self, token): + self.tokens.append(token) + self.text = self.tokenizer.decode(self.tokens) + + def finalize(self): + self.text = self.tokenizer.decode(self.tokens) + + @property + def last_segment(self): + segment = self.text[self.offset :] + self.offset = len(self.text) + return segment + + +def _supports_input_embeddings(model: nn.Module) -> bool: + try: + return "input_embeddings" in inspect.signature(model.__call__).parameters + except (TypeError, ValueError): + return False + + +def _eos_ids(tokenizer) -> set[int]: + if hasattr(tokenizer, "eos_token_ids"): + return set(tokenizer.eos_token_ids) + eos_token_id = getattr(tokenizer, "eos_token_id", None) + return set() if eos_token_id is None else {eos_token_id} + + +def _encode(tokenizer, prompt: Union[str, mx.array, List[int]]) -> mx.array: + if isinstance(prompt, mx.array): + return prompt + if isinstance(prompt, str): + bos_token = getattr(tokenizer, "bos_token", None) + prompt = tokenizer.encode( + prompt, + add_special_tokens=bos_token is None or not prompt.startswith(bos_token), + ) + return mx.array(prompt) + + +def generate_step( + prompt: mx.array, + model: nn.Module, + *, + max_tokens: int = 256, + sampler: Optional[Callable[[mx.array], mx.array]] = None, + logits_processors: Optional[List[Callable[[mx.array, mx.array], mx.array]]] = None, + max_kv_size: Optional[int] = None, + prompt_cache: Optional[Any] = None, + prefill_step_size: int = 2048, + prompt_progress_callback: Optional[Callable[[int, int], None]] = None, + input_embeddings: Optional[mx.array] = None, +) -> Generator[Tuple[mx.array, mx.array], None, None]: + if input_embeddings is not None: + if not _supports_input_embeddings(model): + raise ValueError("Model does not support input embeddings.") + if len(prompt) and len(prompt) != len(input_embeddings): + raise ValueError("prompt and input_embeddings must have the same length.") + elif not len(prompt): + raise ValueError("Either input_embeddings or prompt must be provided.") + + if prompt_cache is None: + prompt_cache = make_prompt_cache(model, max_kv_size=max_kv_size) + prompt_progress_callback = prompt_progress_callback or (lambda *_: None) + sampler = sampler or (lambda logits: mx.argmax(logits, axis=-1)) + tokens = None + + def model_call(input_tokens, embeddings=None): + if embeddings is None: + return model(input_tokens, cache=prompt_cache) + return model(input_tokens, cache=prompt_cache, input_embeddings=embeddings) + + def step(input_tokens, embeddings=None): + nonlocal tokens + with mx.stream(generation_stream): + logits = model_call( + input_tokens[None], + embeddings[None] if embeddings is not None else None, + )[:, -1, :] + if logits_processors and len(input_tokens): + tokens = ( + mx.concatenate([tokens, input_tokens]) + if tokens is not None + else input_tokens + ) + for processor in logits_processors: + logits = processor(tokens, logits) + logprobs = logits - mx.logsumexp(logits, keepdims=True) + return sampler(logprobs), logprobs.squeeze(0) + + with mx.stream(generation_stream): + total = len(input_embeddings) if input_embeddings is not None else len(prompt) + processed = 0 + prompt_progress_callback(processed, total) + while total - processed > 1: + count = min(prefill_step_size, total - processed - 1) + model_call( + prompt[:count][None], + ( + input_embeddings[:count][None] + if input_embeddings is not None + else None + ), + ) + mx.eval([cache.state for cache in prompt_cache]) + processed += count + prompt_progress_callback(processed, total) + prompt = prompt[count:] + if input_embeddings is not None: + input_embeddings = input_embeddings[count:] + mx.clear_cache() + token, logprobs = step(prompt, input_embeddings) + + mx.async_eval(token, logprobs) + count = 0 + while count != max_tokens: + next_token, next_logprobs = step(token) + mx.async_eval(next_token, next_logprobs) + if count == 0: + mx.eval(token) + prompt_progress_callback(total, total) + yield token.item(), logprobs + if count % 256 == 0: + mx.clear_cache() + token, logprobs = next_token, next_logprobs + count += 1 + + +def stream_generate( + model: nn.Module, + tokenizer, + prompt: Union[str, mx.array, List[int]], + max_tokens: int = 256, + **kwargs, +) -> Generator[GenerationResponse, None, None]: + if kwargs.pop("draft_model", None) is not None: + raise ValueError("Speculative decoding is not implemented in mlx-audio.lm.") + + prompt = _encode(tokenizer, prompt) + detokenizer = _NaiveDetokenizer(tokenizer) + eos_ids = _eos_ids(tokenizer) + last_token = last_logprobs = None + prompt_tps = 0.0 + index = -1 + finish_reason = "length" + + with wired_limit(model, [generation_stream]): + started = time.perf_counter() + for index, (token, logprobs) in enumerate( + generate_step(prompt, model, max_tokens=max_tokens, **kwargs) + ): + if index == 0: + prompt_tps = prompt.size / (time.perf_counter() - started) + started = time.perf_counter() + last_token, last_logprobs = token, logprobs + if token in eos_ids: + finish_reason = "stop" + break + detokenizer.add_token(token) + if index + 1 == max_tokens: + break + yield GenerationResponse( + text=detokenizer.last_segment, + token=token, + logprobs=logprobs, + from_draft=False, + prompt_tokens=prompt.size, + prompt_tps=prompt_tps, + generation_tokens=index + 1, + generation_tps=(index + 1) / (time.perf_counter() - started), + peak_memory=mx.get_peak_memory() / 1e9, + ) + + detokenizer.finalize() + yield GenerationResponse( + text=detokenizer.last_segment, + token=last_token, + logprobs=last_logprobs, + from_draft=False, + prompt_tokens=prompt.size, + prompt_tps=prompt_tps, + generation_tokens=index + 1, + generation_tps=(index + 1) / (time.perf_counter() - started), + peak_memory=mx.get_peak_memory() / 1e9, + finish_reason=finish_reason, + ) + + +def generate( + model: nn.Module, tokenizer, prompt: Union[str, List[int]], verbose=False, **kwargs +) -> str: + text = "" + for response in stream_generate(model, tokenizer, prompt, **kwargs): + if verbose: + print(response.text, end="", flush=True) + text += response.text + if verbose: + print() + return text diff --git a/mlx_audio/lm/load.py b/mlx_audio/lm/load.py new file mode 100644 index 000000000..889fc1820 --- /dev/null +++ b/mlx_audio/lm/load.py @@ -0,0 +1,32 @@ +from pathlib import Path + +from mlx_audio.utils import ( + apply_quantization, + get_model_path, + load_config, + load_weights, +) + +from .models import gemma3 + +_MODELS = {"gemma3": gemma3} + + +def load_lm(model_id: str): + model_path = get_model_path(model_id) + config = load_config(model_path) + model_type = config.get("model_type") + module = _MODELS.get(model_type) + if module is None: + supported = ", ".join(sorted(_MODELS)) + raise ValueError( + f"Unsupported embedded language model {model_type!r}; supported: {supported}" + ) + model = module.Model(module.ModelArgs.from_dict(config)) + weights = load_weights(Path(model_path)) + apply_quantization(model, config, weights) + model.load_weights(list(model.sanitize(weights).items())) + + from transformers import AutoTokenizer + + return model, AutoTokenizer.from_pretrained(model_path) diff --git a/mlx_audio/lm/models/__init__.py b/mlx_audio/lm/models/__init__.py new file mode 100644 index 000000000..2d5529dfe --- /dev/null +++ b/mlx_audio/lm/models/__init__.py @@ -0,0 +1 @@ +"""Transformer components and backbones vendored from mlx-lm.""" diff --git a/mlx_audio/lm/models/activations.py b/mlx_audio/lm/models/activations.py new file mode 100644 index 000000000..9803e935e --- /dev/null +++ b/mlx_audio/lm/models/activations.py @@ -0,0 +1,9 @@ +from functools import partial + +import mlx.core as mx +import mlx.nn as nn + + +@partial(mx.compile, shapeless=True) +def swiglu(gate, x): + return nn.silu(gate) * x diff --git a/mlx_audio/lm/models/bailing_moe.py b/mlx_audio/lm/models/bailing_moe.py new file mode 100644 index 000000000..e29e9c19d --- /dev/null +++ b/mlx_audio/lm/models/bailing_moe.py @@ -0,0 +1,403 @@ +# Copyright © 2025 Apple Inc. +# Vendored verbatim from mlx-lm v0.31.3 (ed1fca4cef15a824c5f1702c80f70b4cffc8e4dd), +# mlx_lm/models/bailing_moe.py. MIT licensed. + +from dataclasses import dataclass +from functools import partial +from typing import Any, Dict, Optional, Union + +import mlx.core as mx +import mlx.nn as nn + +from .activations import swiglu +from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention +from .rope_utils import initialize_rope +from .switch_layers import SwitchGLU + + +@dataclass +class ModelArgs(BaseModelArgs): + model_type: str + hidden_size: int + intermediate_size: int + max_position_embeddings: int + moe_intermediate_size: int + num_experts: int + num_shared_experts: int + norm_topk_prob: bool + num_attention_heads: int + num_experts_per_tok: int + num_hidden_layers: int + num_key_value_heads: int + rms_norm_eps: float + rope_theta: float + vocab_size: int + first_k_dense_replace: int + rope_scaling: Optional[Dict[str, Union[float, str]]] = None + use_bias: bool = False + use_qkv_bias: bool = False + norm_head: bool = False + norm_softmax: bool = False + use_qk_norm: bool = False + tie_word_embeddings: bool = False + partial_rotary_factor: float = 1.0 + rotary_dim: Optional[int] = None + moe_router_enable_expert_bias: bool = False + moe_router_enable_routed_scaling: bool = True + routed_scaling_factor: float = 1.0 + score_function: str = "softmax" + n_group: int = 1 + topk_group: int = 4 + moe_shared_expert_intermediate_size: Optional[int] = None + moe_router_enable_shared_expert: bool = True + + +@partial(mx.compile, shapeless=True) +def aggregate_expert_outputs(expert_outputs, scores): + return ( + (expert_outputs * scores[..., None]).sum(axis=-2).astype(expert_outputs.dtype) + ) + + +class BailingMoeMLP(nn.Module): + def __init__(self, args: ModelArgs, intermediate_size: Optional[int] = None): + super().__init__() + self.intermediate_size = ( + intermediate_size + if intermediate_size is not None + else args.intermediate_size + ) + + self.gate_proj = nn.Linear( + args.hidden_size, self.intermediate_size, bias=args.use_bias + ) + self.down_proj = nn.Linear( + self.intermediate_size, args.hidden_size, bias=args.use_bias + ) + self.up_proj = nn.Linear( + args.hidden_size, self.intermediate_size, bias=args.use_bias + ) + + def __call__(self, x) -> mx.array: + return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x))) + + +class BailingMoeAttention(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.use_qk_norm = args.use_qk_norm + self.num_attention_heads = args.num_attention_heads + self.num_key_value_heads = args.num_key_value_heads + self.head_dim = args.hidden_size // self.num_attention_heads + self.scale = self.head_dim**-0.5 + + self.query_key_value = nn.Linear( + args.hidden_size, + (self.num_attention_heads + 2 * self.num_key_value_heads) * self.head_dim, + bias=args.use_qkv_bias, + ) + self.dense = nn.Linear( + self.num_attention_heads * self.head_dim, + args.hidden_size, + bias=args.use_bias, + ) + + if args.use_qk_norm: + self.key_layernorm = nn.RMSNorm(self.head_dim, eps=args.rms_norm_eps) + self.query_layernorm = nn.RMSNorm(self.head_dim, eps=args.rms_norm_eps) + + if (rope_dim := args.rotary_dim) is None: + rope_dim = int(self.head_dim * args.partial_rotary_factor) + self.rope = initialize_rope( + rope_dim, + args.rope_theta, + traditional=False, + scaling_config=args.rope_scaling, + max_position_embeddings=args.max_position_embeddings, + ) + + def __call__( + self, + x: mx.array, + mask: Optional[mx.array] = None, + cache: Optional[Any] = None, + ) -> mx.array: + B, L, D = x.shape + + qkv = self.query_key_value(x) + + q_size = self.num_attention_heads * self.head_dim + kv_size = self.num_key_value_heads * self.head_dim + q, k, v = mx.split(qkv, [q_size, q_size + kv_size], axis=-1) + + queries = q.reshape(B, L, self.num_attention_heads, self.head_dim).transpose( + 0, 2, 1, 3 + ) + keys = k.reshape(B, L, self.num_key_value_heads, self.head_dim).transpose( + 0, 2, 1, 3 + ) + values = v.reshape(B, L, self.num_key_value_heads, self.head_dim).transpose( + 0, 2, 1, 3 + ) + + if self.use_qk_norm: + queries = self.query_layernorm(queries) + keys = self.key_layernorm(keys) + + if cache is not None: + queries = self.rope(queries, offset=cache.offset) + keys = self.rope(keys, offset=cache.offset) + keys, values = cache.update_and_fetch(keys, values) + else: + queries = self.rope(queries) + keys = self.rope(keys) + + output = scaled_dot_product_attention( + queries, keys, values, cache=cache, scale=self.scale, mask=mask + ) + + output = output.transpose(0, 2, 1, 3).reshape(B, L, -1) + return self.dense(output) + + +@mx.compile +def group_expert_select( + gates, + e_score_correction_bias, + top_k, + n_group, + topk_group, + routed_scaling_factor, + norm_topk_prob, + score_function, +): + + in_type = gates.dtype + if score_function == "sigmoid": + scores = mx.sigmoid(gates.astype(mx.float32)) + else: + scores = mx.softmax(gates.astype(mx.float32), axis=-1) + orig_scores = scores + if e_score_correction_bias is not None: + scores = scores + e_score_correction_bias + if n_group > 1: + scores = mx.unflatten(scores, axis=-1, shape=(n_group, -1)) + group_scores = mx.topk(scores, 2, axis=-1).sum(axis=-1, keepdims=True) + k = n_group - topk_group + group_idx = mx.argpartition(group_scores, kth=k - 1, axis=-2)[..., :k, :] + scores = mx.put_along_axis( + scores, mx.stop_gradient(group_idx), mx.array(0.0, scores.dtype), axis=-2 + ) + scores = mx.flatten(scores, -2, -1) + + k = top_k + inds = mx.argpartition(scores, kth=-k, axis=-1)[..., -k:] + scores = mx.take_along_axis(orig_scores, inds, axis=-1) + if top_k > 1 and norm_topk_prob: + denominator = scores.sum(axis=-1, keepdims=True) + 1e-20 + scores = scores / denominator + scores = scores * routed_scaling_factor + + return inds, scores.astype(in_type) + + +class BailingMoeGate(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.norm_topk_prob = args.norm_topk_prob + + self.top_k = args.num_experts_per_tok + self.n_group = args.n_group + self.topk_group = args.topk_group + self.routed_scaling_factor = args.routed_scaling_factor + self.enable_routed_scaling = args.moe_router_enable_routed_scaling + + self.gate_proj = nn.Linear(args.hidden_size, args.num_experts, bias=False) + self.expert_bias = ( + mx.zeros((args.num_experts,)) + if args.moe_router_enable_expert_bias + else None + ) + self.score_function = args.score_function + + def __call__(self, x): + return group_expert_select( + self.gate_proj(x), + self.expert_bias, + self.top_k, + self.n_group, + self.topk_group, + self.routed_scaling_factor, + self.norm_topk_prob, + self.score_function, + ) + + +class BailingMoeSparseMoeBlock(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.args = args + self.num_experts_per_tok = args.num_experts_per_tok + self.switch_mlp = SwitchGLU( + args.hidden_size, + args.moe_intermediate_size, + args.num_experts, + bias=args.use_bias, + ) + self.gate = BailingMoeGate(args) + shared_dim = ( + args.moe_shared_expert_intermediate_size or args.moe_intermediate_size + ) + self.shared_experts = ( + BailingMoeMLP( + args=args, + intermediate_size=shared_dim * args.num_shared_experts, + ) + if args.num_shared_experts > 0 and args.moe_router_enable_shared_expert + else None + ) + + def __call__(self, x): + topk_idx, topk_weight = self.gate(x) + out = self.switch_mlp(x, topk_idx) + out = aggregate_expert_outputs(out, topk_weight) + if self.shared_experts is not None: + out = out + self.shared_experts(x) + return out + + +class BailingMoeDecoderLayer(nn.Module): + def __init__(self, args: ModelArgs, layer_idx: int): + super().__init__() + self.attention = BailingMoeAttention(args) + + self.mlp = ( + BailingMoeSparseMoeBlock(args) + if ( + args.num_experts is not None and layer_idx >= args.first_k_dense_replace + ) + else BailingMoeMLP(args) + ) + self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps) + self.post_attention_layernorm = nn.RMSNorm( + args.hidden_size, eps=args.rms_norm_eps + ) + + def __call__( + self, + x: mx.array, + mask: Optional[mx.array] = None, + cache: Optional[Any] = None, + ) -> mx.array: + r = self.attention(self.input_layernorm(x), mask, cache) + h = x + r + r = self.mlp(self.post_attention_layernorm(h)) + return h + r + + +class BailingMoeModel(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.word_embeddings = nn.Embedding(args.vocab_size, args.hidden_size) + self.layers = [ + BailingMoeDecoderLayer(args, layer_idx=i) + for i in range(args.num_hidden_layers) + ] + self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps) + + def __call__( + self, + inputs: mx.array, + cache: Optional[Any] = None, + ): + h = self.word_embeddings(inputs) + + if cache is None: + cache = [None] * len(self.layers) + + mask = create_attention_mask(h, cache[0]) + + for layer, c in zip(self.layers, cache): + h = layer(h, mask, c) + + return self.norm(h) + + +class Model(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.args = args + self.norm_head = args.norm_head + self.model_type = args.model_type + self.model = BailingMoeModel(args) + if not args.tie_word_embeddings: + self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False) + + def __call__( + self, + inputs: mx.array, + cache=None, + ): + out = self.model(inputs, cache) + if self.args.tie_word_embeddings: + out = self.model.word_embeddings.as_linear(out) + else: + out = self.lm_head(out) + return out + + def sanitize(self, weights): + if self.args.tie_word_embeddings: + weights.pop("lm_head.weight", None) + + if self.norm_head: + w = weights["lm_head.weight"] + dtype = w.dtype + weight_norm = ( + mx.linalg.norm(w.astype(mx.float32), axis=0, keepdims=True) + 1e-7 + ) + weights["lm_head.weight"] = (w / weight_norm).astype(dtype) + + for l in range(self.args.num_hidden_layers): + prefix = f"model.layers.{l}" + + if l >= self.args.first_k_dense_replace: + for m in ["gate_proj", "down_proj", "up_proj"]: + for k in ["weight", "scales", "biases"]: + if f"{prefix}.mlp.experts.0.{m}.{k}" in weights: + to_join = [ + weights.pop(f"{prefix}.mlp.experts.{e}.{m}.{k}") + for e in range(self.args.num_experts) + ] + weights[f"{prefix}.mlp.switch_mlp.{m}.{k}"] = mx.stack( + to_join + ) + + if f"{prefix}.mlp.gate.weight" in weights: + gate_weight = weights.pop(f"{prefix}.mlp.gate.weight") + weights[f"{prefix}.mlp.gate.gate_proj.weight"] = gate_weight + + if f"{prefix}.mlp.gate.bias" in weights: + gate_bias = weights.pop(f"{prefix}.mlp.gate.bias") + weights[f"{prefix}.mlp.gate.gate_proj.bias"] = gate_bias + + return weights + + @property + def quant_predicate(self): + def predicate(path, _): + if path.endswith("mlp.gate.gate_proj"): + return {"group_size": 64, "bits": 8} + return True + + return predicate + + @property + def cast_predicate(self): + def predicate(k): + return "expert_bias" not in k + + return predicate + + @property + def layers(self): + return self.model.layers diff --git a/mlx_audio/lm/models/base.py b/mlx_audio/lm/models/base.py new file mode 100644 index 000000000..6784acfb1 --- /dev/null +++ b/mlx_audio/lm/models/base.py @@ -0,0 +1,83 @@ +# Copyright © 2023-2024 Apple Inc. +# Vendored from mlx-lm v0.31.3 (ed1fca4cef15a824c5f1702c80f70b4cffc8e4dd), +# mlx_lm/models/base.py. Modified: dropped quantized_scaled_dot_product_attention +# and the quantized branch of scaled_dot_product_attention (no cache in mlx-audio +# exposes .bits). MIT licensed. + +import inspect +from dataclasses import dataclass +from typing import Optional + +import mlx.core as mx + + +@dataclass +class BaseModelArgs: + @classmethod + def from_dict(cls, params): + return cls( + **{ + k: v + for k, v in params.items() + if k in inspect.signature(cls).parameters + } + ) + + +def create_causal_mask( + N: int, + offset: int = 0, + window_size: Optional[int] = None, + right_padding: Optional[mx.array] = None, + left_padding: Optional[mx.array] = None, +): + rinds = mx.arange(offset + N) + linds = mx.arange(offset, offset + N) if offset else rinds + linds = linds[:, None] + rinds = rinds[None] + mask = linds >= rinds + if window_size is not None: + mask = mask & (linds < rinds + window_size) + if right_padding is not None: + mask = mask & (rinds < mx.expand_dims((offset + N) - right_padding, (1, 2, 3))) + if left_padding is not None: + mask = mask & (mx.expand_dims(left_padding, (1, 2, 3)) <= rinds) + return mask + + +def create_attention_mask( + h, cache=None, window_size: Optional[int] = None, return_array: bool = False +): + N = h.shape[1] + if cache and hasattr(cache, "make_mask"): + return cache.make_mask(N, return_array=return_array, window_size=window_size) + if N == 1: + return None + if return_array or (window_size and N > window_size): + return create_causal_mask(N, window_size=window_size) + return "causal" + + +def create_ssm_mask(h, cache=None): + if cache and hasattr(cache, "make_mask"): + return cache.make_mask(h.shape[1]) + return None + + +def scaled_dot_product_attention( + queries, + keys, + values, + cache, + scale: float, + mask: Optional[mx.array], + sinks: Optional[mx.array] = None, +) -> mx.array: + return mx.fast.scaled_dot_product_attention( + queries, + keys, + values, + scale=scale, + mask=mask, + sinks=sinks, + ) diff --git a/mlx_audio/lm/models/cache.py b/mlx_audio/lm/models/cache.py new file mode 100644 index 000000000..b038b438c --- /dev/null +++ b/mlx_audio/lm/models/cache.py @@ -0,0 +1,725 @@ +# Copyright © 2023-2024 Apple Inc. +# Vendored from mlx-lm v0.31.3 (ed1fca4cef15a824c5f1702c80f70b4cffc8e4dd), +# mlx_lm/models/cache.py. Modified: kept only the caches mlx-audio uses +# (KVCache, RotatingKVCache, BatchKVCache, ArraysCache) and dropped +# QuantizedKVCache, ConcatenateKVCache, ChunkedKVCache, CacheList, +# BatchRotatingKVCache, the prompt-cache save/load/trim helpers, the prompt +# trie/LRU cache, and the to_quantized methods. MIT licensed. + +from typing import Any, List, Optional + +import mlx.core as mx +import mlx.nn as nn + +from .base import create_causal_mask + + +def make_prompt_cache( + model: nn.Module, + max_kv_size: Optional[int] = None, +) -> List[Any]: + """ + Construct the model's cache for use in generation. + + This function will defer the cache construction to the model if it has a + ``make_cache`` method, otherwise it will make a default KV cache. + + Args: + model (nn.Module): The language model. + max_kv_size (Optional[int]): If provided and the model does not have a + ``make_cache`` method, a ``RotatingKVCache`` is used with a maximum + size of ``max_kv_size`` + """ + if hasattr(model, "make_cache"): + return model.make_cache() + + num_layers = len(model.layers) + if max_kv_size is not None: + return [ + RotatingKVCache(max_size=max_kv_size, keep=4) for _ in range(num_layers) + ] + else: + return [KVCache() for _ in range(num_layers)] + + +def create_attention_mask( + N: int, offset: int, return_array: bool, window_size: Optional[int] +): + if window_size is not None: + return create_causal_mask(N, offset, window_size=window_size) + elif N == 1: + return None + elif return_array: + return create_causal_mask(N, offset, window_size=window_size) + else: + return "causal" + + +class _BaseCache: + @property + def state(self): + return [] + + @state.setter + def state(self, v): + if v is not None and v: + raise ValueError("This cache has no state but a state was set.") + + @property + def meta_state(self): + return "" + + @meta_state.setter + def meta_state(self, v): + if v is not None and v: + raise ValueError("This cache has no meta_state but a meta_state was set.") + + def is_trimmable(self): + return False + + def size(self): + """ + Return the size (i.e. sequence length) of the cache. + + Not every cache is required to implement this, in which case the size + will always be 0 (though the cache may not be empty). + """ + return 0 + + @property + def nbytes(self): + """Return the size of this cache in bytes""" + raise NotImplementedError("Cache sub-class must implement nbytes") + + def empty(self): + """ + Return if the cache is empty or not. + """ + raise NotImplementedError("Cache sub-class must implement this.") + + @classmethod + def from_state(cls, state, meta_state): + # Create an instance of cls without calling __init__ + obj = cls.__new__(cls) + obj.state = state + obj.meta_state = meta_state + return obj + + +class KVCache(_BaseCache): + step = 256 + + def __init__(self): + self.keys = None + self.values = None + self.offset = 0 + + def update_and_fetch(self, keys, values): + prev = self.offset + if self.keys is None or (prev + keys.shape[2]) > self.keys.shape[2]: + B, n_kv_heads, _, k_head_dim = keys.shape + v_head_dim = values.shape[3] + n_steps = (self.step + keys.shape[2] - 1) // self.step + k_shape = (B, n_kv_heads, n_steps * self.step, k_head_dim) + v_shape = (B, n_kv_heads, n_steps * self.step, v_head_dim) + new_k = mx.zeros(k_shape, keys.dtype) + new_v = mx.zeros(v_shape, values.dtype) + if self.keys is not None: + if prev % self.step != 0: + self.keys = self.keys[..., :prev, :] + self.values = self.values[..., :prev, :] + self.keys = mx.concatenate([self.keys, new_k], axis=2) + self.values = mx.concatenate([self.values, new_v], axis=2) + else: + self.keys, self.values = new_k, new_v + + self.offset += keys.shape[2] + self.keys[..., prev : self.offset, :] = keys + self.values[..., prev : self.offset, :] = values + return self.keys[..., : self.offset, :], self.values[..., : self.offset, :] + + def size(self): + return self.offset + + @property + def state(self): + if self.offset == self.keys.shape[2]: + return self.keys, self.values + else: + return ( + self.keys[..., : self.offset, :], + self.values[..., : self.offset, :], + ) + + @state.setter + def state(self, v): + self.keys, self.values = v + self.offset = self.keys.shape[2] + + def is_trimmable(self): + return True + + def trim(self, n): + n = min(self.offset, n) + self.offset -= n + return n + + def make_mask(self, *args, **kwargs): + return create_attention_mask(*args, offset=self.offset, **kwargs) + + @classmethod + def merge(_, caches): + return BatchKVCache.merge(caches) + + def empty(self): + return self.keys is None + + @property + def nbytes(self): + if self.keys is None: + return 0 + return self.keys.nbytes + self.values.nbytes + + +class RotatingKVCache(_BaseCache): + step = 256 + + def __init__(self, max_size, keep=0): + self.keep = keep + self.keys = None + self.values = None + self.offset = 0 + self.max_size = max_size + self._idx = 0 + + def _trim(self, trim_size, v, append=None): + to_cat = [] + if trim_size > 0: + to_cat = [v[..., : self.keep, :], v[..., trim_size + self.keep :, :]] + else: + to_cat = [v] + if append is not None: + to_cat.append(append) + return mx.concatenate(to_cat, axis=2) + + def _temporal_order(self, v): + """ + Rearrange the cache into temporal order, slicing off the end if unused. + """ + if self._idx == v.shape[2]: + return v + elif self._idx < self.offset: + return mx.concatenate( + [ + v[..., : self.keep, :], + v[..., self._idx :, :], + v[..., self.keep : self._idx, :], + ], + axis=2, + ) + else: + return v[..., : self._idx, :] + + def _update_concat(self, keys, values): + if self.keys is None: + self.keys = keys + self.values = values + else: + # Put the keys/values in temporal order to + # preserve context + self.keys = self._temporal_order(self.keys) + self.values = self._temporal_order(self.values) + self._idx = self.keys.shape[2] + + # The largest size is self.max_size + S - 1 to ensure + # every token gets at least self.max_size context + trim_size = self._idx - self.max_size + 1 + self.keys = self._trim(trim_size, self.keys, keys) + self.values = self._trim(trim_size, self.values, values) + self.offset += keys.shape[2] + self._idx = self.keys.shape[2] + return self.keys, self.values + + def _update_in_place(self, keys, values): + # May not have hit the max size yet, so potentially + # keep growing the cache + B, n_kv_heads, S, k_head_dim = keys.shape + prev = self.offset + if self.keys is None or ( + prev >= self.keys.shape[2] and self.keys.shape[2] < self.max_size + ): + v_head_dim = values.shape[3] + new_size = min(self.step, self.max_size - prev) + k_shape = (B, n_kv_heads, new_size, k_head_dim) + v_shape = (B, n_kv_heads, new_size, v_head_dim) + new_k = mx.zeros(k_shape, keys.dtype) + new_v = mx.zeros(v_shape, values.dtype) + if self.keys is not None: + self.keys = mx.concatenate([self.keys, new_k], axis=2) + self.values = mx.concatenate([self.values, new_v], axis=2) + else: + self.keys, self.values = new_k, new_v + self._idx = prev + + # Trim if needed + trim_size = self.keys.shape[2] - self.max_size + if trim_size > 0: + self.keys = self._trim(trim_size, self.keys) + self.values = self._trim(trim_size, self.values) + self._idx = self.max_size + + # Rotate + if self._idx == self.max_size: + self._idx = self.keep + + # Assign + self.keys[..., self._idx : self._idx + S, :] = keys + self.values[..., self._idx : self._idx + S, :] = values + self.offset += S + self._idx += S + + # If the buffer is not full, slice off the end + if self.offset < self.max_size: + return self.keys[..., : self.offset, :], self.values[..., : self.offset, :] + return self.keys, self.values + + def update_and_fetch(self, keys, values): + if keys.shape[2] == 1: + return self._update_in_place(keys, values) + return self._update_concat(keys, values) + + def size(self): + return min(self.offset, self.max_size) + + @property + def state(self): + if self.offset < self.keys.shape[2]: + return self.keys[..., : self.offset, :], self.values[..., : self.offset, :] + else: + return self.keys, self.values + + @state.setter + def state(self, v): + self.keys, self.values = v + + @property + def meta_state(self): + return tuple(map(str, (self.keep, self.max_size, self.offset, self._idx))) + + @meta_state.setter + def meta_state(self, v): + self.keep, self.max_size, self.offset, self._idx = map( + int, + v, + ) + + def is_trimmable(self): + return self.offset < self.max_size + + def trim(self, n): + n = min(self.offset, n) + self.offset -= n + self._idx -= n + return n + + def make_mask( + self, N: int, window_size: Optional[int] = None, return_array: bool = False + ): + if N > 1: + window_size = window_size or self.max_size + offset = min(self.max_size - 1, self.offset) + if offset + N > window_size or return_array: + return create_causal_mask(N, offset, window_size=window_size) + else: + return "causal" + else: + if window_size is None: + return None + # May need a mask for when window_size < max_size + if self.offset >= window_size and self.max_size > window_size: + idx = self._idx + if idx >= self.max_size: + idx = 0 + if self.offset < self.max_size: + mask_size = self.offset + 1 + else: + mask_size = self.max_size + mask = mx.arange(mask_size) >= (mask_size - window_size) + mask = mx.roll(mask, shift=idx + 1) + return mask + + def empty(self): + return self.keys is None + + @property + def nbytes(self): + if self.keys is None: + return 0 + return self.keys.nbytes + self.values.nbytes + + +class ArraysCache(_BaseCache): + def __new__(cls, *args, **kwargs): + instance = super().__new__(cls) + instance.left_padding = None + instance.lengths = None + return instance + + def __init__(self, size, left_padding: Optional[List[int]] = None): + self.cache = [None] * size + if left_padding: + self.left_padding = mx.array(left_padding) + + @property + def batch_size(self): + for c in self.cache: + if c is not None: + return c.shape[0] + if self.left_padding is not None: + return self.left_padding.size + elif self.lengths is not None: + return self.lengths.size + else: + return 1 + + def __setitem__(self, idx, value): + self.cache[idx] = value + + def __getitem__(self, idx): + return self.cache[idx] + + @property + def state(self): + return self.cache + + @state.setter + def state(self, v): + self.cache = v + + def filter(self, batch_indices): + """ + In-place filter to keep just the given indices in the cache. + """ + self.cache = [c[batch_indices] if c is not None else None for c in self.cache] + if self.left_padding is not None: + self.left_padding = self.left_padding[batch_indices] + if self.lengths is not None: + self.lengths = self.lengths[batch_indices] + + def extend(self, other): + """ + In-place extend this cache with the other cache. + """ + + a_batch = self.batch_size + b_batch = other.batch_size + + def cat(a, b): + shape = dtype = None + if a is not None: + shape = a.shape + dtype = a.dtype + if b is not None: + shape = b.shape + dtype = b.dtype + + if shape is None: + return None + + if a is None: + a = mx.zeros((a_batch,) + shape[1:], dtype=dtype) + if b is None: + b = mx.zeros((b_batch,) + shape[1:], dtype=dtype) + + return mx.concatenate([a, b]) + + self.cache = [cat(c, o) for c, o in zip(self.cache, other.cache)] + self.left_padding = cat(self.left_padding, other.left_padding) + self.lengths = cat(self.lengths, other.lengths) + + def extract(self, idx): + cache = ArraysCache(len(self.cache)) + cache.cache = [c[idx : idx + 1] for c in self.cache] + return cache + + def prepare(self, lengths=None, **kwargs): + self.lengths = mx.array(lengths) + + def finalize(self): + self.lengths = None + self.left_padding = None + + def advance(self, N): + if self.lengths is not None: + self.lengths -= N + if self.left_padding is not None: + self.left_padding -= N + + def make_mask(self, N: int): + if self.left_padding is not None: + pos = mx.arange(N) + return pos >= self.left_padding[:, None] + elif self.lengths is not None: + pos = mx.arange(N) + return pos < self.lengths[:, None] + else: + return None + + @classmethod + def merge(cls, caches): + n_state = len(caches[0].cache) + B = len(caches) + cache = cls(n_state) + + # All caches are empty so return early + if all(c.empty() for c in caches): + cache.left_padding = mx.array([0] * B) + return cache + + for e in range(n_state): + c_init = next(iter(c[e] for c in caches if c[e] is not None)) + shape = list(c_init.shape) + shape[0] = B + cache[e] = mx.zeros(shape, c_init.dtype) + for i in range(B): + if caches[i][e] is None: + continue + cache[e][i : i + 1] = caches[i][e] + return cache + + def empty(self): + return self.cache[0] is None + + @property + def nbytes(self): + return sum(c.nbytes for c in self.cache if c is not None) + + +def dynamic_roll(x, shifts, axis): + n = x.shape[axis] + expand_shifts = (...,) + (None,) * (x.ndim - axis) + expand_indices = expand_shifts[:-1] + idx = (mx.arange(n)[expand_indices] - shifts[expand_shifts]) % n + rolled = mx.take_along_axis(x, idx, axis=axis) + return rolled + + +class BatchKVCache(_BaseCache): + step = 256 + + def __init__(self, left_padding: List[int]): + """ + The BatchKV cache expects inputs to be left-padded. + + E.g. the following prompts: + + [1, 3, 5] + [7] + [2, 6, 8, 9] + + Should be padded like so: + + [0, 1, 3, 5] + [0, 0, 0, 7] + [2, 6, 8, 9] + + And ``left_padding`` specifies the amount of padding for each. + In this case, ``left_padding = [1, 3, 0]``. + """ + self.keys = None + self.values = None + self.left_padding = mx.array(left_padding) + self.offset = mx.array([-l for l in left_padding]) + self._idx = 0 + + self._right_padding = None + + def update_and_fetch(self, keys, values): + prev = self._idx + if self.keys is None or (prev + keys.shape[2]) > self.keys.shape[2]: + B, n_kv_heads, _, k_head_dim = keys.shape + v_head_dim = values.shape[3] + n_steps = (self.step + keys.shape[2] - 1) // self.step + k_shape = (B, n_kv_heads, n_steps * self.step, k_head_dim) + v_shape = (B, n_kv_heads, n_steps * self.step, v_head_dim) + new_k = mx.zeros(k_shape, keys.dtype) + new_v = mx.zeros(v_shape, values.dtype) + if self.keys is not None: + if prev % self.step != 0: + self.keys = self.keys[..., :prev, :] + self.values = self.values[..., :prev, :] + self.keys = mx.concatenate([self.keys, new_k], axis=2) + self.values = mx.concatenate([self.values, new_v], axis=2) + else: + self.keys, self.values = new_k, new_v + + self.offset += keys.shape[2] + self._idx += keys.shape[2] + self.keys[..., prev : self._idx, :] = keys + self.values[..., prev : self._idx, :] = values + return self.keys[..., : self._idx, :], self.values[..., : self._idx, :] + + def prepare(self, *, left_padding=None, lengths=None, right_padding=None): + if left_padding is not None: + if self.keys is not None: + raise ValueError( + "Left padding can only be added to an empty BatchKVCache" + ) + left_padding = mx.array(left_padding) + self.left_padding += left_padding + self.offset -= left_padding + + if right_padding is not None and max(right_padding) > 0: + self._right_padding = mx.array(right_padding) + + def finalize(self): + if self._right_padding is not None: + padding = self._right_padding + self.keys = dynamic_roll(self.keys, padding[:, None], axis=2) + self.values = dynamic_roll(self.values, padding[:, None], axis=2) + self.offset -= padding + self.left_padding += padding + self._right_padding = None + + @property + def state(self): + k, v = self.keys, self.values + if self._idx < k.shape[2]: + k = k[..., : self._idx, :] + v = v[..., : self._idx, :] + return k, v, self.offset, self.left_padding + + @state.setter + def state(self, v): + self.keys, self.values, self.offset, self.left_padding = v + self._idx = self.keys.shape[2] + + def is_trimmable(self): + return True + + def trim(self, n): + n = min(self._idx, n) + self._idx -= n + self.offset -= n + return n + + def make_mask(self, N: int, return_array: bool = False, **kwargs): + return create_causal_mask( + N, offset=self._idx, left_padding=self.left_padding, **kwargs + ) + + def filter(self, batch_indices): + """ + In-place filter to keep just the given indices in the cache. + """ + if self.keys is not None: + self.keys = self.keys[batch_indices] + self.values = self.values[batch_indices] + self.offset = self.offset[batch_indices] + self.left_padding = self.left_padding[batch_indices] + + # Shift left to reduce padding + min_left_pad = self.left_padding.min().item() + if min_left_pad > 0: + if self.keys is not None: + self.keys = self.keys[..., min_left_pad:, :] + self.values = self.values[..., min_left_pad:, :] + self._idx -= min_left_pad + self.left_padding -= min_left_pad + + def extend(self, other): + """ + In-place extend this cache with the other cache. + """ + if self.keys is None and other.keys is None: + self.left_padding = mx.concatenate([self.left_padding, other.left_padding]) + self.offset = mx.concatenate([self.offset, other.offset]) + return + + max_idx = max(self._idx, other._idx) + L1 = L2 = 0 + if self.keys is not None: + B, H, L1, D = self.keys.shape + M = self.values.shape[3] + if other.keys is not None: + B, H, L2, D = other.keys.shape + M = other.values.shape[3] + max_size = max(L1, L2) + + # Pad the keys and values so they are right-justified + # with the index and the same size + def pad(c): + k, v = c.keys, c.values + if k is None: + Bc = c.offset.shape[0] + k = mx.array([]).reshape(Bc, H, 0, D) + v = mx.array([]).reshape(Bc, H, 0, M) + left = max_idx - c._idx + right = max_size - k.shape[2] - left + if right < 0: + k = k[..., :right, :] + v = v[..., :right, :] + right = 0 + if left != 0 or right != 0: + pad = [(0, 0), (0, 0), (left, right), (0, 0)] + k = mx.pad(k, pad) + v = mx.pad(v, pad) + left_padding = c.left_padding + left + return k, v, c.offset, left_padding + + self.keys, self.values, self.offset, self.left_padding = map( + mx.concatenate, zip(*(pad(self), pad(other))) + ) + self._idx = max_idx + + def extract(self, idx): + cache = KVCache() + padding = self.left_padding[idx].item() + cache.keys = mx.contiguous(self.keys[idx : idx + 1, :, padding : self._idx]) + cache.values = mx.contiguous(self.values[idx : idx + 1, :, padding : self._idx]) + cache.offset = cache.keys.shape[2] + return cache + + @classmethod + def merge(cls, caches): + lengths = [c.size() for c in caches] + max_length = max(lengths) + + # No cache has content so make an empty one + if max_length == 0: + return BatchKVCache([0] * len(caches)) + + padding = [max_length - l for l in lengths] + B = len(caches) + H = max(c.keys.shape[1] for c in caches if c.keys is not None) + Dk = max(c.keys.shape[3] for c in caches if c.keys is not None) + Dv = max(c.values.shape[3] for c in caches if c.values is not None) + dt = next(iter(c.keys.dtype for c in caches if c.keys is not None)) + + keys = mx.zeros((B, H, max_length, Dk), dtype=dt) + values = mx.zeros((B, H, max_length, Dv), dtype=dt) + for i, (p, c) in enumerate(zip(padding, caches)): + if c.keys is None: + continue + keys[i : i + 1, :, p : p + c.offset] = c.keys[..., : c.offset, :] + values[i : i + 1, :, p : p + c.offset] = c.values[..., : c.offset, :] + + cache = cls(padding) + cache.keys = keys + cache.values = values + cache.offset += keys.shape[2] + cache._idx = keys.shape[2] + + return cache + + def size(self): + return self._idx + + def empty(self): + return self.keys is None + + @property + def nbytes(self): + if self.keys is None: + return 0 + return self.keys.nbytes + self.values.nbytes diff --git a/mlx_audio/lm/models/gemma3.py b/mlx_audio/lm/models/gemma3.py new file mode 100644 index 000000000..1e378d0af --- /dev/null +++ b/mlx_audio/lm/models/gemma3.py @@ -0,0 +1,65 @@ +# Copyright © 2025 Apple Inc. +# Vendored verbatim from mlx-lm v0.31.3 (ed1fca4cef15a824c5f1702c80f70b4cffc8e4dd), +# mlx_lm/models/gemma3.py. MIT licensed. + +from dataclasses import dataclass +from typing import Optional + +import mlx.core as mx +import mlx.nn as nn +from mlx.utils import tree_flatten, tree_unflatten + +from . import gemma3_text +from .base import BaseModelArgs + + +@dataclass +class ModelArgs(BaseModelArgs): + model_type: str + text_config: dict + vocab_size: int = 262208 + + def __post_init__(self): + self.text_config["vocab_size"] = self.vocab_size + self.text_config["num_attention_heads"] = self.text_config.get( + "num_attention_heads", 8 + ) + self.text_config["num_key_value_heads"] = self.text_config.get( + "num_key_value_heads", 4 + ) + + +class Model(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.args = args + self.model_type = args.model_type + self.language_model = gemma3_text.Model( + gemma3_text.ModelArgs.from_dict(args.text_config) + ) + + def __call__( + self, + inputs: mx.array, + cache=None, + input_embeddings: Optional[mx.array] = None, + ): + return self.language_model( + inputs, cache=cache, input_embeddings=input_embeddings + ) + + def sanitize(self, weights): + weights = tree_unflatten(list(weights.items())) + weights.pop("vision_tower", None) + weights.pop("multi_modal_projector", None) + lm_weights = dict(tree_flatten(weights["language_model"])) + lm_weights = self.language_model.sanitize(lm_weights) + weights["language_model"] = tree_unflatten(list(lm_weights.items())) + return dict(tree_flatten(weights)) + + @property + def layers(self): + return self.language_model.layers + + def make_cache(self): + return self.language_model.make_cache() diff --git a/mlx_audio/lm/models/gemma3_text.py b/mlx_audio/lm/models/gemma3_text.py new file mode 100644 index 000000000..abb9f9b8d --- /dev/null +++ b/mlx_audio/lm/models/gemma3_text.py @@ -0,0 +1,259 @@ +# Copyright © 2025 Apple Inc. +# Vendored verbatim from mlx-lm v0.31.3 (ed1fca4cef15a824c5f1702c80f70b4cffc8e4dd), +# mlx_lm/models/gemma3_text.py. MIT licensed. + +from dataclasses import dataclass +from functools import partial +from typing import Any, Dict, Optional + +import mlx.core as mx +import mlx.nn as nn + +from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention +from .cache import KVCache, RotatingKVCache +from .rope_utils import initialize_rope + + +@dataclass +class ModelArgs(BaseModelArgs): + model_type: str + hidden_size: int = 1152 + num_hidden_layers: int = 26 + intermediate_size: int = 6912 + num_attention_heads: int = 4 + head_dim: int = 256 + rms_norm_eps: float = 1.0e-6 + vocab_size: int = 262144 + num_key_value_heads: int = 1 + rope_theta: float = 1_000_000.0 + rope_local_base_freq: float = 10_000.0 + query_pre_attn_scalar: float = 256 + sliding_window: int = 512 + sliding_window_pattern: int = 6 + max_position_embeddings: int = 32768 + rope_scaling: Dict = None + + +class Attention(nn.Module): + def __init__(self, args: ModelArgs, layer_idx: int): + super().__init__() + + dim = args.hidden_size + self.n_heads = n_heads = args.num_attention_heads + self.n_kv_heads = n_kv_heads = args.num_key_value_heads + self.repeats = n_heads // n_kv_heads + self.head_dim = head_dim = args.head_dim + self.layer_idx = layer_idx + + self.scale = args.query_pre_attn_scalar**-0.5 + + self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=False) + self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False) + self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False) + self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=False) + + self.q_norm = RMSNorm(dims=head_dim, eps=args.rms_norm_eps) + self.k_norm = RMSNorm(dims=head_dim, eps=args.rms_norm_eps) + self.is_sliding = (layer_idx + 1) % args.sliding_window_pattern != 0 + + if self.is_sliding: + self.rope = initialize_rope( + dims=head_dim, + base=args.rope_local_base_freq, + traditional=False, + ) + else: + self.rope = initialize_rope( + dims=head_dim, + base=args.rope_theta, + traditional=False, + max_position_embeddings=args.max_position_embeddings, + scaling_config=args.rope_scaling, + ) + + def __call__( + self, + x: mx.array, + mask: Optional[mx.array] = None, + cache: Optional[Any] = None, + ) -> mx.array: + B, L, _ = x.shape + queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x) + queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3) + + keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3) + values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3) + + queries = self.q_norm(queries) + keys = self.k_norm(keys) + + if cache is not None: + queries = self.rope(queries, offset=cache.offset) + keys = self.rope(keys, offset=cache.offset) + keys, values = cache.update_and_fetch(keys, values) + else: + queries = self.rope(queries) + keys = self.rope(keys) + + # Sliding window + output = scaled_dot_product_attention( + queries, keys, values, cache=cache, scale=self.scale, mask=mask + ) + output = output.transpose(0, 2, 1, 3).reshape(B, L, -1) + return self.o_proj(output) + + +class RMSNorm(nn.Module): + def __init__(self, dims: int, eps: float = 1e-5): + super().__init__() + self.weight = mx.ones((dims,)) + self.eps = eps + + def __call__(self, x): + return mx.fast.rms_norm(x, 1.0 + self.weight, self.eps) + + +class MLP(nn.Module): + def __init__(self, dim, hidden_dim): + super().__init__() + self.gate_proj = nn.Linear(dim, hidden_dim, bias=False) + self.down_proj = nn.Linear(hidden_dim, dim, bias=False) + self.up_proj = nn.Linear(dim, hidden_dim, bias=False) + + def __call__(self, x) -> mx.array: + return self.down_proj(nn.gelu_approx(self.gate_proj(x)) * self.up_proj(x)) + + +@partial(mx.compile, shapeless=True) +def clip_residual(x, y): + if x.dtype != mx.float16: + return x + y + bound = mx.finfo(mx.float16).max + return mx.clip(x.astype(mx.float32) + y.astype(mx.float32), -bound, bound).astype( + mx.float16 + ) + + +class TransformerBlock(nn.Module): + def __init__(self, args: ModelArgs, layer_idx: int): + super().__init__() + self.num_attention_heads = args.num_attention_heads + self.hidden_size = args.hidden_size + self.self_attn = Attention(args, layer_idx) + self.mlp = MLP(args.hidden_size, args.intermediate_size) + self.input_layernorm = RMSNorm(args.hidden_size, eps=args.rms_norm_eps) + self.post_attention_layernorm = RMSNorm(args.hidden_size, eps=args.rms_norm_eps) + self.pre_feedforward_layernorm = RMSNorm( + args.hidden_size, eps=args.rms_norm_eps + ) + self.post_feedforward_layernorm = RMSNorm( + args.hidden_size, eps=args.rms_norm_eps + ) + + def __call__( + self, + x: mx.array, + mask: Optional[mx.array] = None, + cache: Optional[Any] = None, + ) -> mx.array: + r = self.self_attn(self.input_layernorm(x), mask, cache) + h = clip_residual(x, self.post_attention_layernorm(r)) + r = self.mlp(self.pre_feedforward_layernorm(h)) + out = clip_residual(h, self.post_feedforward_layernorm(r)) + return out + + +class Gemma3Model(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.args = args + self.window_size = args.sliding_window + self.sliding_window_pattern = args.sliding_window_pattern + self.vocab_size = args.vocab_size + self.num_hidden_layers = args.num_hidden_layers + assert self.vocab_size > 0 + self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size) + self.layers = [ + TransformerBlock(args=args, layer_idx=layer_idx) + for layer_idx in range(args.num_hidden_layers) + ] + self.norm = RMSNorm(args.hidden_size, eps=args.rms_norm_eps) + + def __call__( + self, + inputs: mx.array, + cache=None, + input_embeddings: Optional[mx.array] = None, + ): + if input_embeddings is not None: + h = input_embeddings + else: + h = self.embed_tokens(inputs) + h *= mx.array(self.args.hidden_size**0.5, mx.bfloat16).astype(h.dtype) + + if cache is None: + cache = [None] * len(self.layers) + + global_mask = create_attention_mask(h, cache[self.sliding_window_pattern - 1]) + + if self.sliding_window_pattern > 1: + sliding_window_mask = create_attention_mask( + h, + cache[0], + window_size=self.window_size, + ) + else: + sliding_window_mask = None + for i, (layer, c) in enumerate(zip(self.layers, cache)): + is_global = ( + i % self.sliding_window_pattern == self.sliding_window_pattern - 1 + ) + mask = global_mask if is_global else sliding_window_mask + h = layer(h, mask, c) + + return self.norm(h) + + +class Model(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.args = args + self.model_type = args.model_type + self.model = Gemma3Model(args) + self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False) + self.tie_word_embeddings = False + + def __call__( + self, + inputs: mx.array, + cache=None, + input_embeddings: Optional[mx.array] = None, + ): + out = self.model(inputs, cache, input_embeddings) + if self.tie_word_embeddings: + out = self.model.embed_tokens.as_linear(out) + else: + out = self.lm_head(out) + return out + + def sanitize(self, weights): + if "lm_head.weight" not in weights: + self.tie_word_embeddings = True + self.pop("lm_head") + return weights + + @property + def layers(self): + return self.model.layers + + def make_cache(self): + caches = [] + for i in range(self.args.num_hidden_layers): + if ( + i % self.args.sliding_window_pattern + == self.args.sliding_window_pattern - 1 + ): + caches.append(KVCache()) + else: + caches.append(RotatingKVCache(max_size=self.args.sliding_window)) + return caches diff --git a/mlx_audio/lm/models/gpt2.py b/mlx_audio/lm/models/gpt2.py new file mode 100644 index 000000000..fb6660774 --- /dev/null +++ b/mlx_audio/lm/models/gpt2.py @@ -0,0 +1,202 @@ +# Copyright © 2023 - 2024 Apple Inc. +# Vendored verbatim from mlx-lm v0.31.3 (ed1fca4cef15a824c5f1702c80f70b4cffc8e4dd), +# mlx_lm/models/gpt2.py. MIT licensed. + +from dataclasses import dataclass +from typing import Any, Optional + +import mlx.core as mx +import mlx.nn as nn + +from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention + + +@dataclass +class ModelArgs(BaseModelArgs): + model_type: str + n_ctx: int + n_embd: int + n_head: int + n_layer: int + n_positions: int + layer_norm_epsilon: float + vocab_size: int + num_key_value_heads: int = None + + def __post_init__(self): + if self.num_key_value_heads is None: + self.num_key_value_heads = self.n_head + + +class Attention(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + + assert args.n_embd % args.n_head == 0, "n_embd must be divisible by n_head" + + self.n_embd = args.n_embd + self.n_head = args.n_head + self.head_dim = self.n_embd // self.n_head + + self.scale = self.head_dim**-0.5 + + self.c_attn = nn.Linear(self.n_embd, 3 * self.n_embd, bias=True) + self.c_proj = nn.Linear(self.n_embd, self.n_embd, bias=True) + + def __call__( + self, + x: mx.array, + mask: Optional[mx.array] = None, + cache: Optional[Any] = None, + ) -> mx.array: + B, L, D = x.shape + + qkv = self.c_attn(x) + queries, keys, values = mx.split(qkv, 3, axis=-1) + + # Prepare the queries, keys and values for the attention computation + queries = queries.reshape(B, L, self.n_head, -1).transpose(0, 2, 1, 3) + keys = keys.reshape(B, L, self.n_head, -1).transpose(0, 2, 1, 3) + values = values.reshape(B, L, self.n_head, -1).transpose(0, 2, 1, 3) + + if cache is not None: + keys, values = cache.update_and_fetch(keys, values) + + output = scaled_dot_product_attention( + queries, keys, values, cache=cache, scale=self.scale, mask=mask + ) + + output = output.transpose(0, 2, 1, 3).reshape(B, L, -1) + return self.c_proj(output) + + +class MLP(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + + self.n_embd = args.n_embd + self.c_fc = nn.Linear(self.n_embd, 4 * self.n_embd) + self.c_proj = nn.Linear(4 * self.n_embd, self.n_embd) + + def __call__(self, x) -> mx.array: + return self.c_proj(nn.gelu_approx(self.c_fc(x))) + + +class TransformerBlock(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + + self.n_head = args.n_head + self.n_embd = args.n_embd + self.layer_norm_epsilon = args.layer_norm_epsilon + self.attn = Attention(args) + self.mlp = MLP(args) + self.ln_1 = nn.LayerNorm( + self.n_embd, + eps=self.layer_norm_epsilon, + ) + self.ln_2 = nn.LayerNorm(self.n_embd, eps=self.layer_norm_epsilon) + + def __call__( + self, + x: mx.array, + mask: Optional[mx.array] = None, + cache: Optional[Any] = None, + ) -> mx.array: + r = self.attn(self.ln_1(x), mask, cache) + h = x + r + r = self.mlp(self.ln_2(h)) + out = h + r + return out + + +class GPT2Model(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.n_embd = args.n_embd + self.n_positions = args.n_positions + self.vocab_size = args.vocab_size + self.n_layer = args.n_layer + self.layer_norm_epsilon = args.layer_norm_epsilon + assert self.vocab_size > 0 + self.wte = nn.Embedding(self.vocab_size, self.n_embd) + self.wpe = nn.Embedding(self.n_positions, self.n_embd) + self.h = [TransformerBlock(args=args) for _ in range(self.n_layer)] + self.ln_f = nn.LayerNorm(self.n_embd, eps=self.layer_norm_epsilon) + + def __call__( + self, + inputs: mx.array, + cache=None, + ): + _, L = inputs.shape + + hidden_states = self.wte(inputs) + + if cache is None: + cache = [None] * len(self.h) + + offset = 0 + if cache[0] is not None: + offset = cache[0].offset + + offset = mx.array(offset) + position_ids = mx.arange(L) + offset[..., None] + + hidden_states += self.wpe(position_ids) + + mask = create_attention_mask(hidden_states, cache[0]) + + for layer, c in zip(self.h, cache): + hidden_states = layer(hidden_states, mask, cache=c) + + return self.ln_f(hidden_states) + + +class Model(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.args = args + self.model_type = args.model_type + self.model = GPT2Model(args) + + def __call__( + self, + inputs: mx.array, + cache=None, + ): + out = self.model(inputs, cache) + out = self.model.wte.as_linear(out) + return out + + def sanitize(self, weights): + new_weights = {} + for i in range(self.args.n_layer): + if f"h.{i}.attn.bias" in weights: + del weights[f"h.{i}.attn.bias"] + if f"h.{i}.attn.c_attn.weight" in weights: + weights[f"h.{i}.attn.c_attn.weight"] = weights[ + f"h.{i}.attn.c_attn.weight" + ].transpose(1, 0) + if f"h.{i}.attn.c_proj.weight" in weights: + weights[f"h.{i}.attn.c_proj.weight"] = weights[ + f"h.{i}.attn.c_proj.weight" + ].transpose(1, 0) + if f"h.{i}.mlp.c_fc.weight" in weights: + weights[f"h.{i}.mlp.c_fc.weight"] = weights[ + f"h.{i}.mlp.c_fc.weight" + ].transpose(1, 0) + if f"h.{i}.mlp.c_proj.weight" in weights: + weights[f"h.{i}.mlp.c_proj.weight"] = weights[ + f"h.{i}.mlp.c_proj.weight" + ].transpose(1, 0) + for weight in weights: + if not weight.startswith("model."): + new_weights[f"model.{weight}"] = weights[weight] + else: + new_weights[weight] = weights[weight] + return new_weights + + @property + def layers(self): + return self.model.h diff --git a/mlx_audio/lm/models/granite.py b/mlx_audio/lm/models/granite.py new file mode 100644 index 000000000..3c21ac289 --- /dev/null +++ b/mlx_audio/lm/models/granite.py @@ -0,0 +1,195 @@ +# Copyright © 2023-2024 Apple Inc. +# Vendored verbatim from mlx-lm v0.31.3 (ed1fca4cef15a824c5f1702c80f70b4cffc8e4dd), +# mlx_lm/models/granite.py. MIT licensed. + +from dataclasses import dataclass +from typing import Any, Dict, Optional, Union + +import mlx.core as mx +import mlx.nn as nn + +from .activations import swiglu +from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention +from .rope_utils import initialize_rope + + +@dataclass +class ModelArgs(BaseModelArgs): + model_type: str + hidden_size: int + num_hidden_layers: int + intermediate_size: int + num_attention_heads: int + rms_norm_eps: float + vocab_size: int + logits_scaling: float + attention_multiplier: float + embedding_multiplier: float + residual_multiplier: float + max_position_embeddings: int + num_key_value_heads: int + attention_bias: bool + mlp_bias: bool + rope_theta: float + rope_scaling: Optional[Dict[str, Union[float, str]]] = None + tie_word_embeddings: bool = True + + +class Attention(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + + dim = args.hidden_size + self.n_heads = n_heads = args.num_attention_heads + self.n_kv_heads = n_kv_heads = args.num_key_value_heads + + self.head_dim = head_dim = args.hidden_size // n_heads + + self.scale = args.attention_multiplier + attention_bias = args.attention_bias + self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=attention_bias) + self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=attention_bias) + self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=attention_bias) + self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=attention_bias) + + self.rope = initialize_rope( + self.head_dim, + args.rope_theta, + False, + args.rope_scaling, + args.max_position_embeddings, + ) + + def __call__( + self, + x: mx.array, + mask: Optional[mx.array] = None, + cache: Optional[Any] = None, + ) -> mx.array: + B, L, D = x.shape + + queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x) + + # Prepare the queries, keys and values for the attention computation + queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3) + keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3) + values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3) + + if cache is not None: + queries = self.rope(queries, offset=cache.offset) + keys = self.rope(keys, offset=cache.offset) + keys, values = cache.update_and_fetch(keys, values) + else: + queries = self.rope(queries) + keys = self.rope(keys) + + output = scaled_dot_product_attention( + queries, keys, values, cache=cache, scale=self.scale, mask=mask + ) + + output = output.transpose(0, 2, 1, 3).reshape(B, L, -1) + return self.o_proj(output) + + +class MLP(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + + dim = args.hidden_size + hidden_dim = args.intermediate_size + if hasattr(args, "mlp_bias"): + mlp_bias = args.mlp_bias + else: + mlp_bias = False + + self.gate_proj = nn.Linear(dim, hidden_dim, bias=mlp_bias) + self.down_proj = nn.Linear(hidden_dim, dim, bias=mlp_bias) + self.up_proj = nn.Linear(dim, hidden_dim, bias=mlp_bias) + + def __call__(self, x) -> mx.array: + return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x))) + + +class TransformerBlock(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.num_attention_heads = args.num_attention_heads + self.hidden_size = args.hidden_size + self.self_attn = Attention(args) + self.mlp = MLP(args) + self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps) + self.post_attention_layernorm = nn.RMSNorm( + args.hidden_size, eps=args.rms_norm_eps + ) + self.residual_multiplier = args.residual_multiplier + + def __call__( + self, + x: mx.array, + mask: Optional[mx.array] = None, + cache: Optional[Any] = None, + ) -> mx.array: + r = self.self_attn(self.input_layernorm(x), mask, cache) + h = x + r * self.residual_multiplier + r = self.mlp(self.post_attention_layernorm(h)) + out = h + r * self.residual_multiplier + return out + + +class GraniteModel(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.args = args + self.vocab_size = args.vocab_size + self.num_hidden_layers = args.num_hidden_layers + assert self.vocab_size > 0 + self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size) + self.layers = [ + TransformerBlock(args=args) for _ in range(args.num_hidden_layers) + ] + self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps) + self.embedding_multiplier = args.embedding_multiplier + + def __call__( + self, + inputs: mx.array, + cache=None, + ): + h = self.embed_tokens(inputs) * self.embedding_multiplier + + if cache is None: + cache = [None] * len(self.layers) + + mask = create_attention_mask(h, cache[0]) + + for layer, c in zip(self.layers, cache): + h = layer(h, mask, cache=c) + + return self.norm(h) + + +class Model(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.args = args + self.model_type = args.model_type + self.model = GraniteModel(args) + if not args.tie_word_embeddings: + self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False) + self.logits_scaling = args.logits_scaling + + def __call__( + self, + inputs: mx.array, + cache=None, + ): + out = self.model(inputs, cache) + if self.args.tie_word_embeddings: + out = self.model.embed_tokens.as_linear(out) + else: + out = self.lm_head(out) + return out / self.logits_scaling + + @property + def layers(self): + return self.model.layers diff --git a/mlx_audio/lm/models/lfm2.py b/mlx_audio/lm/models/lfm2.py new file mode 100644 index 000000000..80781127e --- /dev/null +++ b/mlx_audio/lm/models/lfm2.py @@ -0,0 +1,318 @@ +# Copyright © 2025 Apple Inc. +# Vendored verbatim from mlx-lm v0.31.3 (ed1fca4cef15a824c5f1702c80f70b4cffc8e4dd), +# mlx_lm/models/lfm2.py. MIT licensed. +from dataclasses import dataclass +from typing import Any, List, Optional + +import mlx.core as mx +import mlx.nn as nn + +from .activations import swiglu +from .base import ( + BaseModelArgs, + create_attention_mask, + create_ssm_mask, + scaled_dot_product_attention, +) +from .cache import ArraysCache, KVCache + + +@dataclass +class ModelArgs(BaseModelArgs): + model_type: str + vocab_size: int + hidden_size: int + num_hidden_layers: int + num_attention_heads: int + num_key_value_heads: int + max_position_embeddings: int + norm_eps: float + conv_bias: bool + conv_L_cache: int + block_dim: int + block_ff_dim: int + block_multiple_of: int + block_ffn_dim_multiplier: float + block_auto_adjust_ff_dim: bool + rope_theta: float = 1000000.0 + rope_parameters: Optional[dict] = None + full_attn_idxs: Optional[List[int]] = None + layer_types: Optional[List[str]] = None + + def __post_init__(self): + if self.rope_parameters is not None and "rope_theta" in self.rope_parameters: + self.rope_theta = self.rope_parameters["rope_theta"] + if self.num_key_value_heads is None: + self.num_key_value_heads = self.num_attention_heads + if self.full_attn_idxs is None: + self.full_attn_idxs = [ + i + for i, layer_type in enumerate(self.layer_types) + if layer_type == "full_attention" + ] + + +class Attention(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + + dim = args.hidden_size + self.n_heads = n_heads = args.num_attention_heads + self.n_kv_heads = n_kv_heads = args.num_key_value_heads + + self.head_dim = head_dim = args.hidden_size // n_heads + + self.scale = head_dim**-0.5 + + self.q_layernorm = nn.RMSNorm(head_dim, eps=args.norm_eps) + self.k_layernorm = nn.RMSNorm(head_dim, eps=args.norm_eps) + + self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=False) + self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False) + self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False) + self.out_proj = nn.Linear(n_heads * head_dim, dim, bias=False) + + self.rope = nn.RoPE( + self.head_dim, + base=args.rope_theta, + traditional=False, + ) + + def __call__( + self, + x: mx.array, + mask: Optional[mx.array] = None, + cache: Optional[Any] = None, + ) -> mx.array: + B, L, D = x.shape + + queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x) + + queries = self.q_layernorm(queries.reshape(B, L, self.n_heads, -1)).transpose( + 0, 2, 1, 3 + ) + keys = self.k_layernorm(keys.reshape(B, L, self.n_kv_heads, -1)).transpose( + 0, 2, 1, 3 + ) + values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3) + + if cache is not None: + queries = self.rope(queries, offset=cache.offset) + keys = self.rope(keys, offset=cache.offset) + keys, values = cache.update_and_fetch(keys, values) + else: + queries = self.rope(queries) + keys = self.rope(keys) + + output = scaled_dot_product_attention( + queries, keys, values, cache=cache, mask=mask, scale=self.scale + ) + output = output.transpose(0, 2, 1, 3).reshape(B, L, -1) + return self.out_proj(output) + + +class ShortConv(nn.Module): + def __init__( + self, + args: ModelArgs, + layer_idx: int, + ): + super().__init__() + self.args = args + self.layer_idx = layer_idx + self.L_cache = args.conv_L_cache + self.bias = args.conv_bias + + self.conv = nn.Conv1d( + in_channels=args.hidden_size, + out_channels=args.hidden_size, + kernel_size=self.L_cache, + groups=args.hidden_size, + bias=self.bias, + ) + self.in_proj = nn.Linear(args.hidden_size, 3 * args.hidden_size, bias=self.bias) + self.out_proj = nn.Linear(args.hidden_size, args.hidden_size, bias=self.bias) + + def __call__( + self, + x: mx.array, + mask: Optional[mx.array] = None, + cache: Optional[Any] = None, + ): + BCx = self.in_proj(x) + B, C, x = mx.split(BCx, 3, axis=-1) + Bx = B * x + if mask is not None: + Bx = mx.where(mask[..., None], Bx, 0) + + if cache is not None: + if cache[0] is None: + state = mx.zeros( + (Bx.shape[0], self.L_cache - 1, self.args.hidden_size), + dtype=Bx.dtype, + ) + else: + state = cache[0] + Bx = mx.concatenate([state, Bx], axis=1) + n_keep = self.L_cache - 1 + t = x.shape[1] + if cache.lengths is not None: + ends = mx.clip(cache.lengths, 0, t) + positions = (ends[:, None] + mx.arange(n_keep))[..., None] + cache[0] = mx.take_along_axis(Bx, positions, axis=1) + else: + cache[0] = Bx[:, -n_keep:, :] + cache.advance(t) + else: + Bx = mx.pad(Bx, [(0, 0), (self.L_cache - 1, 0), (0, 0)]) + + conv_out = self.conv(Bx) + + y = C * conv_out + return self.out_proj(y) + + +class MLP(nn.Module): + def __init__( + self, + dim: int, + ff_dim: int, + multiple_of: int, + auto_adjust_ff_dim: bool, + ffn_dim_multiplier: Optional[float], + ): + super().__init__() + if auto_adjust_ff_dim: + ff_dim = int(2 * ff_dim / 3) + if ffn_dim_multiplier is not None: + ff_dim = int(ffn_dim_multiplier * ff_dim) + ff_dim = multiple_of * ((ff_dim + multiple_of - 1) // multiple_of) + + self.w1 = nn.Linear(dim, ff_dim, bias=False) + self.w3 = nn.Linear(dim, ff_dim, bias=False) + self.w2 = nn.Linear(ff_dim, dim, bias=False) + + def __call__(self, x) -> mx.array: + return self.w2(swiglu(self.w1(x), self.w3(x))) + + +class Lfm2DecoderLayer(nn.Module): + def __init__(self, args: ModelArgs, layer_idx: int): + super().__init__() + self.is_attention_layer = layer_idx in args.full_attn_idxs + + if self.is_attention_layer: + self.self_attn = Attention(args) + else: + self.conv = ShortConv(args, layer_idx) + self.feed_forward = MLP( + dim=args.block_dim, + ff_dim=args.block_ff_dim, + multiple_of=args.block_multiple_of, + auto_adjust_ff_dim=args.block_auto_adjust_ff_dim, + ffn_dim_multiplier=args.block_ffn_dim_multiplier, + ) + + self.operator_norm = nn.RMSNorm(args.hidden_size, eps=args.norm_eps) + self.ffn_norm = nn.RMSNorm(args.hidden_size, eps=args.norm_eps) + + def __call__( + self, + x: mx.array, + mask: Optional[mx.array] = None, + cache: Optional[Any] = None, + ) -> mx.array: + + if self.is_attention_layer: + r = self.self_attn(self.operator_norm(x), mask=mask, cache=cache) + else: + r = self.conv( + self.operator_norm(x), + mask=mask, + cache=cache, + ) + h = x + r + out = h + self.feed_forward(self.ffn_norm(h)) + return out + + +class Lfm2Model(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.args = args + self.vocab_size = args.vocab_size + self.num_hidden_layers = args.num_hidden_layers + self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size) + self.layers = [ + Lfm2DecoderLayer(args, layer_idx=i) for i in range(args.num_hidden_layers) + ] + + self.embedding_norm = nn.RMSNorm(args.hidden_size, eps=args.norm_eps) + + self.fa_idx = args.full_attn_idxs[0] + self.conv_idx = 0 + for i in range(args.num_hidden_layers): + if i in args.full_attn_idxs: + self.conv_idx += 1 + else: + break + + def __call__( + self, + inputs: mx.array, + cache=None, + input_embeddings: Optional[mx.array] = None, + ): + if input_embeddings is not None: + h = input_embeddings + else: + h = self.embed_tokens(inputs) + + if cache is None: + cache = [None] * len(self.layers) + + attn_mask = create_attention_mask(h, cache[self.fa_idx]) + conv_mask = create_ssm_mask(h, cache[self.conv_idx]) + + for layer, c in zip(self.layers, cache): + mask = attn_mask if layer.is_attention_layer else conv_mask + h = layer(h, mask, cache=c) + + return self.embedding_norm(h) + + +class Model(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.args = args + self.model_type = args.model_type + self.model = Lfm2Model(args) + + def __call__( + self, + inputs: mx.array, + cache=None, + input_embeddings: Optional[mx.array] = None, + ): + out = self.model(inputs, cache, input_embeddings) + return self.model.embed_tokens.as_linear(out) + + def sanitize(self, weights): + sanitized_weights = {} + for name, param in weights.items(): + if "conv.weight" in name: + if param.shape[-1] > param.shape[1]: + param = param.transpose(0, 2, 1) + + sanitized_weights[name] = param + return sanitized_weights + + @property + def layers(self): + return self.model.layers + + def make_cache(self): + return [ + KVCache() if l.is_attention_layer else ArraysCache(size=1) + for l in self.layers + ] diff --git a/mlx_audio/lm/models/llama.py b/mlx_audio/lm/models/llama.py new file mode 100644 index 000000000..916de5754 --- /dev/null +++ b/mlx_audio/lm/models/llama.py @@ -0,0 +1,276 @@ +# Copyright © 2023-2024 Apple Inc. +# Vendored verbatim from mlx-lm v0.31.3 (ed1fca4cef15a824c5f1702c80f70b4cffc8e4dd), +# mlx_lm/models/llama.py. MIT licensed. + +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Union + +import mlx.core as mx +import mlx.nn as nn +from mlx.nn.layers.distributed import shard_linear + +from .activations import swiglu +from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention +from .cache import KVCache, RotatingKVCache +from .rope_utils import initialize_rope + + +@dataclass +class ModelArgs(BaseModelArgs): + model_type: str + hidden_size: int + num_hidden_layers: int + intermediate_size: int + num_attention_heads: int + rms_norm_eps: float + vocab_size: int + head_dim: Optional[int] = None + max_position_embeddings: Optional[int] = None + num_key_value_heads: Optional[int] = None + attention_bias: bool = False + mlp_bias: bool = False + rope_theta: float = 10000 + rope_traditional: bool = False + rope_scaling: Optional[Dict[str, Union[float, str]]] = None + tie_word_embeddings: bool = True + layer_types: Optional[List[str]] = None + sliding_window: Optional[int] = None + + def __post_init__(self): + if self.num_key_value_heads is None: + self.num_key_value_heads = self.num_attention_heads + + if self.layer_types is None: + self.layer_types = ["full_attention"] * self.num_hidden_layers + + +class Attention(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + + dim = args.hidden_size + self.n_heads = n_heads = args.num_attention_heads + self.n_kv_heads = n_kv_heads = args.num_key_value_heads + + self.head_dim = head_dim = args.head_dim or args.hidden_size // n_heads + + self.scale = head_dim**-0.5 + if hasattr(args, "attention_bias"): + attention_bias = args.attention_bias + else: + attention_bias = False + + self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=attention_bias) + self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=attention_bias) + self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=attention_bias) + self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=attention_bias) + + self.rope = initialize_rope( + self.head_dim, + args.rope_theta, + args.rope_traditional, + args.rope_scaling, + args.max_position_embeddings, + ) + + def __call__( + self, + x: mx.array, + mask: Optional[mx.array] = None, + cache: Optional[Any] = None, + ) -> mx.array: + B, L, D = x.shape + + queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x) + + # Prepare the queries, keys and values for the attention computation + queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3) + keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3) + values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3) + + if cache is not None: + queries = self.rope(queries, offset=cache.offset) + keys = self.rope(keys, offset=cache.offset) + keys, values = cache.update_and_fetch(keys, values) + else: + queries = self.rope(queries) + keys = self.rope(keys) + + output = scaled_dot_product_attention( + queries, keys, values, cache=cache, scale=self.scale, mask=mask + ) + + output = output.transpose(0, 2, 1, 3).reshape(B, L, -1) + return self.o_proj(output) + + +class MLP(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + + dim = args.hidden_size + hidden_dim = args.intermediate_size + if hasattr(args, "mlp_bias"): + mlp_bias = args.mlp_bias + else: + mlp_bias = False + + self.gate_proj = nn.Linear(dim, hidden_dim, bias=mlp_bias) + self.down_proj = nn.Linear(hidden_dim, dim, bias=mlp_bias) + self.up_proj = nn.Linear(dim, hidden_dim, bias=mlp_bias) + + def __call__(self, x) -> mx.array: + return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x))) + + +class TransformerBlock(nn.Module): + def __init__(self, args: ModelArgs, use_sliding: bool = False): + super().__init__() + self.num_attention_heads = args.num_attention_heads + self.hidden_size = args.hidden_size + self.use_sliding = use_sliding + self.self_attn = Attention(args) + self.mlp = MLP(args) + self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps) + self.post_attention_layernorm = nn.RMSNorm( + args.hidden_size, eps=args.rms_norm_eps + ) + self.args = args + + def __call__( + self, + x: mx.array, + mask: Optional[mx.array] = None, + cache: Optional[Any] = None, + ) -> mx.array: + r = self.self_attn(self.input_layernorm(x), mask, cache) + h = x + r + r = self.mlp(self.post_attention_layernorm(h)) + out = h + r + return out + + +class LlamaModel(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.args = args + self.vocab_size = args.vocab_size + self.num_hidden_layers = args.num_hidden_layers + self.layer_types = args.layer_types + self.sliding_window = args.sliding_window + assert self.vocab_size > 0 + self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size) + self.layers = [ + TransformerBlock(args=args, use_sliding=layer_type == "sliding_attention") + for layer_type in self.layer_types + ] + self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps) + self.fa_idx = self.layer_types.index("full_attention") + self.swa_idx = None + for e, l in enumerate(self.layers): + if l.use_sliding: + self.swa_idx = e + break + + def __call__( + self, + inputs: mx.array, + cache=None, + input_embeddings: Optional[mx.array] = None, + ): + if input_embeddings is not None: + h = input_embeddings + else: + h = self.embed_tokens(inputs) + + if cache is None: + cache = [None] * len(self.layers) + + fa_mask = create_attention_mask(h, cache[self.fa_idx]) + if self.swa_idx is not None: + swa_mask = create_attention_mask( + h, cache[self.swa_idx], window_size=self.sliding_window + ) + + for layer, cache in zip(self.layers, cache): + mask = swa_mask if layer.use_sliding else fa_mask + h = layer(h, mask, cache=cache) + + return self.norm(h) + + +class Model(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.args = args + self.model_type = args.model_type + self.model = LlamaModel(args) + if not args.tie_word_embeddings: + self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False) + + def __call__( + self, + inputs: mx.array, + cache=None, + input_embeddings: Optional[mx.array] = None, + ): + out = self.model(inputs, cache, input_embeddings) + if self.args.tie_word_embeddings: + out = self.model.embed_tokens.as_linear(out) + else: + out = self.lm_head(out) + return out + + def sanitize(self, weights): + # Remove unused precomputed rotary freqs + weights = { + k: v for k, v in weights.items() if "self_attn.rotary_emb.inv_freq" not in k + } + if self.args.tie_word_embeddings: + weights.pop("lm_head.weight", None) + return weights + + def shard(self, group: Optional[mx.distributed.Group] = None): + group = group or mx.distributed.init() + N = group.size() + for layer in self.model.layers: + # Shard the self attention + layer.self_attn.q_proj = shard_linear( + layer.self_attn.q_proj, "all-to-sharded", group=group + ) + layer.self_attn.k_proj = shard_linear( + layer.self_attn.k_proj, "all-to-sharded", group=group + ) + layer.self_attn.v_proj = shard_linear( + layer.self_attn.v_proj, "all-to-sharded", group=group + ) + layer.self_attn.o_proj = shard_linear( + layer.self_attn.o_proj, "sharded-to-all", group=group + ) + layer.self_attn.n_heads //= N + layer.self_attn.n_kv_heads //= N + + # Shard the MLP + layer.mlp.gate_proj = shard_linear( + layer.mlp.gate_proj, "all-to-sharded", group=group + ) + layer.mlp.down_proj = shard_linear( + layer.mlp.down_proj, "sharded-to-all", group=group + ) + layer.mlp.up_proj = shard_linear( + layer.mlp.up_proj, "all-to-sharded", group=group + ) + + @property + def layers(self): + return self.model.layers + + def make_cache(self): + return [ + ( + RotatingKVCache(max_size=self.model.sliding_window) + if layer.use_sliding + else KVCache() + ) + for layer in self.layers + ] diff --git a/mlx_audio/lm/models/qwen2.py b/mlx_audio/lm/models/qwen2.py new file mode 100644 index 000000000..de4b3ed31 --- /dev/null +++ b/mlx_audio/lm/models/qwen2.py @@ -0,0 +1,223 @@ +# Copyright © 2023-2024 Apple Inc. +# Vendored verbatim from mlx-lm v0.31.3 (ed1fca4cef15a824c5f1702c80f70b4cffc8e4dd), +# mlx_lm/models/qwen2.py. MIT licensed. + +from dataclasses import dataclass +from typing import Any, Dict, Optional, Union + +import mlx.core as mx +import mlx.nn as nn +from mlx.nn.layers.distributed import shard_linear + +from .activations import swiglu +from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention +from .rope_utils import initialize_rope + + +@dataclass +class ModelArgs(BaseModelArgs): + model_type: str + hidden_size: int + num_hidden_layers: int + intermediate_size: int + num_attention_heads: int + rms_norm_eps: float + vocab_size: int + num_key_value_heads: int + max_position_embeddings: int = 32768 + rope_theta: float = 1000000 + rope_traditional: bool = False + rope_scaling: Optional[Dict[str, Union[float, str]]] = None + tie_word_embeddings: bool = True + + +class Attention(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + + dim = args.hidden_size + self.n_heads = n_heads = args.num_attention_heads + assert args.num_key_value_heads is not None + self.n_kv_heads = n_kv_heads = args.num_key_value_heads + + head_dim = args.hidden_size // n_heads + self.scale = head_dim**-0.5 + + self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=True) + self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=True) + self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=True) + self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=False) + + self.rope = initialize_rope( + head_dim, + base=args.rope_theta, + traditional=args.rope_traditional, + scaling_config=args.rope_scaling, + max_position_embeddings=args.max_position_embeddings, + ) + + def __call__( + self, + x: mx.array, + mask: Optional[mx.array] = None, + cache: Optional[Any] = None, + ) -> mx.array: + B, L, D = x.shape + + queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x) + + # Prepare the queries, keys and values for the attention computation + queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3) + keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3) + values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3) + + if cache is not None: + queries = self.rope(queries, offset=cache.offset) + keys = self.rope(keys, offset=cache.offset) + keys, values = cache.update_and_fetch(keys, values) + else: + queries = self.rope(queries) + keys = self.rope(keys) + + output = scaled_dot_product_attention( + queries, keys, values, cache=cache, scale=self.scale, mask=mask + ) + output = output.transpose(0, 2, 1, 3).reshape(B, L, -1) + return self.o_proj(output) + + +class MLP(nn.Module): + def __init__(self, dim, hidden_dim): + super().__init__() + self.gate_proj = nn.Linear(dim, hidden_dim, bias=False) + self.down_proj = nn.Linear(hidden_dim, dim, bias=False) + self.up_proj = nn.Linear(dim, hidden_dim, bias=False) + + def __call__(self, x) -> mx.array: + return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x))) + + +class TransformerBlock(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.num_attention_heads = args.num_attention_heads + self.hidden_size = args.hidden_size + self.self_attn = Attention(args) + self.mlp = MLP(args.hidden_size, args.intermediate_size) + self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps) + self.post_attention_layernorm = nn.RMSNorm( + args.hidden_size, eps=args.rms_norm_eps + ) + self.args = args + + def __call__( + self, + x: mx.array, + mask: Optional[mx.array] = None, + cache: Optional[Any] = None, + ) -> mx.array: + r = self.self_attn(self.input_layernorm(x), mask, cache) + h = x + r + r = self.mlp(self.post_attention_layernorm(h)) + out = h + r + return out + + +class Qwen2Model(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.args = args + self.vocab_size = args.vocab_size + self.num_hidden_layers = args.num_hidden_layers + assert self.vocab_size > 0 + self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size) + self.layers = [ + TransformerBlock(args=args) for _ in range(args.num_hidden_layers) + ] + self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps) + + def __call__( + self, + inputs: mx.array, + cache=None, + input_embeddings: Optional[mx.array] = None, + ): + if input_embeddings is not None: + h = input_embeddings + else: + h = self.embed_tokens(inputs) + + if cache is None: + cache = [None] * len(self.layers) + mask = create_attention_mask(h, cache[0]) + + for layer, c in zip(self.layers, cache): + h = layer(h, mask, c) + + return self.norm(h) + + +class Model(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.args = args + self.model_type = args.model_type + self.model = Qwen2Model(args) + if not args.tie_word_embeddings: + self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False) + + def __call__( + self, + inputs: mx.array, + cache=None, + input_embeddings: Optional[mx.array] = None, + ): + out = self.model(inputs, cache, input_embeddings) + if self.args.tie_word_embeddings: + out = self.model.embed_tokens.as_linear(out) + else: + out = self.lm_head(out) + return out + + def sanitize(self, weights): + if self.args.tie_word_embeddings: + weights.pop("lm_head.weight", None) + # Remove unused precomputed rotary freqs + return { + k: v for k, v in weights.items() if "self_attn.rotary_emb.inv_freq" not in k + } + + def shard(self, group: Optional[mx.distributed.Group] = None): + group = group or mx.distributed.init() + N = group.size() + for layer in self.model.layers: + # Shard the self attention + layer.self_attn.q_proj = shard_linear( + layer.self_attn.q_proj, "all-to-sharded", group=group + ) + layer.self_attn.k_proj = shard_linear( + layer.self_attn.k_proj, "all-to-sharded", group=group + ) + layer.self_attn.v_proj = shard_linear( + layer.self_attn.v_proj, "all-to-sharded", group=group + ) + layer.self_attn.o_proj = shard_linear( + layer.self_attn.o_proj, "sharded-to-all", group=group + ) + layer.self_attn.n_heads //= N + layer.self_attn.n_kv_heads //= N + + # Shard the MLP + layer.mlp.gate_proj = shard_linear( + layer.mlp.gate_proj, "all-to-sharded", group=group + ) + layer.mlp.down_proj = shard_linear( + layer.mlp.down_proj, "sharded-to-all", group=group + ) + layer.mlp.up_proj = shard_linear( + layer.mlp.up_proj, "all-to-sharded", group=group + ) + + @property + def layers(self): + return self.model.layers diff --git a/mlx_audio/lm/models/qwen3.py b/mlx_audio/lm/models/qwen3.py new file mode 100644 index 000000000..692cc43ce --- /dev/null +++ b/mlx_audio/lm/models/qwen3.py @@ -0,0 +1,225 @@ +# Copyright © 2023-2024 Apple Inc. +# Vendored verbatim from mlx-lm v0.31.3 (ed1fca4cef15a824c5f1702c80f70b4cffc8e4dd), +# mlx_lm/models/qwen3.py. MIT licensed. + +from dataclasses import dataclass +from typing import Any, Dict, Optional, Union + +import mlx.core as mx +import mlx.nn as nn +from mlx.nn.layers.distributed import shard_linear + +from .activations import swiglu +from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention +from .rope_utils import initialize_rope + + +@dataclass +class ModelArgs(BaseModelArgs): + model_type: str + hidden_size: int + num_hidden_layers: int + intermediate_size: int + num_attention_heads: int + rms_norm_eps: float + vocab_size: int + num_key_value_heads: int + max_position_embeddings: int + rope_theta: float + head_dim: int + tie_word_embeddings: bool + rope_scaling: Optional[Dict[str, Union[float, str]]] = None + + +class Attention(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + + dim = args.hidden_size + self.n_heads = n_heads = args.num_attention_heads + assert args.num_key_value_heads is not None + self.n_kv_heads = n_kv_heads = args.num_key_value_heads + + head_dim = args.head_dim + self.scale = head_dim**-0.5 + + self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=False) + self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False) + self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False) + self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=False) + + self.q_norm = nn.RMSNorm(head_dim, eps=args.rms_norm_eps) + self.k_norm = nn.RMSNorm(head_dim, eps=args.rms_norm_eps) + self.rope = initialize_rope( + head_dim, + base=args.rope_theta, + traditional=False, + scaling_config=args.rope_scaling, + max_position_embeddings=args.max_position_embeddings, + ) + + def __call__( + self, + x: mx.array, + mask: Optional[mx.array] = None, + cache: Optional[Any] = None, + ) -> mx.array: + B, L, D = x.shape + + queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x) + + queries = self.q_norm(queries.reshape(B, L, self.n_heads, -1)).transpose( + 0, 2, 1, 3 + ) + keys = self.k_norm(keys.reshape(B, L, self.n_kv_heads, -1)).transpose( + 0, 2, 1, 3 + ) + values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3) + + if cache is not None: + queries = self.rope(queries, offset=cache.offset) + keys = self.rope(keys, offset=cache.offset) + keys, values = cache.update_and_fetch(keys, values) + else: + queries = self.rope(queries) + keys = self.rope(keys) + + output = scaled_dot_product_attention( + queries, keys, values, cache=cache, scale=self.scale, mask=mask + ) + output = output.transpose(0, 2, 1, 3).reshape(B, L, -1) + return self.o_proj(output) + + +class MLP(nn.Module): + def __init__(self, dim, hidden_dim): + super().__init__() + self.gate_proj = nn.Linear(dim, hidden_dim, bias=False) + self.down_proj = nn.Linear(hidden_dim, dim, bias=False) + self.up_proj = nn.Linear(dim, hidden_dim, bias=False) + + def __call__(self, x) -> mx.array: + return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x))) + + +class TransformerBlock(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.num_attention_heads = args.num_attention_heads + self.hidden_size = args.hidden_size + self.self_attn = Attention(args) + self.mlp = MLP(args.hidden_size, args.intermediate_size) + self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps) + self.post_attention_layernorm = nn.RMSNorm( + args.hidden_size, eps=args.rms_norm_eps + ) + self.args = args + + def __call__( + self, + x: mx.array, + mask: Optional[mx.array] = None, + cache: Optional[Any] = None, + ) -> mx.array: + r = self.self_attn(self.input_layernorm(x), mask, cache) + h = x + r + r = self.mlp(self.post_attention_layernorm(h)) + out = h + r + return out + + +class Qwen3Model(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.args = args + self.vocab_size = args.vocab_size + self.num_hidden_layers = args.num_hidden_layers + assert self.vocab_size > 0 + self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size) + self.layers = [ + TransformerBlock(args=args) for _ in range(args.num_hidden_layers) + ] + self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps) + + def __call__( + self, + inputs: mx.array, + cache=None, + input_embeddings: Optional[mx.array] = None, + ): + if input_embeddings is not None: + h = input_embeddings + else: + h = self.embed_tokens(inputs) + + if cache is None: + cache = [None] * len(self.layers) + mask = create_attention_mask(h, cache[0]) + + for layer, c in zip(self.layers, cache): + h = layer(h, mask, c) + + return self.norm(h) + + +class Model(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.args = args + self.model_type = args.model_type + self.model = Qwen3Model(args) + if not args.tie_word_embeddings: + self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False) + + def __call__( + self, + inputs: mx.array, + cache=None, + input_embeddings: Optional[mx.array] = None, + ): + out = self.model(inputs, cache, input_embeddings) + if self.args.tie_word_embeddings: + out = self.model.embed_tokens.as_linear(out) + else: + out = self.lm_head(out) + return out + + def sanitize(self, weights): + if self.args.tie_word_embeddings: + weights.pop("lm_head.weight", None) + return weights + + def shard(self, group: Optional[mx.distributed.Group] = None): + group = group or mx.distributed.init() + N = group.size() + for layer in self.model.layers: + # Shard the self attention + layer.self_attn.q_proj = shard_linear( + layer.self_attn.q_proj, "all-to-sharded", group=group + ) + layer.self_attn.k_proj = shard_linear( + layer.self_attn.k_proj, "all-to-sharded", group=group + ) + layer.self_attn.v_proj = shard_linear( + layer.self_attn.v_proj, "all-to-sharded", group=group + ) + layer.self_attn.o_proj = shard_linear( + layer.self_attn.o_proj, "sharded-to-all", group=group + ) + layer.self_attn.n_heads //= N + layer.self_attn.n_kv_heads //= N + + # Shard the MLP + layer.mlp.gate_proj = shard_linear( + layer.mlp.gate_proj, "all-to-sharded", group=group + ) + layer.mlp.down_proj = shard_linear( + layer.mlp.down_proj, "sharded-to-all", group=group + ) + layer.mlp.up_proj = shard_linear( + layer.mlp.up_proj, "all-to-sharded", group=group + ) + + @property + def layers(self): + return self.model.layers diff --git a/mlx_audio/lm/models/rope_utils.py b/mlx_audio/lm/models/rope_utils.py new file mode 100644 index 000000000..aedd5f62c --- /dev/null +++ b/mlx_audio/lm/models/rope_utils.py @@ -0,0 +1,310 @@ +# Copyright © 2023-2024 Apple Inc. +# Vendored verbatim from mlx-lm v0.31.3 (ed1fca4cef15a824c5f1702c80f70b4cffc8e4dd), +# mlx_lm/models/rope_utils.py. MIT licensed. + +import math +from typing import List, Optional, Union + +import mlx.core as mx +import mlx.nn as nn + + +class SuScaledRoPE(nn.Module): + def __init__( + self, + dims: int, + base: float = 10000.0, + max_position_embeddings: int = 131072, + original_max_position_embeddings: int = 4096, + short_factor: Union[List[float], float] = 1.0, + long_factor: Union[List[float], float] = 1.0, + short_mscale: float = None, + long_mscale: float = None, + ): + """ + Su Scaled Rotary Embedding layer. + + Args: + dims (int): The feature dimensions to be rotated. + base (int, optional): Base for the exponential scaling. + max_position_embeddings (int, optional): The maximum sequence + length that this model was trained with. This is used to determine + the size of the original RoPE embeddings when using long scaling. + Default: ``131072``. + original_max_position_embeddings (int, optional): The maximum + sequence length that this model was trained with. This is used to + determine the size of the original RoPE embeddings when using long + scaling. Default: ``4096``. + short_factor (float or list[float], optional): List of scaling + factors for sequences of length lesser than + ``original_max_position_embeddings``. Default: ``1.0``. + long_factor (float or list[float], optional): List of scaling + factors for sequences of length greater than + ``original_max_position_embeddings``. Default: ``1.0``. + short_mscale (float, optional): Scale the input prior to embedding. + long_mscale (float, optional): Scale the input prior to embedding. + """ + super().__init__() + self.original_max_position_embeddings = original_max_position_embeddings + self.dim = dims + + freqs = base ** (mx.arange(0, dims, 2, dtype=mx.float32) / dims) + self._freqs = mx.array(long_factor, dtype=mx.float32) * freqs + + def default_scale(factor): + return math.sqrt( + 1 + math.log(factor) / math.log(original_max_position_embeddings) + ) + + factor = max_position_embeddings / original_max_position_embeddings + self._scale = long_mscale or (1.0 if factor <= 1.0 else default_scale(factor)) + + def __call__(self, x, offset: Union[int, mx.array] = 0): + x = x[...] + x[..., : self.dim] = self._scale * x[..., : self.dim] + return mx.fast.rope( + x, + self.dim, + traditional=False, + base=None, + scale=1.0, + offset=offset, + freqs=self._freqs, + ) + + +class Llama3RoPE(nn.Module): + def __init__( + self, + dims: int, + max_position_embeddings: int = 2048, + traditional: bool = False, + base: float = 10000, + scaling_config: dict = None, + ): + super().__init__() + self.dims = dims + self.max_position_embeddings = max_position_embeddings + self.traditional = traditional + + factor = scaling_config["factor"] + low_freq_factor = scaling_config.get("low_freq_factor", 1.0) + high_freq_factor = scaling_config.get("high_freq_factor", 4.0) + old_context_len = scaling_config.get( + "original_max_position_embeddings", + 8192, + ) + + low_freq_wavelen = old_context_len / low_freq_factor + high_freq_wavelen = old_context_len / high_freq_factor + + freqs = base ** (mx.arange(0, dims, 2) / dims) + wavelens = 2 * mx.pi * freqs + + freqs = mx.where(wavelens > low_freq_wavelen, freqs * factor, freqs) + is_medium_freq = (wavelens > high_freq_wavelen) & (wavelens < low_freq_wavelen) + smooth_factors = (old_context_len / wavelens - low_freq_factor) / ( + high_freq_factor - low_freq_factor + ) + smooth_freqs = freqs / ((1 - smooth_factors) / factor + smooth_factors) + self._freqs = mx.where(is_medium_freq, smooth_freqs, freqs) + + def extra_repr(self): + return ( + f"{self.dims}, traditional={self.traditional}, " + f"max_position_embeddings={self.max_position_embeddings}" + ) + + def __call__(self, x, offset: int = 0): + return mx.fast.rope( + x, + self.dims, + traditional=self.traditional, + base=None, + scale=1.0, + offset=offset, + freqs=self._freqs, + ) + + +class YarnRoPE(nn.Module): + def __init__( + self, + dims, + traditional=False, + max_position_embeddings=2048, + base=10000, + scaling_factor=1.0, + original_max_position_embeddings=4096, + beta_fast=32, + beta_slow=1, + mscale=1, + mscale_all_dim=0, + ): + super().__init__() + + def yarn_find_correction_dim(num_rotations): + return ( + dims + * math.log( + original_max_position_embeddings / (num_rotations * 2 * math.pi) + ) + ) / (2 * math.log(base)) + + def yarn_find_correction_range(): + low = math.floor(yarn_find_correction_dim(beta_fast)) + high = math.ceil(yarn_find_correction_dim(beta_slow)) + return max(low, 0), min(high, dims - 1) + + def yarn_get_mscale(scale=1, mscale=1): + if scale <= 1: + return 1.0 + return 0.1 * mscale * math.log(scale) + 1.0 + + def yarn_linear_ramp_mask(min_val, max_val, dim): + if min_val == max_val: + max_val += 0.001 # Prevent singularity + + linear_func = (mx.arange(dim, dtype=mx.float32) - min_val) / ( + max_val - min_val + ) + return mx.clip(linear_func, 0, 1) + + self.mscale = yarn_get_mscale(scaling_factor, mscale) / yarn_get_mscale( + scaling_factor, mscale_all_dim + ) + freq_extra = base ** (mx.arange(0, dims, 2, dtype=mx.float32) / dims) + freq_inter = scaling_factor * freq_extra + low, high = yarn_find_correction_range() + freq_mask = 1.0 - yarn_linear_ramp_mask(low, high, dims // 2) + self._freqs = (freq_inter * freq_extra) / ( + freq_inter * freq_mask + freq_extra * (1 - freq_mask) + ) + self.dims = dims + self.traditional = traditional + + def __call__(self, x, offset=0): + if self.mscale != 1.0: + x = x[...] + x[..., : self.dims] = self.mscale * x[..., : self.dims] + return mx.fast.rope( + x, + self.dims, + traditional=self.traditional, + base=None, + scale=1.0, + offset=offset, + freqs=self._freqs, + ) + + +class ProportionalRoPE(nn.Module): + def __init__( + self, + dims: int, + rotated_dims: int, + traditional: bool = False, + base: float = 10000.0, + factor: float = 1.0, + ): + super().__init__() + self.dims = dims + self.traditional = traditional + + if rotated_dims > dims: + raise ValueError("rotated_dims should be smaller than dims") + + exponents = mx.arange(0, rotated_dims, 2, dtype=mx.float32) / dims + self._freqs = mx.concatenate( + [ + factor * (base**exponents), + mx.full(((dims - rotated_dims) // 2,), mx.inf), + ] + ) + + def __call__(self, x, offset=0): + return mx.fast.rope( + x, + self.dims, + traditional=self.traditional, + base=None, + scale=1.0, + offset=offset, + freqs=self._freqs, + ) + + +def initialize_rope( + dims, + base, + traditional, + scaling_config: Optional[dict] = None, + max_position_embeddings: Optional[int] = None, +): + if scaling_config is not None: + rope_type = scaling_config.get("type") or scaling_config.get( + "rope_type", "default" + ) + else: + rope_type = "default" + + if rope_type in ["default", "linear"]: + scale = 1 / scaling_config["factor"] if rope_type == "linear" else 1.0 + return nn.RoPE(dims, traditional=traditional, base=base, scale=scale) + + elif rope_type == "llama3": + return Llama3RoPE( + dims=dims, + max_position_embeddings=max_position_embeddings, + traditional=traditional, + base=base, + scaling_config=scaling_config, + ) + elif rope_type in ("yarn", "deepseek_yarn", "telechat3-yarn"): + scaling_factor = scaling_config["factor"] + rope_kwargs = { + key: scaling_config[key] + for key in [ + "original_max_position_embeddings", + "beta_fast", + "beta_slow", + "mscale", + "mscale_all_dim", + ] + if key in scaling_config + } + return YarnRoPE( + dims=dims, + max_position_embeddings=max_position_embeddings, + traditional=traditional, + scaling_factor=scaling_factor, + base=base, + **rope_kwargs, + ) + elif rope_type == "longrope": + return SuScaledRoPE( + dims=dims, + base=base, + max_position_embeddings=max_position_embeddings, + original_max_position_embeddings=scaling_config[ + "original_max_position_embeddings" + ], + short_factor=scaling_config["short_factor"], + long_factor=scaling_config["long_factor"], + ) + elif rope_type == "proportional": + return ProportionalRoPE( + dims=dims, + rotated_dims=int(dims * scaling_config.get("partial_rotary_factor", 1.0)), + traditional=traditional, + base=base, + factor=scaling_config.get("factor", 1.0), + ) + elif rope_type == "mrope": + mrope_section = scaling_config.get("mrope_section", []) + assert ( + len(mrope_section) == 3 + ), f"MRoPE currently only supports 3 sections, got {len(mrope_section)}." + return nn.RoPE(dims, traditional=traditional, base=base) + else: + raise ValueError(f"Unsupported RoPE type {rope_type}") diff --git a/mlx_audio/lm/models/switch_layers.py b/mlx_audio/lm/models/switch_layers.py new file mode 100644 index 000000000..a0b85d6a2 --- /dev/null +++ b/mlx_audio/lm/models/switch_layers.py @@ -0,0 +1,238 @@ +# Copyright © 2023-2024 Apple Inc. +# Vendored verbatim from mlx-lm v0.31.3 (ed1fca4cef15a824c5f1702c80f70b4cffc8e4dd), +# mlx_lm/models/switch_layers.py. MIT licensed. + +import math +from functools import partial + +import mlx.core as mx +import mlx.nn as nn + +from .activations import swiglu + + +def _gather_sort(x, indices): + *_, M = indices.shape + indices = indices.flatten() + order = mx.argsort(indices) + inv_order = mx.argsort(order) + return x.flatten(0, -3)[order // M], indices[order], inv_order + + +def _scatter_unsort(x, inv_order, shape=None): + x = x[inv_order] + if shape is not None: + x = mx.unflatten(x, 0, shape) + return x + + +class QuantizedSwitchLinear(nn.Module): + def __init__( + self, + input_dims: int, + output_dims: int, + num_experts: int, + bias: bool = True, + group_size: int = 64, + bits: int = 4, + mode: str = "affine", + ): + super().__init__() + + scale = math.sqrt(1 / input_dims) + self.weight, self.scales, *biases = mx.quantize( + mx.random.uniform( + low=-scale, + high=scale, + shape=(num_experts, output_dims, input_dims), + ), + group_size=group_size, + bits=bits, + mode=mode, + ) + self.biases = biases[0] if biases else None + + if bias: + self.bias = mx.zeros((num_experts, output_dims)) + + self.group_size = group_size + self.bits = bits + self.mode = mode + + # Freeze this model's parameters + self.freeze() + + @property + def input_dims(self): + return self.scales.shape[2] * self.group_size + + @property + def output_dims(self): + return self.weight.shape[1] + + @property + def num_experts(self): + return self.weight.shape[0] + + def __call__(self, x, indices, sorted_indices=False): + x = mx.gather_qmm( + x, + self["weight"], + self["scales"], + self.get("biases"), + rhs_indices=indices, + transpose=True, + group_size=self.group_size, + bits=self.bits, + mode=self.mode, + sorted_indices=sorted_indices, + ) + if "bias" in self: + x = x + mx.expand_dims(self["bias"][indices], -2) + return x + + +class SwitchLinear(nn.Module): + def __init__( + self, input_dims: int, output_dims: int, num_experts: int, bias: bool = True + ): + super().__init__() + scale = math.sqrt(1 / input_dims) + self.weight = mx.random.uniform( + low=-scale, + high=scale, + shape=(num_experts, output_dims, input_dims), + ) + + if bias: + self.bias = mx.zeros((num_experts, output_dims)) + + @property + def input_dims(self): + return self.weight.shape[2] + + @property + def output_dims(self): + return self.weight.shape[1] + + @property + def num_experts(self): + return self.weight.shape[0] + + def __call__(self, x, indices, sorted_indices=False): + x = mx.gather_mm( + x, + self["weight"].swapaxes(-1, -2), + rhs_indices=indices, + sorted_indices=sorted_indices, + ) + if "bias" in self: + x = x + mx.expand_dims(self["bias"][indices], -2) + return x + + def to_quantized(self, group_size: int = 64, bits: int = 4, mode: str = "affine"): + num_experts, output_dims, input_dims = self.weight.shape + ql = QuantizedSwitchLinear( + input_dims, + output_dims, + num_experts, + False, + group_size, + bits, + mode=mode, + ) + ql.weight, ql.scales, *biases = mx.quantize( + self.weight, group_size, bits, mode=mode + ) + ql.biases = biases[0] if biases else None + + if "bias" in self: + ql.bias = self.bias + return ql + + +class SwiGLU(nn.Module): + def __init__(self): + super().__init__() + + def __call__(self, x, gate): + return swiglu(gate, x) + + +class SwitchGLU(nn.Module): + def __init__( + self, + input_dims: int, + hidden_dims: int, + num_experts: int, + activation=SwiGLU(), + bias: bool = False, + ): + super().__init__() + + self.gate_proj = SwitchLinear(input_dims, hidden_dims, num_experts, bias=bias) + self.up_proj = SwitchLinear(input_dims, hidden_dims, num_experts, bias=bias) + self.down_proj = SwitchLinear(hidden_dims, input_dims, num_experts, bias=bias) + self.activation = activation + + def __call__(self, x, indices) -> mx.array: + x = mx.expand_dims(x, (-2, -3)) + + # When we have many tokens, then sort them to make sure that the access + # of different experts is in order. + do_sort = indices.size >= 64 + idx = indices + inv_order = None + if do_sort: + x, idx, inv_order = _gather_sort(x, indices) + if self.training: + idx = mx.stop_gradient(idx) + x_up = self.up_proj(x, idx, sorted_indices=do_sort) + x_gate = self.gate_proj(x, idx, sorted_indices=do_sort) + x = self.down_proj( + self.activation(x_up, x_gate), + idx, + sorted_indices=do_sort, + ) + + if do_sort: + x = _scatter_unsort(x, inv_order, indices.shape) + + return x.squeeze(-2) + + +class SwitchMLP(nn.Module): + def __init__( + self, + input_dims: int, + hidden_dims: int, + num_experts: int, + activation=nn.GELU(approx="precise"), + bias: bool = False, + ): + super().__init__() + + self.fc1 = SwitchLinear(input_dims, hidden_dims, num_experts, bias=bias) + self.fc2 = SwitchLinear(hidden_dims, input_dims, num_experts, bias=bias) + self.activation = activation + + def __call__(self, x, indices) -> mx.array: + x = mx.expand_dims(x, (-2, -3)) + + # When we have many tokens, then sort them to make sure that the access + # of different experts is in order. + do_sort = indices.size >= 64 + idx = indices + inv_order = None + if do_sort: + x, idx, inv_order = _gather_sort(x, indices) + if self.training: + idx = mx.stop_gradient(idx) + x = self.fc1(x, idx, sorted_indices=do_sort) + x = self.activation(x) + x = self.fc2(x, idx, sorted_indices=do_sort) + + if do_sort: + x = _scatter_unsort(x, inv_order, indices.shape) + + return x.squeeze(-2) diff --git a/mlx_audio/lm/sample_utils.py b/mlx_audio/lm/sample_utils.py new file mode 100644 index 000000000..bd8b04207 --- /dev/null +++ b/mlx_audio/lm/sample_utils.py @@ -0,0 +1,371 @@ +# Copyright © 2023-2024 Apple Inc. +# Vendored from mlx-lm v0.31.3 (ed1fca4cef15a824c5f1702c80f70b4cffc8e4dd), +# mlx_lm/sample_utils.py. Modified: apply_min_p passes a bool array rather than +# a Python bool to mx.put_along_axis, which upstream raises a TypeError on +# whenever min_p > 0 and min_tokens_to_keep > 1. MIT licensed. + +import math +from functools import partial +from typing import Callable, Dict, List, Optional + +import mlx.core as mx + + +def make_sampler( + temp: float = 0.0, + top_p: float = 0.0, + min_p: float = 0.0, + min_tokens_to_keep: int = 1, + top_k: int = 0, + xtc_probability: float = 0.0, + xtc_threshold: float = 0.0, + xtc_special_tokens: List[int] = [], +) -> Callable[[mx.array], mx.array]: + """ + Make a sampler function for use with ``generate_step``. + + Args: + temp (float): The temperature for sampling, if 0 the argmax is used. + Default: ``0``. + top_p (float, optional): Nulceus sampling, higher means model considers + more less likely words. + min_p (float, optional): The minimum value (scaled by the top token's + probability) that a token probability must have to be considered. + min_tokens_to_keep (int, optional): Minimum number of tokens that cannot + be filtered by min_p sampling. + top_k (int, optional): The top k tokens ranked by probability to constrain + the sampling to. + xtc_probability (float, optional): The probability of applying XTC + sampling. + xtc_threshold (float, optional): The threshold the probs need to reach + for being sampled. + xtc_special_tokens (list(int), optional): List of special tokens IDs to + be excluded from XTC sampling. + + + Returns: + Callable[mx.array, mx.array]: + A sampler which takes log-probabilities and returns tokens. + """ + if temp == 0: + return lambda x: mx.argmax(x, axis=-1) + + # Create sampler chain + sampling_methods = [] + if top_p > 0 and top_p < 1.0: + sampling_methods.append(lambda x: apply_top_p(x, top_p)) + if min_p != 0.0: + sampling_methods.append(lambda x: apply_min_p(x, min_p, min_tokens_to_keep)) + if xtc_probability > 0.0: + sampling_methods.append( + lambda x: apply_xtc(x, xtc_probability, xtc_threshold, xtc_special_tokens) + ) + if top_k > 0: + sampling_methods.append(lambda x: apply_top_k(x, top_k)) + + # Apply the sampling methods + def sampler(logprobs): + for method in sampling_methods: + logprobs = method(logprobs) + # Return the sampled token + return categorical_sampling(logprobs, temp) + + return sampler + + +def make_logits_processors( + logit_bias: Optional[Dict[int, float]] = None, + repetition_penalty: Optional[float] = None, + repetition_context_size: Optional[int] = 20, + presence_penalty: Optional[float] = None, + presence_context_size: Optional[int] = 20, + frequency_penalty: Optional[float] = None, + frequency_context_size: Optional[int] = 20, +): + """ + Make logits processors for use with ``generate_step``. + + Args: + repetition_penalty (float, optional): A (sign-aware) multiplicative + penalty for repeating tokens. + repetition_context_size (int, optional): The number of tokens to + consider for repetition penalty. Default: ``20``. + presence_penalty (float, optional): An additive penalty to reduce + repeating tokens. + presence_context_size (int, optional): The number of tokens to consider + for the presence penalty. Default: ``20``. + frequency_penalty (float, optional): An additive penalty to reduce + repeating tokens. The tokens are penalized proportionally to their + frequency. + frequency_context_size (int, optional): The number of tokens to consider + for the frequency penalty. Default: ``20``. + logit_bias (dictionary, optional): Additive logit bias. + + Returns: + List[Callable[[mx.array, mx.array], mx.array]]: + A list of logits processors. Each processor in the list is a + callable which takes an array of tokens and an array of logits + and returns the updated logits. + """ + logits_processors = [] + if logit_bias: + indices = mx.array(list(logit_bias.keys())) + values = mx.array(list(logit_bias.values())) + + def logit_bias_processor(_, logits): + return logits.at[:, indices].add(values) + + logits_processors.append(logit_bias_processor) + + repetition_penalties = [ + (make_repetition_penalty, repetition_penalty, repetition_context_size), + (make_presence_penalty, presence_penalty, presence_context_size), + (make_frequency_penalty, frequency_penalty, frequency_context_size), + ] + + for make_penalty, penalty, context_size in repetition_penalties: + if penalty is not None and penalty != 0: + logits_processors.append(make_penalty(penalty, context_size)) + + return logits_processors + + +@partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state) +def apply_top_k( + logprobs: mx.array, + top_k: int, +) -> mx.array: + """ + Sample from only the top K tokens ranked by probability. + + Args: + logprobs: A vector of log probabilities. + top_k (int): Top k tokens to sample from. + """ + vocab_size = logprobs.shape[-1] + if not isinstance(top_k, int) or not (0 < top_k < vocab_size): + raise ValueError( + f"`top_k` has to be an integer in the (0, {vocab_size}] interval," + f" but is {top_k}." + ) + mask_idx = mx.argpartition(-logprobs, kth=top_k - 1, axis=-1)[..., top_k:] + masked_logprobs = mx.put_along_axis( + logprobs, mask_idx, mx.array(-float("inf"), logprobs.dtype), axis=-1 + ) + return masked_logprobs + + +@partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state) +def apply_min_p( + logprobs: mx.array, + min_p: float, + min_tokens_to_keep: int = 1, +) -> mx.array: + """ + Apply min-p sampling to the logprobs. + + Min-p keeps all tokens that are above a minimum probability, scaled by the + probability of the most likely token. As a result, the filter is more + aggressive given a very high-probability token. + + Args: + logprobs: A vector of log probabilities. + min_p (float): Minimum token probability. Typical values are in the + 0.01-0.2 range, comparably selective as setting `top_p` in the + 0.99-0.8 range. + min_tokens_to_keep (int, optional): Minimum number of tokens that cannot + be filtered. Default: ``1``. + + """ + if not (0 <= min_p <= 1.0): + raise ValueError( + f"`min_p` has to be a float in the [0, 1] interval, but is {min_p}" + ) + if not isinstance(min_tokens_to_keep, int) or (min_tokens_to_keep < 1): + raise ValueError( + f"`min_tokens_to_keep` has to be a positive integer, but is {min_tokens_to_keep}" + ) + + # Mask tokens that have a probability less than the max(p) * min_p + top_logprobs = mx.max(logprobs, axis=-1, keepdims=True) + scaled_min_p = top_logprobs + math.log(min_p) + tokens_to_remove = logprobs < scaled_min_p + + # Ensure at least min_tokens_to_keep survive the filter + if min_tokens_to_keep > 1: + top_indices = mx.argpartition(logprobs, kth=-min_tokens_to_keep, axis=-1) + top_indices = top_indices[..., -min_tokens_to_keep:] + tokens_to_remove = mx.put_along_axis( + tokens_to_remove, + top_indices, + mx.zeros(top_indices.shape, dtype=mx.bool_), + axis=-1, + ) + + return mx.where(tokens_to_remove, -float("inf"), logprobs) + + +@partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state) +def apply_top_p(logprobs: mx.array, top_p: float) -> mx.array: + """ + Apply top-p (nucleus) sampling to logits. + + Args: + logprobs: A vector of log probabilities. + top_p: The cumulative probability threshold for top-p filtering. + Returns: + token selected based on the top-p criterion. + """ + # referenced implementation from https://github.com/huggingface/transformers/blob/main/src/transformers/generation/logits_process.py#L449-L460 + probs = mx.exp(logprobs) + # sort in ascending order + sorted_indices = mx.argsort(logprobs, axis=-1) + sorted_probs = mx.take_along_axis(probs, sorted_indices, axis=-1) + + cumulative_probs = mx.cumsum(sorted_probs, axis=-1) + + # Rearrange cumulative probs back to original order + inverse_indices = mx.put_along_axis( + mx.zeros_like(sorted_indices), + sorted_indices, + mx.arange(sorted_indices.shape[-1], dtype=sorted_indices.dtype), + axis=-1, + ) + cumulative_probs = mx.take_along_axis(cumulative_probs, inverse_indices, axis=-1) + + # select tokens with cumulative probs below threshold + return mx.where( + cumulative_probs > 1 - top_p, + logprobs, + -float("inf"), + ) + + +@partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state) +def apply_xtc( + logits: mx.array, + xtc_probability: float, + xtc_threshold: float, + xtc_special_tokens: List[int], +) -> mx.array: + """ + Apply XTC sampling to the logits. + + Args: + logits: The logits from the model's output. + xtc_probability (float): Probability of XTC sampling to happen for each token + xtc_threshold (float): The threshold the probs need to reach for being sampled. + special_tokens_ids (list(int)): List of special tokens IDs to be excluded from XTC sampling. + """ + if not (0 <= xtc_threshold <= 0.5): + raise ValueError( + f"`threshold` has to be a float in the [0, 0.5] interval, but is {xtc_threshold}" + ) + if not (0 <= xtc_probability <= 1.0): + raise ValueError( + f"`probability` has to be a float in the [0, 1] interval, but is {xtc_probability}" + ) + + probs = mx.softmax(logits, -1) + mask = probs > mx.where(probs > xtc_threshold, probs, mx.inf).min() + if xtc_special_tokens: + mask[..., xtc_special_tokens] = False + + return mx.where( + mx.random.uniform(0, 1) > xtc_probability, + logits, + mx.where(mask, -mx.inf, logits), + ) + + +@partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state) +def categorical_sampling(logits, temp): + return mx.random.categorical(logits * (1 / temp)) + + +def make_repetition_penalty(penalty: float, context_size: int = 20): + """ + Make repetition penalty processor. + + Paper: https://arxiv.org/abs/1909.05858 + + Args: + penalty (float): The repetition penalty factor to be applied. + context_size (int): The number of previous tokens to use. + Default: ``20``. + + Returns: + Callable[[mx.array, List[int]], mx.array]: + The repetition penalty processor. + """ + if penalty < 0 or not isinstance(penalty, (int, float)): + raise ValueError(f"penalty must be a non-negative float, got {penalty}") + + def repetition_penalty_processor(tokens, logits): + if len(tokens) > 0: + tokens = tokens[-context_size:] + selected_logits = logits[:, tokens] + selected_logits = mx.where( + selected_logits < 0, + selected_logits * penalty, + selected_logits / penalty, + ) + logits[:, tokens] = selected_logits + return logits + + return repetition_penalty_processor + + +def make_presence_penalty(penalty: float, context_size: int = 20): + """ + Make a presence penalty processor. + + Corresponds to the OpenAI option with the same name. Namely, subtracts + ``penalty`` from a logit if the token has occured at least once in the + ``context_size`` previous tokens. + + Args: + penalty (float): The presence penalty to be applied. + context_size (int): The number of previous tokens to use. + Default: ``20``. + + Returns: + Callable[[mx.array, List[int]], mx.array] + """ + + def presence_penalty_processor(tokens, logits): + if len(tokens) > 0: + tokens = tokens[-context_size:] + logits[:, tokens] -= penalty + return logits + + return presence_penalty_processor + + +def make_frequency_penalty(penalty: float, context_size: int = 20): + """ + Make a frequency penalty processor. + + Corresponds to the OpenAI option with the same name. Namely, subtracts + ``penalty`` from a logit for every time that the token has occured in the + ``context_size`` previous tokens. + + The difference with the presence penalty is that the more often a token + occurs the more it will be penalized. + + Args: + penalty (float): The frequency penalty to be applied. + context_size (int): The number of previous tokens to use. + Default: ``20``. + + Returns: + Callable[[mx.array, List[int]], mx.array] + """ + + def frequency_penalty_processor(tokens, logits): + if len(tokens) > 0: + tokens = tokens[-context_size:] + logits = logits.at[:, tokens].subtract(penalty) + return logits + + return frequency_penalty_processor diff --git a/mlx_audio/sts/models/lfm_audio/config.py b/mlx_audio/sts/models/lfm_audio/config.py index 8c9a2a7f3..595128024 100644 --- a/mlx_audio/sts/models/lfm_audio/config.py +++ b/mlx_audio/sts/models/lfm_audio/config.py @@ -1,11 +1,10 @@ # Copyright (c) 2025 Prince Canuma and contributors (https://github.com/Blaizzy/mlx-audio) -from dataclasses import dataclass, field +from dataclasses import dataclass, field, fields from typing import Any, Dict, List, Optional -from mlx_lm.models.lfm2 import ModelArgs as LFM2Config - from mlx_audio.base import BaseModelArgs +from mlx_audio.lm.models.lfm2 import ModelArgs as LFM2Config @dataclass @@ -151,12 +150,13 @@ def from_dict(cls, config_dict: Dict[str, Any]) -> "LFM2AudioConfig": ) } + known = {f.name for f in fields(cls)} return cls( preprocessor=preprocessor, encoder=encoder, lfm=lfm, depthformer=depthformer, - **config_dict, + **{k: v for k, v in config_dict.items() if k in known}, ) def to_dict(self) -> Dict[str, Any]: diff --git a/mlx_audio/sts/models/lfm_audio/model.py b/mlx_audio/sts/models/lfm_audio/model.py index 05818b847..263866e94 100644 --- a/mlx_audio/sts/models/lfm_audio/model.py +++ b/mlx_audio/sts/models/lfm_audio/model.py @@ -12,8 +12,9 @@ import mlx.core as mx import mlx.nn as nn from huggingface_hub import snapshot_download -from mlx_lm.models.cache import ArraysCache, KVCache -from mlx_lm.models.lfm2 import Lfm2Model + +from mlx_audio.lm.models.cache import ArraysCache, KVCache +from mlx_audio.lm.models.lfm2 import Lfm2Model from ....base import check_array_shape from .config import DepthformerConfig, LFM2AudioConfig diff --git a/mlx_audio/sts/voice_pipeline.py b/mlx_audio/sts/voice_pipeline.py index 4dffc94af..ec63abf5a 100644 --- a/mlx_audio/sts/voice_pipeline.py +++ b/mlx_audio/sts/voice_pipeline.py @@ -10,8 +10,6 @@ import mlx.core as mx import numpy as np import sounddevice as sd -from mlx_lm.generate import generate as generate_text -from mlx_lm.utils import load as load_llm from mlx_audio.sts.audio_player import AudioPlayer from mlx_audio.tts.utils import load_model as load_tts @@ -399,6 +397,13 @@ def __init__(self, model_name: str, *, system_prompt: str): self.tokenizer = None def load(self) -> None: + try: + from mlx_lm.utils import load as load_llm + except ImportError as exc: + raise ImportError( + "The in-process LLM responder needs mlx-lm, please run `pip install -U mlx-lm` first. " + ) from exc + self.llm, self.tokenizer = load_llm(self.model_name) def generate(self, transcript: str, context: Optional[list[dict]] = None) -> str: @@ -411,6 +416,8 @@ def generate(self, transcript: str, context: Optional[list[dict]] = None) -> str prompt = self.tokenizer.apply_chat_template( messages, tokenize=False, enable_thinking=False, add_generation_prompt=True ) + from mlx_lm.generate import generate as generate_text + return generate_text(self.llm, self.tokenizer, prompt, verbose=False).strip() diff --git a/mlx_audio/stt/models/cohere_asr/cohere_asr.py b/mlx_audio/stt/models/cohere_asr/cohere_asr.py index 86daf7388..5f52522fe 100644 --- a/mlx_audio/stt/models/cohere_asr/cohere_asr.py +++ b/mlx_audio/stt/models/cohere_asr/cohere_asr.py @@ -6,8 +6,8 @@ import mlx.nn as nn import numpy as np from mlx.utils import tree_flatten -from mlx_lm.models.cache import KVCache +from mlx_audio.lm.models.cache import KVCache from mlx_audio.stt.models.base import STTOutput from .audio import CohereAudioFrontend diff --git a/mlx_audio/stt/models/fun_asr_nano/fun_asr_nano.py b/mlx_audio/stt/models/fun_asr_nano/fun_asr_nano.py index c5f01e88f..59246712f 100644 --- a/mlx_audio/stt/models/fun_asr_nano/fun_asr_nano.py +++ b/mlx_audio/stt/models/fun_asr_nano/fun_asr_nano.py @@ -364,7 +364,7 @@ def layers(self): return self.llm.model.layers def make_cache(self) -> List[Any]: - from mlx_lm.models.cache import KVCache + from mlx_audio.lm.models.cache import KVCache return [KVCache() for _ in range(self.config.text_config.num_hidden_layers)] @@ -496,7 +496,7 @@ def stream_generate( itn: bool = True, prefill_step_size: int = 2048, ): - from mlx_lm.generate import generate_step + from mlx_audio.lm.generate import generate_step hotwords = self._resolve_hotwords(hotwords, context) input_ids, inputs_embeds = self._build_inputs_embeds( @@ -538,7 +538,7 @@ def _generate_single_chunk( generated_tokens = [] eos_token_ids = {151643, 151645} - from mlx_lm.generate import generate_step + from mlx_audio.lm.generate import generate_step for token, _ in generate_step( prompt=input_ids[0], @@ -585,8 +585,7 @@ def generate( verbose: bool = False, **kwargs, ) -> STTOutput: - from mlx_lm.sample_utils import make_logits_processors, make_sampler - + from mlx_audio.lm.sample_utils import make_logits_processors, make_sampler from mlx_audio.stt.utils import load_audio del verbose, kwargs diff --git a/mlx_audio/stt/models/glmasr/glmasr.py b/mlx_audio/stt/models/glmasr/glmasr.py index 3c03577f2..5173bf00d 100644 --- a/mlx_audio/stt/models/glmasr/glmasr.py +++ b/mlx_audio/stt/models/glmasr/glmasr.py @@ -356,7 +356,7 @@ def __init__(self, config: LlamaConfig): self.config = config self.model_type = config.model_type - from mlx_lm.models.llama import LlamaModel + from mlx_audio.lm.models.llama import LlamaModel self.model = LlamaModel(config) @@ -613,7 +613,7 @@ def stream_generate( Yields: Tuple of (token, logprobs) """ - from mlx_lm.generate import generate_step + from mlx_audio.lm.generate import generate_step input_embeddings = self._merge_audio_text_embeddings( input_ids=input_ids, @@ -754,7 +754,7 @@ def generate( verbose=verbose, ) - from mlx_lm.sample_utils import make_sampler + from mlx_audio.lm.sample_utils import make_sampler start_time = time.time() @@ -973,8 +973,7 @@ def stream_transcribe( Yields: StreamingResult objects with text, timing, and status information. """ - from mlx_lm.sample_utils import make_sampler - + from mlx_audio.lm.sample_utils import make_sampler from mlx_audio.stt.utils import load_audio # Load audio diff --git a/mlx_audio/stt/models/granite_speech/granite_speech.py b/mlx_audio/stt/models/granite_speech/granite_speech.py index 3c5fa7cb4..aa61bf5f0 100644 --- a/mlx_audio/stt/models/granite_speech/granite_speech.py +++ b/mlx_audio/stt/models/granite_speech/granite_speech.py @@ -7,11 +7,11 @@ import mlx.core as mx import mlx.nn as nn import numpy as np -from mlx_lm.models.base import create_attention_mask -from mlx_lm.models.cache import KVCache -from mlx_lm.models.granite import Model as GraniteLM -from mlx_lm.models.granite import ModelArgs as GraniteModelArgs +from mlx_audio.lm.models.base import create_attention_mask +from mlx_audio.lm.models.cache import KVCache +from mlx_audio.lm.models.granite import Model as GraniteLM +from mlx_audio.lm.models.granite import ModelArgs as GraniteModelArgs from mlx_audio.stt.models.base import STTOutput from .config import EncoderConfig, ModelConfig, ProjectorConfig @@ -661,8 +661,8 @@ def generate( start_time = time.time() - from mlx_lm.generate import generate_step - from mlx_lm.sample_utils import make_logits_processors, make_sampler + from mlx_audio.lm.generate import generate_step + from mlx_audio.lm.sample_utils import make_logits_processors, make_sampler audio_data = self._load_audio(audio) input_features, num_audio_tokens = self._extract_features(audio_data) @@ -737,8 +737,8 @@ def _stream_generate( prefill_step_size: int = 2048, verbose: bool = False, ) -> Generator[StreamingResult, None, None]: - from mlx_lm.generate import generate_step - from mlx_lm.sample_utils import make_logits_processors, make_sampler + from mlx_audio.lm.generate import generate_step + from mlx_audio.lm.sample_utils import make_logits_processors, make_sampler audio_data = self._load_audio(audio) input_features, num_audio_tokens = self._extract_features(audio_data) diff --git a/mlx_audio/stt/models/higgs_audio_3/higgs_audio_3.py b/mlx_audio/stt/models/higgs_audio_3/higgs_audio_3.py index df899a9a6..fa3b79e27 100644 --- a/mlx_audio/stt/models/higgs_audio_3/higgs_audio_3.py +++ b/mlx_audio/stt/models/higgs_audio_3/higgs_audio_3.py @@ -6,10 +6,10 @@ import mlx.core as mx import mlx.nn as nn import numpy as np -from mlx_lm.models.cache import KVCache -from mlx_lm.models.qwen3 import ModelArgs as Qwen3Args -from mlx_lm.models.qwen3 import Qwen3Model +from mlx_audio.lm.models.cache import KVCache +from mlx_audio.lm.models.qwen3 import ModelArgs as Qwen3Args +from mlx_audio.lm.models.qwen3 import Qwen3Model from mlx_audio.stt.models.base import STTOutput from .audio import AudioFeatureExtractor @@ -338,8 +338,8 @@ def generate( verbose: bool = False, **kwargs, ) -> STTOutput: - from mlx_lm.generate import generate_step - from mlx_lm.sample_utils import make_logits_processors, make_sampler + from mlx_audio.lm.generate import generate_step + from mlx_audio.lm.sample_utils import make_logits_processors, make_sampler if not hasattr(self, "_tokenizer"): raise RuntimeError("Tokenizer not initialized. Call post_load_hook first.") diff --git a/mlx_audio/stt/models/moss_music/convert.py b/mlx_audio/stt/models/moss_music/convert.py index d2940ae85..62ce28d97 100644 --- a/mlx_audio/stt/models/moss_music/convert.py +++ b/mlx_audio/stt/models/moss_music/convert.py @@ -102,7 +102,7 @@ def _quantize( q_bits: int, q_group_size: int, ) -> tuple[dict[str, Any], dict[str, mx.array]]: - from mlx_lm.utils import quantize_model + from mlx_audio.lm.convert import quantize_model from .config import ModelConfig from .moss_music import Model diff --git a/mlx_audio/stt/models/moss_music/moss_music.py b/mlx_audio/stt/models/moss_music/moss_music.py index ea4d1f72f..729c3ce9a 100644 --- a/mlx_audio/stt/models/moss_music/moss_music.py +++ b/mlx_audio/stt/models/moss_music/moss_music.py @@ -8,11 +8,11 @@ import mlx.core as mx import mlx.nn as nn import numpy as np -from mlx_lm.models.base import create_attention_mask -from mlx_lm.models.cache import KVCache -from mlx_lm.models.qwen3 import ModelArgs as Qwen3Args -from mlx_lm.models.qwen3 import Qwen3Model +from mlx_audio.lm.models.base import create_attention_mask +from mlx_audio.lm.models.cache import KVCache +from mlx_audio.lm.models.qwen3 import ModelArgs as Qwen3Args +from mlx_audio.lm.models.qwen3 import Qwen3Model from mlx_audio.stt.models.base import STTOutput from .config import AudioEncoderConfig, ModelConfig @@ -482,7 +482,7 @@ def _generate_tokens( repetition_context_size: int, prefill_step_size: int, ) -> Generator[int, None, None]: - from mlx_lm.sample_utils import make_logits_processors, make_sampler + from mlx_audio.lm.sample_utils import make_logits_processors, make_sampler sampler = make_sampler(temperature, top_p=top_p, min_p=min_p, top_k=top_k) logits_processors = make_logits_processors( diff --git a/mlx_audio/stt/models/moss_transcribe_diarize/moss_transcribe_diarize.py b/mlx_audio/stt/models/moss_transcribe_diarize/moss_transcribe_diarize.py index 75674685f..5cd5c4941 100644 --- a/mlx_audio/stt/models/moss_transcribe_diarize/moss_transcribe_diarize.py +++ b/mlx_audio/stt/models/moss_transcribe_diarize/moss_transcribe_diarize.py @@ -242,7 +242,7 @@ def sample_rate(self) -> int: return self.config.sample_rate def make_cache(self) -> List[Any]: - from mlx_lm.models.cache import KVCache + from mlx_audio.lm.models.cache import KVCache return [KVCache() for _ in range(self.config.text_config.num_hidden_layers)] @@ -592,7 +592,7 @@ def stream_generate( prefill_step_size: int = 4096, verbose: bool = False, ) -> Generator[Tuple[mx.array, mx.array], None, None]: - from mlx_lm.generate import generate_step + from mlx_audio.lm.generate import generate_step prompt_ids, inputs_embeds, _, _ = self._prepare_generation_inputs(audio, prompt) eos_token_ids = self._eos_token_ids() @@ -664,9 +664,14 @@ def generate( prefill_step_size: int = 4096, verbose: bool = False, stream: bool = False, + hotwords: Optional[List[str]] = None, **kwargs, ) -> Union[STTOutput, Generator[StreamingResult, None, None]]: - from mlx_lm.sample_utils import make_logits_processors, make_sampler + from mlx_audio.lm.sample_utils import make_logits_processors, make_sampler + from mlx_audio.stt.utils import merge_hotwords + + # MOSS biases toward rare vocabulary via the prompt. + prompt = merge_hotwords(prompt, hotwords) sampler = make_sampler(temperature, top_p=top_p, min_p=min_p, top_k=top_k) logits_processors = make_logits_processors( @@ -695,7 +700,7 @@ def generate( generated_tokens = [] gen_start = time.time() - from mlx_lm.generate import generate_step + from mlx_audio.lm.generate import generate_step for token, _ in generate_step( prompt=prompt_ids, diff --git a/mlx_audio/stt/models/nemotron_asr/convert.py b/mlx_audio/stt/models/nemotron_asr/convert.py index 8237f864d..0e440b7b4 100644 --- a/mlx_audio/stt/models/nemotron_asr/convert.py +++ b/mlx_audio/stt/models/nemotron_asr/convert.py @@ -171,7 +171,8 @@ def _quantize(config: dict, weights: dict, q_bits: int, q_group_size: int): are left in the base dtype — the standard mlx quantization predicate. """ from mlx.utils import tree_flatten - from mlx_lm.utils import quantize_model + + from mlx_audio.lm.convert import quantize_model from .nemotron_asr import Model, ModelConfig diff --git a/mlx_audio/stt/models/qwen2_audio/qwen2_audio.py b/mlx_audio/stt/models/qwen2_audio/qwen2_audio.py index 2e1850996..aeecd519b 100644 --- a/mlx_audio/stt/models/qwen2_audio/qwen2_audio.py +++ b/mlx_audio/stt/models/qwen2_audio/qwen2_audio.py @@ -7,11 +7,11 @@ import mlx.core as mx import mlx.nn as nn import numpy as np -from mlx_lm.models.base import create_attention_mask -from mlx_lm.models.cache import KVCache -from mlx_lm.models.qwen2 import Model as Qwen2LM -from mlx_lm.models.qwen2 import ModelArgs as Qwen2ModelArgs +from mlx_audio.lm.models.base import create_attention_mask +from mlx_audio.lm.models.cache import KVCache +from mlx_audio.lm.models.qwen2 import Model as Qwen2LM +from mlx_audio.lm.models.qwen2 import ModelArgs as Qwen2ModelArgs from mlx_audio.stt.models.base import STTOutput from .config import EncoderConfig, ModelConfig @@ -477,8 +477,8 @@ def generate( verbose=verbose, ) - from mlx_lm.generate import generate_step - from mlx_lm.sample_utils import make_logits_processors, make_sampler + from mlx_audio.lm.generate import generate_step + from mlx_audio.lm.sample_utils import make_logits_processors, make_sampler start_time = time.time() prompt_ids, inputs_embeds, prompt_tokens = self.get_input_embeddings( @@ -542,8 +542,8 @@ def _stream_generate( prefill_step_size: int = 2048, verbose: bool = False, ) -> Generator[StreamingResult, None, None]: - from mlx_lm.generate import generate_step - from mlx_lm.sample_utils import make_logits_processors, make_sampler + from mlx_audio.lm.generate import generate_step + from mlx_audio.lm.sample_utils import make_logits_processors, make_sampler prompt_ids, inputs_embeds, prompt_token_count = self.get_input_embeddings( audio, prompt, verbose diff --git a/mlx_audio/stt/models/qwen3_asr/qwen3_asr.py b/mlx_audio/stt/models/qwen3_asr/qwen3_asr.py index db4dc3bff..c1219a0ea 100644 --- a/mlx_audio/stt/models/qwen3_asr/qwen3_asr.py +++ b/mlx_audio/stt/models/qwen3_asr/qwen3_asr.py @@ -9,9 +9,9 @@ import mlx.core as mx import mlx.nn as nn import numpy as np -from mlx_lm.models.base import create_attention_mask, scaled_dot_product_attention from tqdm import tqdm +from mlx_audio.lm.models.base import create_attention_mask, scaled_dot_product_attention from mlx_audio.stt.models.base import STTOutput from .config import AudioEncoderConfig, ModelConfig, TextConfig @@ -768,7 +768,7 @@ def sample_rate(self) -> int: def make_cache(self) -> List[Any]: """Create KV cache for generation.""" - from mlx_lm.models.cache import KVCache + from mlx_audio.lm.models.cache import KVCache return [KVCache() for _ in range(self.config.text_config.num_hidden_layers)] @@ -959,7 +959,7 @@ def stream_generate( system_prompt: str | None = None, ) -> Generator[Tuple[mx.array, mx.array], None, None]: """Stream generate tokens from audio using mlx_lm generate_step.""" - from mlx_lm.generate import generate_step + from mlx_audio.lm.generate import generate_step if not hasattr(self, "_tokenizer") or not hasattr(self, "_feature_extractor"): raise RuntimeError( @@ -1111,7 +1111,7 @@ def _generate_chunks_batched( across the batch. Chunks within a batch are padded (audio) to equal length so prompts share one length and the plain causal mask stays valid. """ - from mlx_lm.models.cache import KVCache + from mlx_audio.lm.models.cache import KVCache eos_token_ids = self._eos_token_ids() texts = [""] * len(chunks) @@ -1246,6 +1246,7 @@ def generate( verbose: bool = False, stream: bool = False, system_prompt: str | None = None, + hotwords: Optional[List[str]] = None, **kwargs, ) -> Union[STTOutput, Generator[str, None, None]]: """Generate transcription from audio. @@ -1256,7 +1257,13 @@ def generate( chunk_duration: Maximum chunk duration in seconds (default: 1200 = 20 min). min_chunk_duration: Minimum chunk duration in seconds (default: 1.0). stream: If True, return a generator that yields tokens as they are generated. + hotwords: Optional vocabulary/hotword list, folded into ``system_prompt``. """ + from mlx_audio.stt.utils import merge_hotwords + + # Qwen3-ASR biases toward rare vocabulary via the system prompt. + system_prompt = merge_hotwords(system_prompt, hotwords) + # If streaming requested, delegate to stream_transcribe if stream: return self.stream_transcribe( @@ -1277,8 +1284,7 @@ def generate( system_prompt=system_prompt, ) - from mlx_lm.sample_utils import make_logits_processors, make_sampler - + from mlx_audio.lm.sample_utils import make_logits_processors, make_sampler from mlx_audio.stt.utils import load_audio del kwargs @@ -1461,8 +1467,7 @@ def stream_transcribe( Yields: StreamingResult objects with text, timing, and status information. """ - from mlx_lm.sample_utils import make_logits_processors, make_sampler - + from mlx_audio.lm.sample_utils import make_logits_processors, make_sampler from mlx_audio.stt.utils import load_audio if not hasattr(self, "_tokenizer") or not hasattr(self, "_feature_extractor"): diff --git a/mlx_audio/stt/models/vibevoice_asr/vibevoice_asr.py b/mlx_audio/stt/models/vibevoice_asr/vibevoice_asr.py index cbfa7ccfd..f0a0b43f2 100644 --- a/mlx_audio/stt/models/vibevoice_asr/vibevoice_asr.py +++ b/mlx_audio/stt/models/vibevoice_asr/vibevoice_asr.py @@ -50,12 +50,12 @@ def __init__(self, config): # Import and use mlx_lm's Qwen2Model try: - from mlx_lm.models.qwen2 import Qwen2Model + from mlx_audio.lm.models.qwen2 import Qwen2Model self.model = Qwen2Model(config) except ImportError: # Fallback to llama if qwen2 not available - from mlx_lm.models.llama import LlamaModel + from mlx_audio.lm.models.llama import LlamaModel self.model = LlamaModel(config) @@ -567,7 +567,7 @@ def stream_generate( Yields: Tuple of (token, logprobs) """ - from mlx_lm.generate import generate_step + from mlx_audio.lm.generate import generate_step # Get input embeddings with speech merged in input_embeddings = self._merge_speech_text_embeddings( @@ -649,6 +649,7 @@ def generate( prefill_step_size: int = 2048, generation_stream: bool = False, verbose: bool = False, + hotwords: Optional[List[str]] = None, **kwargs, ) -> STTOutput: """ @@ -657,6 +658,7 @@ def generate( Args: audio: Audio path (str) or waveform (mx.array/np.array) context: Optional context string (hotwords, metadata) + hotwords: Optional vocabulary/hotword list, folded into ``context`` sampling_rate: Sample rate of the input waveform. When *audio* is an array not at 24 kHz, provide its actual sample rate so that it is resampled correctly. @@ -675,7 +677,11 @@ def generate( Returns: STTOutput with transcription text and segments """ - from mlx_lm.sample_utils import make_logits_processors, make_sampler + from mlx_audio.lm.sample_utils import make_logits_processors, make_sampler + from mlx_audio.stt.utils import merge_hotwords + + # VibeVoice biases toward rare vocabulary via the context string. + context = merge_hotwords(context, hotwords) start_time = time.time() @@ -787,7 +793,7 @@ def stream_transcribe( Yields: Decoded text chunks as they are generated. """ - from mlx_lm.sample_utils import make_logits_processors, make_sampler + from mlx_audio.lm.sample_utils import make_logits_processors, make_sampler # Preprocess audio audio_tensor = self._preprocess_audio(audio, sampling_rate=sampling_rate) diff --git a/mlx_audio/stt/models/voxtral/voxtral.py b/mlx_audio/stt/models/voxtral/voxtral.py index 6b4a99800..70e5b7fd4 100644 --- a/mlx_audio/stt/models/voxtral/voxtral.py +++ b/mlx_audio/stt/models/voxtral/voxtral.py @@ -203,7 +203,7 @@ def __init__(self, config: ModelConfig): self.config = config self.model_type = config.model_type - from mlx_lm.models.llama import LlamaModel + from mlx_audio.lm.models.llama import LlamaModel self.model = LlamaModel(config) @@ -379,7 +379,7 @@ def stream_generate( verbose: bool = False, ) -> Generator[Tuple[mx.array, mx.array], None, None]: - from mlx_lm.generate import generate_step + from mlx_audio.lm.generate import generate_step input_embeddings = self._merge_input_embeddings( input_ids=input_ids, @@ -445,7 +445,7 @@ def generate( generated = [] - from mlx_lm.sample_utils import make_sampler + from mlx_audio.lm.sample_utils import make_sampler sampler = make_sampler( temperature, diff --git a/mlx_audio/stt/models/voxtral_realtime/decoder.py b/mlx_audio/stt/models/voxtral_realtime/decoder.py index 93eb31a61..f53ebaec0 100644 --- a/mlx_audio/stt/models/voxtral_realtime/decoder.py +++ b/mlx_audio/stt/models/voxtral_realtime/decoder.py @@ -20,7 +20,8 @@ import mlx.core as mx import mlx.nn as nn -from mlx_lm.models.cache import RotatingKVCache + +from mlx_audio.lm.models.cache import RotatingKVCache from .config import DecoderConfig diff --git a/mlx_audio/stt/models/voxtral_realtime/encoder.py b/mlx_audio/stt/models/voxtral_realtime/encoder.py index bdc26f739..f9ab4f8e1 100644 --- a/mlx_audio/stt/models/voxtral_realtime/encoder.py +++ b/mlx_audio/stt/models/voxtral_realtime/encoder.py @@ -198,7 +198,7 @@ def encode_chunks(self, conv_out): Yields: mx.array: [chunk_size, dim] encoded chunk """ - from mlx_lm.models.cache import RotatingKVCache + from mlx_audio.lm.models.cache import RotatingKVCache seq_len = conv_out.shape[0] sw = self.config.sliding_window diff --git a/mlx_audio/stt/models/voxtral_realtime/streaming.py b/mlx_audio/stt/models/voxtral_realtime/streaming.py index c2af11db9..f35d6b7ef 100644 --- a/mlx_audio/stt/models/voxtral_realtime/streaming.py +++ b/mlx_audio/stt/models/voxtral_realtime/streaming.py @@ -354,7 +354,7 @@ class StreamingEncoder: """ def __init__(self, encoder): - from mlx_lm.models.cache import RotatingKVCache + from mlx_audio.lm.models.cache import RotatingKVCache self.encoder = encoder self._sw = encoder.config.sliding_window diff --git a/mlx_audio/stt/models/whisper/whisper.py b/mlx_audio/stt/models/whisper/whisper.py index 6824c6449..24c2b7686 100644 --- a/mlx_audio/stt/models/whisper/whisper.py +++ b/mlx_audio/stt/models/whisper/whisper.py @@ -818,6 +818,7 @@ def generate( append_punctuations: str = "\"'.。,，!！?？:：”)]}、", clip_timestamps: Union[str, List[float]] = "0", hallucination_silence_threshold: Optional[float] = None, + hotwords: Optional[List[str]] = None, **decode_options, ): """ @@ -885,7 +886,14 @@ def generate( ------- A dictionary containing the resulting text ("text") and segment-level details ("segments"), and the spoken language ("language"), which is detected when `decode_options["language"]` is None. + + hotwords: Optional[List[str]] + Vocabulary/hotword list, folded into ``initial_prompt`` to bias decoding. """ + # Whisper biases toward rare vocabulary via the initial prompt. + from mlx_audio.stt.utils import merge_hotwords + + initial_prompt = merge_hotwords(initial_prompt, hotwords) if stream: return self.generate_streaming( diff --git a/mlx_audio/stt/tests/test_fun_asr_nano.py b/mlx_audio/stt/tests/test_fun_asr_nano.py index 569d809b0..22d66a324 100644 --- a/mlx_audio/stt/tests/test_fun_asr_nano.py +++ b/mlx_audio/stt/tests/test_fun_asr_nano.py @@ -108,7 +108,7 @@ def fake_generate_step(**kwargs): yield 151643, None model._build_inputs_embeds = MethodType(fake_build_inputs, model) - generate_module = importlib.import_module("mlx_lm.generate") + generate_module = importlib.import_module("mlx_audio.lm.generate") monkeypatch.setattr(generate_module, "generate_step", fake_generate_step) assert list(model.stream_generate(mx.zeros((1,)), context=" MLX ")) == [] diff --git a/mlx_audio/stt/tests/test_hotwords.py b/mlx_audio/stt/tests/test_hotwords.py new file mode 100644 index 000000000..f967ed61f --- /dev/null +++ b/mlx_audio/stt/tests/test_hotwords.py @@ -0,0 +1,67 @@ +"""Uniform ``hotwords`` support across ASR backends (mlx-vlm #1781, part one). + +Each supporting model folds a structured ``hotwords`` list into its own native +prompt field inside its ``generate`` forward loop; models without a hook ignore +it (silent drop). These tests stay weight-free: they exercise the shared merge +helper and introspect each model's ``generate`` signature. +""" + +import ast +import inspect +from pathlib import Path + +from mlx_audio.stt.utils import merge_hotwords + + +class TestMergeHotwords: + def test_none_is_noop(self): + assert merge_hotwords("base", None) == "base" + assert merge_hotwords(None, None) is None + + def test_empty_or_blank_terms_noop(self): + assert merge_hotwords("base", []) == "base" + assert merge_hotwords("base", ["", " ", None]) == "base" + + def test_terms_only(self): + assert merge_hotwords(None, ["Nativ", "MLX"]) == "Nativ, MLX" + + def test_appends_to_existing_base(self): + assert ( + merge_hotwords("Prior text", ["Nativ", "MLX"]) == "Prior text\nNativ, MLX" + ) + + def test_strips_and_drops_blanks(self): + assert merge_hotwords(None, [" Nativ ", "", "MLX"]) == "Nativ, MLX" + + +# model file -> the native prompt field hotwords must fold into +_NATIVE_FIELD = { + "qwen3_asr/qwen3_asr.py": "system_prompt", + "whisper/whisper.py": "initial_prompt", + "vibevoice_asr/vibevoice_asr.py": "context", + "moss_transcribe_diarize/moss_transcribe_diarize.py": "prompt", +} + +_MODELS_DIR = Path(__file__).resolve().parents[1] / "models" + + +class TestGenerateAcceptsHotwords: + def _generate_args(self, rel_path): + tree = ast.parse((_MODELS_DIR / rel_path).read_text()) + for node in ast.walk(tree): + if isinstance(node, ast.FunctionDef) and node.name == "generate": + args = [a.arg for a in node.args.args + node.args.kwonlyargs] + if "hotwords" in args or "audio" in args: + return args + return [] + + def test_each_supporting_model_exposes_hotwords_and_native_field(self): + for rel_path, native in _NATIVE_FIELD.items(): + args = self._generate_args(rel_path) + assert "hotwords" in args, f"{rel_path}: generate() missing hotwords" + assert native in args, f"{rel_path}: generate() missing {native}" + + def test_merge_helper_is_used_in_each_model(self): + for rel_path in _NATIVE_FIELD: + src = (_MODELS_DIR / rel_path).read_text() + assert "merge_hotwords" in src, f"{rel_path}: does not call merge_hotwords" diff --git a/mlx_audio/stt/tests/test_models.py b/mlx_audio/stt/tests/test_models.py index 644ffeb3d..c5405a678 100644 --- a/mlx_audio/stt/tests/test_models.py +++ b/mlx_audio/stt/tests/test_models.py @@ -435,9 +435,8 @@ def _small_config_dict(self): return _cohere_small_config_dict() def _build_quantized_checkpoint(self, bits: int) -> Path: - from mlx_lm.utils import save_model - from mlx_audio.convert import convert + from mlx_audio.lm.convert import save_model from mlx_audio.stt.models.cohere_asr.cohere_asr import Model, STTOutput self.STTOutput = STTOutput @@ -1803,7 +1802,7 @@ def test_text_model_forward_with_embeddings(self): self.assertEqual(output.shape, (1, 5, self.text_config.hidden_size)) def test_text_model_cached_chunks_match_full_causal_pass(self): - from mlx_lm.models.cache import KVCache + from mlx_audio.lm.models.cache import KVCache model = self.TextModel(self.text_config) input_ids = mx.array([[1, 2, 3, 4, 5]], dtype=mx.int32) diff --git a/mlx_audio/stt/utils.py b/mlx_audio/stt/utils.py index b112ec101..4f8b15040 100644 --- a/mlx_audio/stt/utils.py +++ b/mlx_audio/stt/utils.py @@ -12,6 +12,28 @@ SAMPLE_RATE = 16000 +def merge_hotwords(base: Optional[str], hotwords: Optional[List[str]]) -> Optional[str]: + """Merge a structured vocabulary/hotword list into a model's native prompt. + + ASR backends bias their transcription toward rare words (names, acronyms, + product terms) through a text field that differs per model — ``system_prompt`` + for Qwen3-ASR, ``initial_prompt`` for Whisper, ``context`` for VibeVoice, + ``prompt`` for MOSS, etc. This helper lets each model accept a uniform + ``hotwords`` list and fold it into whichever field it already uses, so callers + (e.g. the mlx-vlm server) can pass one structured list regardless of backend. + + Returns ``base`` unchanged when ``hotwords`` is empty/None (silent no-op), the + joined terms when ``base`` is empty, or ``base`` followed by the terms. + """ + if not hotwords: + return base + terms = [str(t).strip() for t in hotwords if t is not None and str(t).strip()] + if not terms: + return base + joined = ", ".join(terms) + return f"{base}\n{joined}" if base else joined + + @contextlib.contextmanager def wired_limit(model: nn.Module, streams: Optional[List[mx.Stream]] = None): """ diff --git a/mlx_audio/tests/test_lazy_imports.py b/mlx_audio/tests/test_lazy_imports.py index 4fea6eebf..92328caeb 100644 --- a/mlx_audio/tests/test_lazy_imports.py +++ b/mlx_audio/tests/test_lazy_imports.py @@ -73,3 +73,18 @@ def test_codec_no_eager_imports(): text=True, ) assert result.returncode == 0, f"Codec lazy import failed: {result.stderr}" + + +def test_vendored_lm_no_eager_imports(): + code = """ +import sys +import mlx_audio.lm.models.llama +assert "mlx_lm" not in sys.modules, "mlx_lm was eagerly imported" +assert "transformers" not in sys.modules, "transformers was eagerly imported" +""" + result = subprocess.run( + [sys.executable, "-c", code], + capture_output=True, + text=True, + ) + assert result.returncode == 0, f"Vendored LM lazy import failed: {result.stderr}" diff --git a/mlx_audio/tests/test_lm_generate.py b/mlx_audio/tests/test_lm_generate.py new file mode 100644 index 000000000..410f50cc1 --- /dev/null +++ b/mlx_audio/tests/test_lm_generate.py @@ -0,0 +1,85 @@ +import mlx.core as mx +import mlx.nn as nn +import pytest + +from mlx_audio.lm.generate import generate_step, stream_generate + +VOCAB = 17 +EOS = 5 + + +class _Cache: + state = [] + + +class CycleModel(nn.Module): + """Emits token (t + 1) % VOCAB, so a prompt of EOS-1 produces EOS first.""" + + layers = [object()] + + def make_cache(self): + return [_Cache()] + + def __call__(self, tokens, cache=None, input_embeddings=None): + del cache, input_embeddings + return mx.eye(VOCAB)[(tokens + 1) % VOCAB] + + +class Tok: + eos_token_ids = {EOS} + eos_token_id = EOS + bos_token = None + clean_up_tokenization_spaces = False + + def encode(self, text, **kwargs): + return [1] + + def decode(self, ids, **kwargs): + return " ".join(str(i) for i in ids) + + +def responses(prompt, **kwargs): + return list(stream_generate(CycleModel(), Tok(), mx.array(prompt), **kwargs)) + + +def test_final_response_carries_eos_token_and_stop_reason(): + out = responses([EOS - 1], max_tokens=10) + assert out, "expected at least the terminal response" + assert out[-1].finish_reason == "stop" + assert out[-1].token == EOS + + +def test_first_token_eos_still_yields_a_final_response(): + """A caller reading finish_reason off the last response must get one.""" + out = responses([EOS - 1], max_tokens=10) + assert len(out) == 1 + assert out[-1].finish_reason == "stop" + + +def test_length_finish_reason_when_eos_never_reached(): + out = responses([EOS + 1], max_tokens=3) + assert out[-1].finish_reason == "length" + assert out[-1].token != EOS + + +def test_eos_token_is_not_a_duplicate_of_the_previous_token(): + """Re-emitting the last audio code instead of EOS shifts codec framing.""" + out = responses([EOS - 3], max_tokens=10) + assert out[-1].token == EOS + if len(out) > 1: + assert out[-1].token != out[-2].token + + +@pytest.mark.parametrize("max_tokens", [1, 2, 5]) +def test_generate_step_respects_max_tokens(max_tokens): + toks = [ + int(t) + for t, _ in generate_step(mx.array([1, 2]), CycleModel(), max_tokens=max_tokens) + ] + assert len(toks) == max_tokens + + +def test_generate_step_negative_max_tokens_is_unbounded(): + stream = generate_step(mx.array([1, 2]), CycleModel(), max_tokens=-1) + produced = [int(pair[0]) for pair, _ in zip(stream, range(40))] + assert len(produced) == 40 diff --git a/mlx_audio/tests/test_no_mlx_lm.py b/mlx_audio/tests/test_no_mlx_lm.py new file mode 100644 index 000000000..bcf07e161 --- /dev/null +++ b/mlx_audio/tests/test_no_mlx_lm.py @@ -0,0 +1,62 @@ +"""mlx-lm is vendored under mlx_audio/lm; only the optional LLM responder may import it.""" + +import ast +import re +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[2] +SOURCE = ROOT / "mlx_audio" +ALLOWED = {"mlx_audio/sts/voice_pipeline.py"} +DYNAMIC = re.compile(r"""(?:importlib\.import_module|__import__)\(\s*["']mlx_lm""") + + +def _sources(): + for path in SOURCE.rglob("*.py"): + rel = path.relative_to(ROOT).as_posix() + if "/tests/" in rel or rel in ALLOWED: + continue + yield rel, path + + +def test_no_mlx_lm_imports_in_source(): + offenders = [] + for rel, path in _sources(): + for node in ast.walk(ast.parse(path.read_text(encoding="utf-8"))): + if isinstance(node, ast.Import): + names = [alias.name for alias in node.names] + elif isinstance(node, ast.ImportFrom): + names = [node.module or ""] + else: + continue + if any(n == "mlx_lm" or n.startswith("mlx_lm.") for n in names): + offenders.append(f"{rel}:{node.lineno}") + assert not offenders, ( + f"mlx_lm imported at {offenders}. mlx-audio vendors this machinery: import " + "from mlx_audio.lm instead (see docs/contributing/adding-a-model.md)." + ) + + +def test_no_dynamic_mlx_lm_imports_in_source(): + offenders = [ + rel + for rel, path in _sources() + if DYNAMIC.search(path.read_text(encoding="utf-8")) + ] + assert not offenders, f"dynamic mlx_lm import in: {offenders}" + + +def test_mlx_lm_is_not_a_core_dependency(): + lines = (ROOT / "pyproject.toml").read_text(encoding="utf-8").splitlines() + section = None + core = [] + for line in lines: + stripped = line.strip() + if stripped.startswith("[") and stripped.endswith("]"): + section = stripped + elif ( + section == "[project]" + and "mlx-lm" in stripped + and not stripped.startswith("#") + ): + core.append(stripped) + assert not core, f"mlx-lm must live in an optional extra, found: {core}" diff --git a/mlx_audio/tts/models/bailingmm/bailingmm.py b/mlx_audio/tts/models/bailingmm/bailingmm.py index 08efa29b2..ba4145f59 100644 --- a/mlx_audio/tts/models/bailingmm/bailingmm.py +++ b/mlx_audio/tts/models/bailingmm/bailingmm.py @@ -8,17 +8,17 @@ import mlx.core as mx import mlx.nn as nn -import mlx_lm.models.bailing_moe as bailing_moe_impl -import mlx_lm.models.qwen2 as qwen2_impl import numpy as np from huggingface_hub import snapshot_download -from mlx_lm.models.bailing_moe import Model as BailingMoeModel -from mlx_lm.models.bailing_moe import ModelArgs as BailingMoeModelArgs -from mlx_lm.models.base import create_attention_mask -from mlx_lm.models.cache import KVCache -from mlx_lm.models.qwen2 import ModelArgs as Qwen2ModelArgs -from mlx_lm.models.qwen2 import Qwen2Model +import mlx_audio.lm.models.bailing_moe as bailing_moe_impl +import mlx_audio.lm.models.qwen2 as qwen2_impl +from mlx_audio.lm.models.bailing_moe import Model as BailingMoeModel +from mlx_audio.lm.models.bailing_moe import ModelArgs as BailingMoeModelArgs +from mlx_audio.lm.models.base import create_attention_mask +from mlx_audio.lm.models.cache import KVCache +from mlx_audio.lm.models.qwen2 import ModelArgs as Qwen2ModelArgs +from mlx_audio.lm.models.qwen2 import Qwen2Model from mlx_audio.tts.models.base import BaseModelArgs, GenerationResult from mlx_audio.utils import load_audio diff --git a/mlx_audio/tts/models/bark/bark.py b/mlx_audio/tts/models/bark/bark.py index 07ed7f324..4d99b008a 100644 --- a/mlx_audio/tts/models/bark/bark.py +++ b/mlx_audio/tts/models/bark/bark.py @@ -11,9 +11,10 @@ import numpy as np import tqdm from mlx.utils import tree_map, tree_unflatten -from mlx_lm.models.base import create_causal_mask from transformers import BertTokenizer +from mlx_audio.lm.models.base import create_causal_mask + from ..base import BaseModelArgs, GenerationResult from .pipeline import Pipeline diff --git a/mlx_audio/tts/models/chatterbox/scripts/convert.py b/mlx_audio/tts/models/chatterbox/scripts/convert.py index fa500c47e..67b61371a 100644 --- a/mlx_audio/tts/models/chatterbox/scripts/convert.py +++ b/mlx_audio/tts/models/chatterbox/scripts/convert.py @@ -46,7 +46,7 @@ pip install torch safetensors huggingface_hub onnx s3tokenizer After conversion, the model only needs: - pip install mlx mlx-lm + pip install mlx """ import argparse diff --git a/mlx_audio/tts/models/chatterbox/scripts/convert_chatterbox.py b/mlx_audio/tts/models/chatterbox/scripts/convert_chatterbox.py index 83820002c..7e8b80cef 100644 --- a/mlx_audio/tts/models/chatterbox/scripts/convert_chatterbox.py +++ b/mlx_audio/tts/models/chatterbox/scripts/convert_chatterbox.py @@ -33,7 +33,7 @@ pip install torch safetensors huggingface_hub onnx s3tokenizer After conversion, the model only needs: - pip install mlx mlx-lm + pip install mlx """ import argparse diff --git a/mlx_audio/tts/models/chatterbox/t3/t3.py b/mlx_audio/tts/models/chatterbox/t3/t3.py index cb28fc478..ff95d4489 100644 --- a/mlx_audio/tts/models/chatterbox/t3/t3.py +++ b/mlx_audio/tts/models/chatterbox/t3/t3.py @@ -2,10 +2,11 @@ import mlx.core as mx import mlx.nn as nn -from mlx_lm.models.cache import make_prompt_cache -from mlx_lm.models.llama import Model as LlamaModel -from mlx_lm.models.llama import ModelArgs as LlamaModelConfig -from mlx_lm.sample_utils import make_logits_processors, make_sampler + +from mlx_audio.lm.models.cache import make_prompt_cache +from mlx_audio.lm.models.llama import Model as LlamaModel +from mlx_audio.lm.models.llama import ModelArgs as LlamaModelConfig +from mlx_audio.lm.sample_utils import make_logits_processors, make_sampler from ..config import LLAMA_CONFIGS, T3Config from .cond_enc import T3Cond, T3CondEnc diff --git a/mlx_audio/tts/models/chatterbox_turbo/models/t3/gpt2.py b/mlx_audio/tts/models/chatterbox_turbo/models/t3/gpt2.py index 854a4584a..f283fe203 100644 --- a/mlx_audio/tts/models/chatterbox_turbo/models/t3/gpt2.py +++ b/mlx_audio/tts/models/chatterbox_turbo/models/t3/gpt2.py @@ -5,7 +5,8 @@ import mlx.core as mx import mlx.nn as nn -from mlx_lm.models.cache import KVCache + +from mlx_audio.lm.models.cache import KVCache @dataclass diff --git a/mlx_audio/tts/models/dia/dia.py b/mlx_audio/tts/models/dia/dia.py index 9218a8b34..255a566b2 100644 --- a/mlx_audio/tts/models/dia/dia.py +++ b/mlx_audio/tts/models/dia/dia.py @@ -6,10 +6,10 @@ import mlx.nn as nn import numpy as np from huggingface_hub import hf_hub_download -from mlx_lm.sample_utils import make_sampler from tqdm import trange from mlx_audio.codec.models import DAC +from mlx_audio.lm.sample_utils import make_sampler from mlx_audio.utils import load_audio from ..base import GenerationResult diff --git a/mlx_audio/tts/models/dramabox/gemma.py b/mlx_audio/tts/models/dramabox/gemma.py index 58bff9bdd..9b5edb737 100644 --- a/mlx_audio/tts/models/dramabox/gemma.py +++ b/mlx_audio/tts/models/dramabox/gemma.py @@ -5,8 +5,9 @@ from typing import Sequence import mlx.core as mx -from mlx_lm import load as mlx_lm_load -from mlx_lm.models.cache import create_causal_mask + +from mlx_audio.lm.load import load_lm +from mlx_audio.lm.models.cache import create_causal_mask @dataclass @@ -16,7 +17,7 @@ class EncodedPrompt: def load_text_encoder(model_id: str): - return mlx_lm_load(model_id) + return load_lm(model_id) def _language_core(model): diff --git a/mlx_audio/tts/models/fish_qwen3_omni/fish_speech.py b/mlx_audio/tts/models/fish_qwen3_omni/fish_speech.py index a02423705..91a3ba168 100644 --- a/mlx_audio/tts/models/fish_qwen3_omni/fish_speech.py +++ b/mlx_audio/tts/models/fish_qwen3_omni/fish_speech.py @@ -9,9 +9,9 @@ import mlx.core as mx import mlx.nn as nn -from mlx_lm.models.base import create_attention_mask -from mlx_lm.models.cache import KVCache +from mlx_audio.lm.models.base import create_attention_mask +from mlx_audio.lm.models.cache import KVCache from mlx_audio.tts.models.base import BatchGenerationResult, GenerationResult from .config import ModelConfig diff --git a/mlx_audio/tts/models/higgs_audio/higgs_audio.py b/mlx_audio/tts/models/higgs_audio/higgs_audio.py index 6f37e0b2d..32e7d12fc 100644 --- a/mlx_audio/tts/models/higgs_audio/higgs_audio.py +++ b/mlx_audio/tts/models/higgs_audio/higgs_audio.py @@ -17,11 +17,12 @@ import mlx.core as mx import mlx.nn as nn -from mlx_lm.models.base import create_causal_mask -from mlx_lm.models.cache import make_prompt_cache -from mlx_lm.models.llama import MLP as LlamaMLP -from mlx_lm.models.llama import Attention as LlamaAttention -from mlx_lm.models.llama import ModelArgs as LlamaModelArgs + +from mlx_audio.lm.models.base import create_causal_mask +from mlx_audio.lm.models.cache import make_prompt_cache +from mlx_audio.lm.models.llama import MLP as LlamaMLP +from mlx_audio.lm.models.llama import Attention as LlamaAttention +from mlx_audio.lm.models.llama import ModelArgs as LlamaModelArgs from .config import HiggsAudioConfig from .generation import ( diff --git a/mlx_audio/tts/models/higgs_audio_v3/config.py b/mlx_audio/tts/models/higgs_audio_v3/config.py index 554fe0e29..16e00ed13 100644 --- a/mlx_audio/tts/models/higgs_audio_v3/config.py +++ b/mlx_audio/tts/models/higgs_audio_v3/config.py @@ -3,7 +3,7 @@ from dataclasses import dataclass, field from typing import Any, Optional -from mlx_lm.models.qwen3 import ModelArgs as Qwen3ModelArgs +from mlx_audio.lm.models.qwen3 import ModelArgs as Qwen3ModelArgs @dataclass diff --git a/mlx_audio/tts/models/higgs_audio_v3/continuous_batching.py b/mlx_audio/tts/models/higgs_audio_v3/continuous_batching.py index 4ce8bef7c..dbd7b4310 100644 --- a/mlx_audio/tts/models/higgs_audio_v3/continuous_batching.py +++ b/mlx_audio/tts/models/higgs_audio_v3/continuous_batching.py @@ -5,8 +5,8 @@ from typing import Optional import mlx.core as mx -from mlx_lm.models.cache import BatchKVCache +from mlx_audio.lm.models.cache import BatchKVCache from mlx_audio.tts.continuous import TTSBatchEvent, TTSBatchItem, TTSBatchOptions from .generation import HiggsSamplerState diff --git a/mlx_audio/tts/models/higgs_audio_v3/generation.py b/mlx_audio/tts/models/higgs_audio_v3/generation.py index 20c38f321..736a818d2 100644 --- a/mlx_audio/tts/models/higgs_audio_v3/generation.py +++ b/mlx_audio/tts/models/higgs_audio_v3/generation.py @@ -5,8 +5,9 @@ import mlx.core as mx import mlx.nn as nn -from mlx_lm.sample_utils import apply_top_k as _apply_top_k_logprobs -from mlx_lm.sample_utils import apply_top_p as _apply_top_p_logprobs + +from mlx_audio.lm.sample_utils import apply_top_k as _apply_top_k_logprobs +from mlx_audio.lm.sample_utils import apply_top_p as _apply_top_p_logprobs STOP_CODE = -1 diff --git a/mlx_audio/tts/models/higgs_audio_v3/model.py b/mlx_audio/tts/models/higgs_audio_v3/model.py index 5a23ca862..a9f9c3266 100644 --- a/mlx_audio/tts/models/higgs_audio_v3/model.py +++ b/mlx_audio/tts/models/higgs_audio_v3/model.py @@ -7,8 +7,9 @@ import mlx.core as mx import mlx.nn as nn import numpy as np -from mlx_lm.models.cache import BatchKVCache, make_prompt_cache -from mlx_lm.models.qwen3 import Qwen3Model + +from mlx_audio.lm.models.cache import BatchKVCache, make_prompt_cache +from mlx_audio.lm.models.qwen3 import Qwen3Model from ..base import BatchGenerationResult, GenerationResult from .config import HiggsAudioV3Config diff --git a/mlx_audio/tts/models/indextts/gpt2.py b/mlx_audio/tts/models/indextts/gpt2.py index da47e83c5..50f365340 100644 --- a/mlx_audio/tts/models/indextts/gpt2.py +++ b/mlx_audio/tts/models/indextts/gpt2.py @@ -1,7 +1,8 @@ import mlx.core as mx import mlx.nn as nn -from mlx_lm.models.base import create_attention_mask -from mlx_lm.models.gpt2 import ModelArgs, TransformerBlock + +from mlx_audio.lm.models.base import create_attention_mask +from mlx_audio.lm.models.gpt2 import ModelArgs, TransformerBlock class GPT2Model(nn.Module): diff --git a/mlx_audio/tts/models/indextts/indextts.py b/mlx_audio/tts/models/indextts/indextts.py index ab542bbd1..5957b843d 100644 --- a/mlx_audio/tts/models/indextts/indextts.py +++ b/mlx_audio/tts/models/indextts/indextts.py @@ -8,10 +8,10 @@ import mlx.nn as nn import sentencepiece as spm import tqdm -from mlx_lm.models.cache import KVCache -from mlx_lm.models.gpt2 import ModelArgs as GPT2Args -from mlx_lm.sample_utils import make_sampler +from mlx_audio.lm.models.cache import KVCache +from mlx_audio.lm.models.gpt2 import ModelArgs as GPT2Args +from mlx_audio.lm.sample_utils import make_sampler from mlx_audio.tts.models.base import GenerationResult from mlx_audio.tts.models.indextts import normalize from mlx_audio.tts.models.indextts.attention import LearnedPositionEncoding diff --git a/mlx_audio/tts/models/kitten_tts/convert.py b/mlx_audio/tts/models/kitten_tts/convert.py index 23672f119..24d17a592 100644 --- a/mlx_audio/tts/models/kitten_tts/convert.py +++ b/mlx_audio/tts/models/kitten_tts/convert.py @@ -9,9 +9,10 @@ import onnx from huggingface_hub import hf_hub_download, snapshot_download from mlx.utils import tree_flatten -from mlx_lm.utils import save_config, save_model from onnx import helper, numpy_helper +from mlx_audio.lm.convert import save_config, save_model + from .kitten_tts import Model, ModelConfig diff --git a/mlx_audio/tts/models/llama/llama.py b/mlx_audio/tts/models/llama/llama.py index 67ce98191..218d74009 100644 --- a/mlx_audio/tts/models/llama/llama.py +++ b/mlx_audio/tts/models/llama/llama.py @@ -3,14 +3,14 @@ from typing import Generator, List, Optional, Union import mlx.core as mx -from mlx_lm.generate import stream_generate -from mlx_lm.models.llama import Model as LlamaModel -from mlx_lm.models.llama import ModelArgs as LlamaModelConfig -from mlx_lm.sample_utils import make_logits_processors, make_sampler from tqdm import tqdm from transformers import AutoTokenizer from mlx_audio.codec.models.snac import SNAC +from mlx_audio.lm.generate import stream_generate +from mlx_audio.lm.models.llama import Model as LlamaModel +from mlx_audio.lm.models.llama import ModelArgs as LlamaModelConfig +from mlx_audio.lm.sample_utils import make_logits_processors, make_sampler from mlx_audio.utils import load_audio from ..base import GenerationResult diff --git a/mlx_audio/tts/models/moss_tts/config.py b/mlx_audio/tts/models/moss_tts/config.py index b24d28104..cbaba4ddc 100644 --- a/mlx_audio/tts/models/moss_tts/config.py +++ b/mlx_audio/tts/models/moss_tts/config.py @@ -3,8 +3,7 @@ from dataclasses import dataclass, replace from typing import Any -from mlx_lm.models.qwen3 import ModelArgs as Qwen3ModelConfig - +from mlx_audio.lm.models.qwen3 import ModelArgs as Qwen3ModelConfig from mlx_audio.tts.models.base import BaseModelArgs from mlx_audio.tts.models.moss_tts_nano.config import GPT2Config diff --git a/mlx_audio/tts/models/moss_tts/moss_tts.py b/mlx_audio/tts/models/moss_tts/moss_tts.py index 760c57b0e..a9f0fdad1 100644 --- a/mlx_audio/tts/models/moss_tts/moss_tts.py +++ b/mlx_audio/tts/models/moss_tts/moss_tts.py @@ -7,9 +7,9 @@ import mlx.core as mx import mlx.nn as nn -from mlx_lm.models.cache import make_prompt_cache -from mlx_lm.models.qwen3 import Qwen3Model +from mlx_audio.lm.models.cache import make_prompt_cache +from mlx_audio.lm.models.qwen3 import Qwen3Model from mlx_audio.tts.models.base import GenerationResult from mlx_audio.tts.models.moss_tts_nano.gpt2 import GPT2Model diff --git a/mlx_audio/tts/models/moss_tts/sampling.py b/mlx_audio/tts/models/moss_tts/sampling.py index a69e94e0c..6b42d01c9 100644 --- a/mlx_audio/tts/models/moss_tts/sampling.py +++ b/mlx_audio/tts/models/moss_tts/sampling.py @@ -2,8 +2,9 @@ import mlx.core as mx import mlx.nn as nn -from mlx_lm.sample_utils import apply_top_k as _apply_top_k_logprobs -from mlx_lm.sample_utils import apply_top_p as _apply_top_p_logprobs + +from mlx_audio.lm.sample_utils import apply_top_k as _apply_top_k_logprobs +from mlx_audio.lm.sample_utils import apply_top_p as _apply_top_p_logprobs def _mask_logits_from_logprobs(logits: mx.array, logprobs: mx.array) -> mx.array: diff --git a/mlx_audio/tts/models/moss_tts_nano/sampling.py b/mlx_audio/tts/models/moss_tts_nano/sampling.py index f855332cd..7bbc273d4 100644 --- a/mlx_audio/tts/models/moss_tts_nano/sampling.py +++ b/mlx_audio/tts/models/moss_tts_nano/sampling.py @@ -2,8 +2,9 @@ import mlx.core as mx import mlx.nn as nn -from mlx_lm.sample_utils import apply_top_k as _apply_top_k_logprobs -from mlx_lm.sample_utils import apply_top_p as _apply_top_p_logprobs + +from mlx_audio.lm.sample_utils import apply_top_k as _apply_top_k_logprobs +from mlx_audio.lm.sample_utils import apply_top_p as _apply_top_p_logprobs def _mask_logits_from_logprobs(logits: mx.array, logprobs: mx.array) -> mx.array: diff --git a/mlx_audio/tts/models/outetts/outetts.py b/mlx_audio/tts/models/outetts/outetts.py index edbd7d41c..3e1c79596 100644 --- a/mlx_audio/tts/models/outetts/outetts.py +++ b/mlx_audio/tts/models/outetts/outetts.py @@ -8,17 +8,18 @@ import mlx.core as mx import mlx.nn as nn -from mlx_lm.generate import stream_generate -from mlx_lm.models.llama import Model as LlamaModel -from mlx_lm.models.llama import ModelArgs as LlamaModelConfig -from mlx_lm.models.qwen2 import Model as Qwen2Model -from mlx_lm.models.qwen2 import ModelArgs as Qwen2ModelConfig -from mlx_lm.models.qwen3 import Model as Qwen3Model -from mlx_lm.models.qwen3 import ModelArgs as Qwen3ModelConfig -from mlx_lm.sample_utils import make_logits_processors, make_sampler from tqdm import tqdm from transformers import AutoTokenizer +from mlx_audio.lm.generate import stream_generate +from mlx_audio.lm.models.llama import Model as LlamaModel +from mlx_audio.lm.models.llama import ModelArgs as LlamaModelConfig +from mlx_audio.lm.models.qwen2 import Model as Qwen2Model +from mlx_audio.lm.models.qwen2 import ModelArgs as Qwen2ModelConfig +from mlx_audio.lm.models.qwen3 import Model as Qwen3Model +from mlx_audio.lm.models.qwen3 import ModelArgs as Qwen3ModelConfig +from mlx_audio.lm.sample_utils import make_logits_processors, make_sampler + from ..base import GenerationResult from .audio_processor import AudioProcessor from .dac_interface import DacInterface diff --git a/mlx_audio/tts/models/pocket_tts/transformer.py b/mlx_audio/tts/models/pocket_tts/transformer.py index 4aa6b6c9b..30c40490b 100644 --- a/mlx_audio/tts/models/pocket_tts/transformer.py +++ b/mlx_audio/tts/models/pocket_tts/transformer.py @@ -1,6 +1,7 @@ import mlx.core as mx import mlx.nn as nn -from mlx_lm.models.cache import KVCache + +from mlx_audio.lm.models.cache import KVCache from .rope import RotaryEmbedding diff --git a/mlx_audio/tts/models/qwen3/qwen3.py b/mlx_audio/tts/models/qwen3/qwen3.py index 008e8d0a0..8382db1f3 100644 --- a/mlx_audio/tts/models/qwen3/qwen3.py +++ b/mlx_audio/tts/models/qwen3/qwen3.py @@ -4,13 +4,13 @@ from typing import List, Optional, Union import mlx.core as mx -from mlx_lm.generate import stream_generate -from mlx_lm.models.qwen3 import Model as Qwen3Model -from mlx_lm.models.qwen3 import ModelArgs as Qwen3ModelConfig -from mlx_lm.sample_utils import make_logits_processors, make_sampler from tqdm import tqdm from mlx_audio.codec.models.snac import SNAC +from mlx_audio.lm.generate import stream_generate +from mlx_audio.lm.models.qwen3 import Model as Qwen3Model +from mlx_audio.lm.models.qwen3 import ModelArgs as Qwen3ModelConfig +from mlx_audio.lm.sample_utils import make_logits_processors, make_sampler from mlx_audio.utils import load_audio from ..base import GenerationResult diff --git a/mlx_audio/tts/models/qwen3_tts/continuous_batching.py b/mlx_audio/tts/models/qwen3_tts/continuous_batching.py index 17426200f..3f3204bb1 100644 --- a/mlx_audio/tts/models/qwen3_tts/continuous_batching.py +++ b/mlx_audio/tts/models/qwen3_tts/continuous_batching.py @@ -7,8 +7,8 @@ from typing import Any, List, Optional, Tuple import mlx.core as mx -from mlx_lm.models.cache import BatchKVCache, KVCache +from mlx_audio.lm.models.cache import BatchKVCache, KVCache from mlx_audio.tts.continuous import TTSBatchEvent, TTSBatchItem, TTSBatchOptions diff --git a/mlx_audio/tts/models/qwen3_tts/qwen3_tts.py b/mlx_audio/tts/models/qwen3_tts/qwen3_tts.py index 148a5f1a6..66abb7111 100644 --- a/mlx_audio/tts/models/qwen3_tts/qwen3_tts.py +++ b/mlx_audio/tts/models/qwen3_tts/qwen3_tts.py @@ -8,15 +8,15 @@ import mlx.core as mx import mlx.nn as nn -from mlx_lm.sample_utils import ( +from tqdm import tqdm + +from mlx_audio.dsp import mel_filters, stft +from mlx_audio.lm.sample_utils import ( apply_min_p, apply_top_k, apply_top_p, categorical_sampling, ) -from tqdm import tqdm - -from mlx_audio.dsp import mel_filters, stft from mlx_audio.tts.continuous import TTSBatchItem, TTSBatchOptions from mlx_audio.tts.models.base import BatchGenerationResult, GenerationResult from mlx_audio.utils import load_audio diff --git a/mlx_audio/tts/models/qwen3_tts/speech_tokenizer.py b/mlx_audio/tts/models/qwen3_tts/speech_tokenizer.py index 2970e7504..f97e5bde3 100644 --- a/mlx_audio/tts/models/qwen3_tts/speech_tokenizer.py +++ b/mlx_audio/tts/models/qwen3_tts/speech_tokenizer.py @@ -8,7 +8,6 @@ import mlx.core as mx import mlx.nn as nn import numpy as np -from mlx_lm.models.cache import KVCache from mlx_audio.codec.models.mimi.mimi import _reset_kv_cache from mlx_audio.codec.models.mimi.modules import ( @@ -21,6 +20,7 @@ SplitResidualVectorQuantizer as MimiSplitRVQ, ) from mlx_audio.codec.models.mimi.modules import TransformerConfig +from mlx_audio.lm.models.cache import KVCache from .config import ( Qwen3TTSTokenizerConfig, diff --git a/mlx_audio/tts/models/qwen3_tts/talker.py b/mlx_audio/tts/models/qwen3_tts/talker.py index ec658b9eb..817b6a2bb 100644 --- a/mlx_audio/tts/models/qwen3_tts/talker.py +++ b/mlx_audio/tts/models/qwen3_tts/talker.py @@ -5,7 +5,8 @@ import mlx.core as mx import mlx.nn as nn -from mlx_lm.models.cache import KVCache + +from mlx_audio.lm.models.cache import KVCache from .config import Qwen3TTSTalkerCodePredictorConfig, Qwen3TTSTalkerConfig diff --git a/mlx_audio/tts/models/sesame/attention.py b/mlx_audio/tts/models/sesame/attention.py index d4f1a6077..2321627f5 100644 --- a/mlx_audio/tts/models/sesame/attention.py +++ b/mlx_audio/tts/models/sesame/attention.py @@ -3,8 +3,9 @@ import mlx.core as mx from mlx import nn -from mlx_lm.models.base import scaled_dot_product_attention -from mlx_lm.models.llama import ModelArgs + +from mlx_audio.lm.models.base import scaled_dot_product_attention +from mlx_audio.lm.models.llama import ModelArgs class Llama3ScaledRoPE(nn.Module): diff --git a/mlx_audio/tts/models/sesame/sesame.py b/mlx_audio/tts/models/sesame/sesame.py index f578951a3..438b08d6e 100644 --- a/mlx_audio/tts/models/sesame/sesame.py +++ b/mlx_audio/tts/models/sesame/sesame.py @@ -9,16 +9,16 @@ import mlx.core as mx import mlx.nn as nn from huggingface_hub import hf_hub_download -from mlx_lm.models.cache import make_prompt_cache -from mlx_lm.models.llama import LlamaModel -from mlx_lm.models.llama import ModelArgs as LlamaModelArgs -from mlx_lm.sample_utils import make_sampler from tokenizers.processors import TemplateProcessing from tqdm import tqdm from transformers import AutoTokenizer from mlx_audio.audio_io import read as audio_read from mlx_audio.codec.models.mimi import Mimi, MimiStreamingDecoder +from mlx_audio.lm.models.cache import make_prompt_cache +from mlx_audio.lm.models.llama import LlamaModel +from mlx_audio.lm.models.llama import ModelArgs as LlamaModelArgs +from mlx_audio.lm.sample_utils import make_sampler from mlx_audio.utils import load_audio, resample_audio from ..base import GenerationResult diff --git a/mlx_audio/tts/models/soprano/soprano.py b/mlx_audio/tts/models/soprano/soprano.py index f2cf1e86d..bf33473d7 100644 --- a/mlx_audio/tts/models/soprano/soprano.py +++ b/mlx_audio/tts/models/soprano/soprano.py @@ -10,13 +10,14 @@ import mlx.core as mx import mlx.nn as nn from huggingface_hub import snapshot_download -from mlx_lm.models.base import create_attention_mask -from mlx_lm.models.cache import KVCache -from mlx_lm.models.qwen3 import ModelArgs as Qwen3ModelConfig -from mlx_lm.models.qwen3 import Qwen3Model -from mlx_lm.sample_utils import make_sampler from transformers import AutoTokenizer +from mlx_audio.lm.models.base import create_attention_mask +from mlx_audio.lm.models.cache import KVCache +from mlx_audio.lm.models.qwen3 import ModelArgs as Qwen3ModelConfig +from mlx_audio.lm.models.qwen3 import Qwen3Model +from mlx_audio.lm.sample_utils import make_sampler + from ..base import BaseModelArgs, GenerationResult from .decoder import SopranoDecoder from .text import clean_text diff --git a/mlx_audio/tts/models/spark/spark.py b/mlx_audio/tts/models/spark/spark.py index f44f139f0..377a02f66 100644 --- a/mlx_audio/tts/models/spark/spark.py +++ b/mlx_audio/tts/models/spark/spark.py @@ -6,11 +6,11 @@ import mlx.core as mx import mlx.nn as nn -from mlx_lm.generate import stream_generate -from mlx_lm.models.qwen2 import Model as Qwen2Model -from mlx_lm.sample_utils import make_logits_processors, make_sampler from tqdm import tqdm +from mlx_audio.lm.generate import stream_generate +from mlx_audio.lm.models.qwen2 import Model as Qwen2Model +from mlx_audio.lm.sample_utils import make_logits_processors, make_sampler from mlx_audio.tts.models.base import BaseModelArgs, GenerationResult from .audio_tokenizer import BiCodecTokenizer diff --git a/mlx_audio/tts/models/voxtral_tts/voxtral_tts.py b/mlx_audio/tts/models/voxtral_tts/voxtral_tts.py index 6bb2b194b..d8353eb93 100644 --- a/mlx_audio/tts/models/voxtral_tts/voxtral_tts.py +++ b/mlx_audio/tts/models/voxtral_tts/voxtral_tts.py @@ -231,8 +231,8 @@ class MistralBackbone(nn.Module): def __init__(self, config: ModelConfig): super().__init__() - from mlx_lm.models.llama import Model as LlamaFullModel - from mlx_lm.models.llama import ModelArgs + from mlx_audio.lm.models.llama import Model as LlamaFullModel + from mlx_audio.lm.models.llama import ModelArgs lm_args = ModelArgs( model_type="llama", @@ -592,7 +592,7 @@ def generate( intermediate results have ``is_streaming_chunk=True`` and the last result additionally has ``is_final_chunk=True``. """ - from mlx_lm.models.cache import make_prompt_cache + from mlx_audio.lm.models.cache import make_prompt_cache if self.tokenizer is None: raise RuntimeError( diff --git a/mlx_audio/tts/models/zonos2/model.py b/mlx_audio/tts/models/zonos2/model.py index c94a6220e..9b0dfc7d2 100644 --- a/mlx_audio/tts/models/zonos2/model.py +++ b/mlx_audio/tts/models/zonos2/model.py @@ -7,9 +7,10 @@ import mlx.core as mx import mlx.nn as nn -from mlx_lm.models.base import create_attention_mask, scaled_dot_product_attention -from mlx_lm.models.cache import BatchKVCache, make_prompt_cache -from mlx_lm.models.switch_layers import SwitchGLU + +from mlx_audio.lm.models.base import create_attention_mask, scaled_dot_product_attention +from mlx_audio.lm.models.cache import BatchKVCache, make_prompt_cache +from mlx_audio.lm.models.switch_layers import SwitchGLU from ..base import BatchGenerationResult, GenerationResult from .config import Zonos2Config diff --git a/mlx_audio/tts/tests/test_models.py b/mlx_audio/tts/tests/test_models.py index a583ac891..33cb2f728 100644 --- a/mlx_audio/tts/tests/test_models.py +++ b/mlx_audio/tts/tests/test_models.py @@ -3668,9 +3668,8 @@ def test_convert_campplus_onnx_to_safetensors_allclose(self): ) def test_qwen2_sliding_window_attention_applies_after_window_boundary(self): - from mlx_lm.models.base import create_attention_mask - from mlx_lm.models.qwen2 import ModelArgs as Qwen2ModelArgs - + from mlx_audio.lm.models.base import create_attention_mask + from mlx_audio.lm.models.qwen2 import ModelArgs as Qwen2ModelArgs from mlx_audio.tts.models.bailingmm.bailingmm import MingQwen2Model args = Qwen2ModelArgs( diff --git a/mlx_audio/tts/utils.py b/mlx_audio/tts/utils.py index 03a80676b..b180f5ff9 100644 --- a/mlx_audio/tts/utils.py +++ b/mlx_audio/tts/utils.py @@ -222,8 +222,13 @@ def convert( quant_predicate: Optional[str] = None, q_mode: str = "affine", ): - from mlx_lm.convert import mixed_quant_predicate_builder - from mlx_lm.utils import dequantize_model, quantize_model, save_config, save_model + from mlx_audio.lm.convert import ( + dequantize_model, + mixed_quant_predicate_builder, + quantize_model, + save_config, + save_model, + ) print("[INFO] Loading") model_path = get_model_path(hf_path, revision=revision) diff --git a/pyproject.toml b/pyproject.toml index b4d0fe221..cb5aa1159 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,7 +21,6 @@ requires-python = ">=3.10" dependencies = [ "huggingface_hub>=1.0", "miniaudio>=1.61", - "mlx-lm>=0.31.1", "mlx>=0.31.1", "numpy>=1.26.4", "scipy>=1.10.0", @@ -54,7 +53,11 @@ server = [ # STS (Speech-to-Speech) dependencies sts = [ "sentencepiece>=0.2.0", - + + # The speech-to-speech pipeline's default responder loads an arbitrary chat + # model, which is mlx-lm's job; every other use is vendored in mlx_audio/lm. + "mlx-lm>=0.31.1", + "webrtcvad>=2.0.10", "setuptools<81", # pinned because >=81 drops pkg_resources which breaks webrtcvad ] @@ -64,6 +67,8 @@ all = [ "mistral-common[audio]", "sentencepiece>=0.2.0", + "mlx-lm>=0.31.1", + "fastapi>=0.95.0", "uvicorn[standard]>=0.22.0", @@ -71,6 +76,17 @@ all = [ "setuptools<81", # pinned because >=81 drops pkg_resources which breaks webrtcvad ] +# Optional in-process LLM responder for the speech-to-speech pipeline. Every +# other mlx-lm use is vendored under mlx_audio/lm. +llm = [ + "mlx-lm>=0.31.1", +] + +# Differential tests against the mlx-lm version mlx_audio/lm was vendored from. +parity = [ + "mlx-lm==0.31.3", +] + # Development dependencies dev = [ "pytest>=7.0.0", diff --git a/tests/vendor_parity/conftest.py b/tests/vendor_parity/conftest.py new file mode 100644 index 000000000..edc231734 --- /dev/null +++ b/tests/vendor_parity/conftest.py @@ -0,0 +1,17 @@ +import mlx.core as mx +import pytest + + +@pytest.fixture(autouse=True) +def _seeded(): + mx.random.seed(0) + + +def assert_exactly_equal(got, want): + assert got.dtype == want.dtype + assert got.shape == want.shape + assert mx.array_equal(got, want).item() + + +def rng_fingerprint(n=8): + return mx.random.uniform(shape=(n,)) diff --git a/tests/vendor_parity/test_parity_base.py b/tests/vendor_parity/test_parity_base.py new file mode 100644 index 000000000..212be4e04 --- /dev/null +++ b/tests/vendor_parity/test_parity_base.py @@ -0,0 +1,99 @@ +import dataclasses +import itertools +from dataclasses import dataclass + +import mlx.core as mx +import pytest +from conftest import assert_exactly_equal + +upstream = pytest.importorskip("mlx_lm.models.base") + +from mlx_audio.lm.models import base as vendored + +SHAPES = list( + itertools.product( + [1, 2, 5, 8, 17], # N + [0, 1, 7, 64], # offset + [None, 1, 4, 8], # window_size + ) +) + + +@pytest.mark.parametrize("n,offset,window", SHAPES) +def test_create_causal_mask_matches_upstream(n, offset, window): + kwargs = dict(offset=offset, window_size=window) + assert_exactly_equal( + vendored.create_causal_mask(n, **kwargs), + upstream.create_causal_mask(n, **kwargs), + ) + + +@pytest.mark.parametrize("left", [None, mx.array([0]), mx.array([2, 5])]) +@pytest.mark.parametrize("right", [None, mx.array([0]), mx.array([3, 0])]) +def test_create_causal_mask_padding_matches_upstream(left, right): + kwargs = dict(offset=3, left_padding=left, right_padding=right) + assert_exactly_equal( + vendored.create_causal_mask(6, **kwargs), + upstream.create_causal_mask(6, **kwargs), + ) + + +@pytest.mark.parametrize("n", [1, 2, 8, 17]) +@pytest.mark.parametrize("window", [None, 4, 32]) +@pytest.mark.parametrize("return_array", [False, True]) +def test_create_attention_mask_returns_identical_type(n, window, return_array): + h = mx.zeros((1, n, 8)) + got = vendored.create_attention_mask( + h, window_size=window, return_array=return_array + ) + want = upstream.create_attention_mask( + h, window_size=window, return_array=return_array + ) + assert type(got) is type(want) + if isinstance(want, mx.array): + assert_exactly_equal(got, want) + else: + assert got == want + + +def test_create_ssm_mask_matches_upstream(): + h = mx.zeros((1, 5, 8)) + assert vendored.create_ssm_mask(h) == upstream.create_ssm_mask(h) + + +@pytest.mark.parametrize("dtype", [mx.float32, mx.bfloat16]) +@pytest.mark.parametrize("mask_kind", [None, "causal", "bool", "additive"]) +@pytest.mark.parametrize("n_repeats", [1, 4]) +def test_scaled_dot_product_attention_matches_upstream(dtype, mask_kind, n_repeats): + b, n_kv, seq, d = 2, 2, 6, 16 + q = mx.random.normal((b, n_kv * n_repeats, seq, d)).astype(dtype) + k = mx.random.normal((b, n_kv, seq, d)).astype(dtype) + v = mx.random.normal((b, n_kv, seq, d)).astype(dtype) + if mask_kind == "bool": + mask = vendored.create_causal_mask(seq) + elif mask_kind == "additive": + mask = mx.where(vendored.create_causal_mask(seq), 0.0, -1e9).astype(dtype) + else: + mask = mask_kind + scale = d**-0.5 + assert_exactly_equal( + vendored.scaled_dot_product_attention(q, k, v, None, scale, mask), + upstream.scaled_dot_product_attention(q, k, v, None, scale, mask), + ) + + +def test_base_model_args_from_dict_matches_upstream(): + @dataclass + class V(vendored.BaseModelArgs): + a: int = 1 + b: str = "x" + + @dataclass + class U(upstream.BaseModelArgs): + a: int = 1 + b: str = "x" + + params = {"a": 7, "b": "y", "unknown": 3} + assert dataclasses.astuple(V.from_dict(params)) == dataclasses.astuple( + U.from_dict(params) + ) diff --git a/tests/vendor_parity/test_parity_cache.py b/tests/vendor_parity/test_parity_cache.py new file mode 100644 index 000000000..e5afb2bfd --- /dev/null +++ b/tests/vendor_parity/test_parity_cache.py @@ -0,0 +1,150 @@ +import mlx.core as mx +import pytest +from conftest import assert_exactly_equal + +upstream = pytest.importorskip("mlx_lm.models.cache") + +from mlx_audio.lm.models import cache as vendored + +B, H, D = 2, 4, 8 + + +def kv(n, seed): + mx.random.seed(seed) + return mx.random.normal((B, H, n, D)), mx.random.normal((B, H, n, D)) + + +def compare_state(a, b): + sa, sb = a.state, b.state + assert len(sa) == len(sb) + for x, y in zip(sa, sb): + if x is None or y is None: + assert x is y + else: + assert_exactly_equal(x, y) + if hasattr(a, "offset"): + assert_exactly_equal(mx.array(a.offset), mx.array(b.offset)) + + +def drive(vc, uc, steps): + for i, n in enumerate(steps): + k, val = kv(n, i) + vk, vv = vc.update_and_fetch(k, val) + uk, uv = uc.update_and_fetch(k, val) + assert_exactly_equal(vk, uk) + assert_exactly_equal(vv, uv) + compare_state(vc, uc) + + +@pytest.mark.parametrize( + "steps", + [ + [1] * 8, + [7, 1, 1, 1], + [256, 1, 1], + [255, 2, 1], + [512], + ], +) +def test_kvcache_update_and_fetch_sequence_matches_upstream(steps): + drive(vendored.KVCache(), upstream.KVCache(), steps) + + +@pytest.mark.parametrize("max_size", [4, 8, 256]) +@pytest.mark.parametrize("keep", [0, 1, 4]) +@pytest.mark.parametrize( + "steps", + [ + [1, 1, 1, 1], + [1] * 12, + [16], + [6, 1, 1, 1, 1, 1], + ], +) +def test_rotating_kvcache_wraparound_matches_upstream(max_size, keep, steps): + if keep >= max_size: + pytest.skip("keep must be smaller than max_size") + drive( + vendored.RotatingKVCache(max_size=max_size, keep=keep), + upstream.RotatingKVCache(max_size=max_size, keep=keep), + steps, + ) + + +@pytest.mark.parametrize("max_size", [4, 8]) +@pytest.mark.parametrize("n", [1, 3, 9]) +def test_rotating_kvcache_make_mask_matches_upstream(max_size, n): + vc = vendored.RotatingKVCache(max_size=max_size, keep=0) + uc = upstream.RotatingKVCache(max_size=max_size, keep=0) + drive(vc, uc, [1] * 6) + got, want = vc.make_mask(n), uc.make_mask(n) + assert type(got) is type(want) + if isinstance(want, mx.array): + assert_exactly_equal(got, want) + else: + assert got == want + + +@pytest.mark.parametrize("wrapped", [False, True]) +def test_batch_kvcache_merge_matches_upstream(wrapped): + """Mirrors higgs_audio_v3 continuous batching: merge extracted single rows.""" + steps = [4, 1, 1] if wrapped else [2] + vc = vendored.BatchKVCache([0, 3]) + uc = upstream.BatchKVCache([0, 3]) + mx.random.seed(0) + for n in steps: + k, val = mx.random.normal((2, H, n, D)), mx.random.normal((2, H, n, D)) + vc.update_and_fetch(k, val) + uc.update_and_fetch(k, val) + vm = vendored.BatchKVCache.merge([vc.extract(i) for i in range(2)]) + um = upstream.BatchKVCache.merge([uc.extract(i) for i in range(2)]) + assert_exactly_equal(vm.keys, um.keys) + assert_exactly_equal(vm.values, um.values) + assert_exactly_equal(mx.array(vm.offset), mx.array(um.offset)) + assert_exactly_equal(mx.array(vm.left_padding), mx.array(um.left_padding)) + + +def test_kvcache_merge_matches_upstream(): + vcs = [vendored.KVCache() for _ in range(2)] + ucs = [upstream.KVCache() for _ in range(2)] + for vc, uc in zip(vcs, ucs): + k, val = kv(3, 1) + vc.update_and_fetch(k[:1], val[:1]) + uc.update_and_fetch(k[:1], val[:1]) + vm, um = vendored.KVCache.merge(vcs), upstream.KVCache.merge(ucs) + assert_exactly_equal(vm.keys, um.keys) + assert_exactly_equal(mx.array(vm.offset), mx.array(um.offset)) + + +def test_arrays_cache_matches_upstream(): + vc, uc = vendored.ArraysCache(size=2), upstream.ArraysCache(size=2) + x = mx.random.normal((B, 3, D)) + vc[0], uc[0] = x, x + assert_exactly_equal(vc[0], uc[0]) + compare_state(vc, uc) + + +class _Toy: + def __init__(self, n): + self.layers = list(range(n)) + + +def test_make_prompt_cache_matches_upstream(): + got = vendored.make_prompt_cache(_Toy(3)) + want = upstream.make_prompt_cache(_Toy(3)) + assert len(got) == len(want) + assert [type(c).__name__ for c in got] == [type(c).__name__ for c in want] + + +def test_make_prompt_cache_with_max_kv_size_matches_upstream(): + got = vendored.make_prompt_cache(_Toy(3), max_kv_size=16) + want = upstream.make_prompt_cache(_Toy(3), max_kv_size=16) + assert [type(c).__name__ for c in got] == [type(c).__name__ for c in want] + + +def test_state_roundtrip_matches_upstream(): + vc, uc = vendored.KVCache(), upstream.KVCache() + drive(vc, uc, [5, 1]) + vc2 = vendored.KVCache.from_state(vc.state, vc.meta_state) + uc2 = upstream.KVCache.from_state(uc.state, uc.meta_state) + compare_state(vc2, uc2) diff --git a/tests/vendor_parity/test_parity_convert.py b/tests/vendor_parity/test_parity_convert.py new file mode 100644 index 000000000..47e76f6fa --- /dev/null +++ b/tests/vendor_parity/test_parity_convert.py @@ -0,0 +1,112 @@ +import json + +import mlx.core as mx +import mlx.nn as nn +import pytest +from conftest import assert_exactly_equal + +upstream_utils = pytest.importorskip("mlx_lm.utils") +upstream_convert = pytest.importorskip("mlx_lm.convert") + +from mlx_audio.lm import convert as vendored + + +class Tiny(nn.Module): + def __init__(self, layers=8): + super().__init__() + self.layers = [ + nn.Sequential(nn.Linear(64, 64), nn.Linear(64, 64)) for _ in range(layers) + ] + self.lm_head = nn.Linear(64, 128) + + def __call__(self, x): + return x + + +class Deep(nn.Module): + """Named so paths contain down_proj / v_proj / lm_head like a real LM.""" + + def __init__(self, layers=8): + super().__init__() + self.layers = [Deep._Block() for _ in range(layers)] + self.lm_head = nn.Linear(64, 128) + + class _Block(nn.Module): + def __init__(self): + super().__init__() + self.v_proj = nn.Linear(64, 64) + self.down_proj = nn.Linear(64, 64) + + def __call__(self, x): + return x + + +def test_quantize_model_matches_upstream(): + v, u = Tiny(), Tiny() + mx.eval(v.parameters()) + u.update(v.parameters()) + cfg = {"model_type": "tiny"} + vm, vc = vendored.quantize_model(v, cfg, 64, 4) + um, uc = upstream_utils.quantize_model(u, cfg, 64, 4) + assert vc == uc + from mlx.utils import tree_flatten + + vw, uw = dict(tree_flatten(vm.parameters())), dict(tree_flatten(um.parameters())) + assert vw.keys() == uw.keys() + for k in vw: + assert_exactly_equal(vw[k], uw[k]) + + +@pytest.mark.parametrize("recipe", ["mixed_2_6", "mixed_3_4", "mixed_3_6", "mixed_4_6"]) +@pytest.mark.parametrize("layers", [4, 8, 32]) +def test_mixed_quant_predicate_matches_upstream(recipe, layers): + model = Deep(layers) + v = vendored.mixed_quant_predicate_builder(recipe, model) + u = upstream_convert.mixed_quant_predicate_builder(recipe, model) + for name, module in model.named_modules(): + if not isinstance(module, nn.Linear): + continue + assert v(name, module) == u(name, module), name + + +def test_save_model_matches_upstream(tmp_path): + v, u = Tiny(), Tiny() + mx.eval(v.parameters()) + u.update(v.parameters()) + mx.eval(u.parameters()) + + vendored.save_model(tmp_path / "v", v) + upstream_utils.save_model(tmp_path / "u", u) + + vi = json.loads((tmp_path / "v" / "model.safetensors.index.json").read_text()) + ui = json.loads((tmp_path / "u" / "model.safetensors.index.json").read_text()) + assert vi["weight_map"] == ui["weight_map"] + assert vi["metadata"]["total_size"] == ui["metadata"]["total_size"] + + vw = mx.load(str(tmp_path / "v" / "model.safetensors")) + uw = mx.load(str(tmp_path / "u" / "model.safetensors")) + assert vw.keys() == uw.keys() + for k in vw: + assert_exactly_equal(vw[k], uw[k]) + + +def test_save_model_shards_like_upstream(): + weights = {f"w{i}": mx.zeros((1024, 1024)) for i in range(4)} + got = vendored.make_shards(weights, max_file_size_gb=1) + want = upstream_utils.make_shards(weights, max_file_size_gb=1) + assert [sorted(s) for s in got] == [sorted(s) for s in want] + + +def test_save_config_matches_upstream(tmp_path): + cfg = { + "b": 2, + "a": 1, + "_name_or_path": "x", + "vision_config": {"y": 1}, + "quantization": {"group_size": 64, "bits": 4}, + } + vendored.save_config(dict(cfg), tmp_path / "v.json") + upstream_utils.save_config(dict(cfg), tmp_path / "u.json") + assert json.loads((tmp_path / "v.json").read_text()) == json.loads( + (tmp_path / "u.json").read_text() + ) diff --git a/tests/vendor_parity/test_parity_gemma3.py b/tests/vendor_parity/test_parity_gemma3.py new file mode 100644 index 000000000..654486a22 --- /dev/null +++ b/tests/vendor_parity/test_parity_gemma3.py @@ -0,0 +1,56 @@ +from dataclasses import asdict + +import pytest +from mlx.utils import tree_flatten + +upstream_gemma3 = pytest.importorskip("mlx_lm.models.gemma3") +upstream_gemma3_text = pytest.importorskip("mlx_lm.models.gemma3_text") + +from mlx_audio.lm.models import gemma3 as vendored_gemma3 +from mlx_audio.lm.models import gemma3_text as vendored_gemma3_text + + +def parameter_shapes(model): + return { + name: parameter.shape for name, parameter in tree_flatten(model.parameters()) + } + + +def text_args(module): + return module.ModelArgs( + model_type="gemma3_text", + hidden_size=16, + num_hidden_layers=1, + intermediate_size=32, + num_attention_heads=2, + head_dim=8, + vocab_size=64, + num_key_value_heads=2, + sliding_window=8, + sliding_window_pattern=1, + ) + + +def gemma3_args(module): + return module.ModelArgs( + model_type="gemma3", + vocab_size=64, + text_config=asdict(text_args(vendored_gemma3_text)), + ) + + +@pytest.mark.parametrize( + ("vendored", "upstream", "make_args"), + [ + (vendored_gemma3_text, upstream_gemma3_text, text_args), + (vendored_gemma3, upstream_gemma3, gemma3_args), + ], +) +def test_config_and_parameter_keys_match_upstream(vendored, upstream, make_args): + vendored_args = make_args(vendored) + upstream_args = make_args(upstream) + + assert asdict(vendored_args) == asdict(upstream_args) + assert parameter_shapes(vendored.Model(vendored_args)) == parameter_shapes( + upstream.Model(upstream_args) + ) diff --git a/tests/vendor_parity/test_parity_generate.py b/tests/vendor_parity/test_parity_generate.py new file mode 100644 index 000000000..0d546606e --- /dev/null +++ b/tests/vendor_parity/test_parity_generate.py @@ -0,0 +1,138 @@ +import mlx.core as mx +import mlx.nn as nn +import pytest +from conftest import assert_exactly_equal + +upstream = pytest.importorskip("mlx_lm.generate") +upstream_llama = pytest.importorskip("mlx_lm.models.llama") +upstream_sample = pytest.importorskip("mlx_lm.sample_utils") + +from mlx_audio.lm import generate as vendored +from mlx_audio.lm import sample_utils as vendored_sample +from mlx_audio.lm.models import llama as vendored_llama + +ARGS = dict( + model_type="llama", + hidden_size=64, + num_hidden_layers=2, + intermediate_size=128, + num_attention_heads=4, + num_key_value_heads=2, + rms_norm_eps=1e-5, + vocab_size=128, +) + + +def paired_models(): + """Two structurally identical llama models sharing the same weights.""" + v = vendored_llama.Model(vendored_llama.ModelArgs(**ARGS)) + u = upstream_llama.Model(upstream_llama.ModelArgs(**ARGS)) + mx.eval(v.parameters()) + u.update(v.parameters()) + mx.eval(u.parameters()) + return v, u + + +def stream(mod, model, sampler=None, processors=None, **kw): + mx.random.seed(7) + return [ + (int(tok), lp) + for tok, lp in mod.generate_step( + mx.array([3, 9, 14, 2, 8]), + model, + max_tokens=16, + sampler=sampler, + logits_processors=processors, + **kw, + ) + ] + + +def compare(got, want): + assert [t for t, _ in got] == [t for t, _ in want], "token streams diverge" + for (_, g), (_, w) in zip(got, want): + assert_exactly_equal(g, w) + + +def test_greedy_stream_matches_upstream(): + v, u = paired_models() + compare(stream(vendored, v), stream(upstream, u)) + + +@pytest.mark.parametrize("prefill", [1, 2, 4, 4096]) +def test_prefill_chunking_matches_upstream(prefill): + v, u = paired_models() + compare( + stream(vendored, v, prefill_step_size=prefill), + stream(upstream, u, prefill_step_size=prefill), + ) + + +def test_sampler_stream_matches_upstream(): + v, u = paired_models() + compare( + stream(vendored, v, sampler=vendored_sample.make_sampler(temp=0.8, top_p=0.9)), + stream(upstream, u, sampler=upstream_sample.make_sampler(temp=0.8, top_p=0.9)), + ) + + +def test_logits_processors_stream_matches_upstream(): + v, u = paired_models() + compare( + stream( + vendored, + v, + processors=vendored_sample.make_logits_processors(repetition_penalty=1.2), + ), + stream( + upstream, + u, + processors=upstream_sample.make_logits_processors(repetition_penalty=1.2), + ), + ) + + +def test_rotating_cache_stream_matches_upstream(): + """max_kv_size forces RotatingKVCache, whose wraparound is the subtlest path.""" + v, u = paired_models() + compare( + stream(vendored, v, max_kv_size=8), + stream(upstream, u, max_kv_size=8), + ) + + +def test_vendored_loop_against_upstream_model(): + """Cross control: isolates a loop bug from a backbone bug.""" + v, u = paired_models() + compare(stream(vendored, u), stream(upstream, u)) + + +def test_max_tokens_boundary_matches_upstream(): + v, u = paired_models() + for n in (1, 2, 5): + got = [ + int(t) for t, _ in vendored.generate_step(mx.array([3, 9]), v, max_tokens=n) + ] + want = [ + int(t) for t, _ in upstream.generate_step(mx.array([3, 9]), u, max_tokens=n) + ] + assert got == want and len(got) == n + + +def test_input_embeddings_stream_matches_upstream(): + v, u = paired_models() + mx.random.seed(3) + emb = mx.random.normal((5, ARGS["hidden_size"])) + got = [ + int(t) + for t, _ in vendored.generate_step( + mx.array([], dtype=mx.int32), v, max_tokens=8, input_embeddings=emb + ) + ] + want = [ + int(t) + for t, _ in upstream.generate_step( + mx.array([], dtype=mx.int32), u, max_tokens=8, input_embeddings=emb + ) + ] + assert got == want diff --git a/tests/vendor_parity/test_parity_gpt2_granite.py b/tests/vendor_parity/test_parity_gpt2_granite.py new file mode 100644 index 000000000..c5a91c10d --- /dev/null +++ b/tests/vendor_parity/test_parity_gpt2_granite.py @@ -0,0 +1,67 @@ +from dataclasses import asdict + +import pytest +from mlx.utils import tree_flatten + +upstream_gpt2 = pytest.importorskip("mlx_lm.models.gpt2") +upstream_granite = pytest.importorskip("mlx_lm.models.granite") + +from mlx_audio.lm.models import gpt2 as vendored_gpt2 +from mlx_audio.lm.models import granite as vendored_granite + + +def parameter_shapes(model): + return { + name: parameter.shape for name, parameter in tree_flatten(model.parameters()) + } + + +def gpt2_args(module): + return module.ModelArgs( + model_type="gpt2", + n_ctx=128, + n_embd=16, + n_head=2, + n_layer=1, + n_positions=128, + layer_norm_epsilon=1e-5, + vocab_size=64, + ) + + +def granite_args(module): + return module.ModelArgs( + model_type="granite", + hidden_size=16, + num_hidden_layers=1, + intermediate_size=32, + num_attention_heads=2, + num_key_value_heads=2, + rms_norm_eps=1e-5, + vocab_size=64, + logits_scaling=1.0, + attention_multiplier=16**-0.5, + embedding_multiplier=1.0, + residual_multiplier=1.0, + max_position_embeddings=128, + attention_bias=False, + mlp_bias=False, + rope_theta=10000.0, + ) + + +@pytest.mark.parametrize( + ("vendored", "upstream", "make_args"), + [ + (vendored_gpt2, upstream_gpt2, gpt2_args), + (vendored_granite, upstream_granite, granite_args), + ], +) +def test_config_and_parameter_keys_match_upstream(vendored, upstream, make_args): + vendored_args = make_args(vendored) + upstream_args = make_args(upstream) + + assert asdict(vendored_args) == asdict(upstream_args) + assert parameter_shapes(vendored.Model(vendored_args)) == parameter_shapes( + upstream.Model(upstream_args) + ) diff --git a/tests/vendor_parity/test_parity_lfm2_bailing.py b/tests/vendor_parity/test_parity_lfm2_bailing.py new file mode 100644 index 000000000..e277ecb85 --- /dev/null +++ b/tests/vendor_parity/test_parity_lfm2_bailing.py @@ -0,0 +1,76 @@ +from dataclasses import asdict + +import pytest +from mlx.utils import tree_flatten + +upstream_bailing = pytest.importorskip("mlx_lm.models.bailing_moe") +upstream_lfm2 = pytest.importorskip("mlx_lm.models.lfm2") + +from mlx_audio.lm.models import bailing_moe as vendored_bailing +from mlx_audio.lm.models import lfm2 as vendored_lfm2 + + +def parameter_shapes(model): + return { + name: parameter.shape for name, parameter in tree_flatten(model.parameters()) + } + + +def lfm2_args(module): + return module.ModelArgs( + model_type="lfm2", + vocab_size=64, + hidden_size=16, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + max_position_embeddings=128, + norm_eps=1e-5, + conv_bias=False, + conv_L_cache=4, + block_dim=16, + block_ff_dim=32, + block_multiple_of=8, + block_ffn_dim_multiplier=1.0, + block_auto_adjust_ff_dim=False, + full_attn_idxs=[0], + layer_types=["full_attention"], + ) + + +def bailing_args(module): + return module.ModelArgs( + model_type="bailing_moe", + hidden_size=16, + intermediate_size=32, + max_position_embeddings=128, + moe_intermediate_size=16, + num_experts=2, + num_shared_experts=0, + norm_topk_prob=True, + num_attention_heads=2, + num_experts_per_tok=1, + num_hidden_layers=1, + num_key_value_heads=2, + rms_norm_eps=1e-5, + rope_theta=10000.0, + vocab_size=64, + first_k_dense_replace=1, + ) + + +@pytest.mark.parametrize( + ("vendored", "upstream", "make_args"), + [ + (vendored_lfm2, upstream_lfm2, lfm2_args), + (vendored_bailing, upstream_bailing, bailing_args), + ], +) +def test_config_and_parameter_keys_match_upstream(vendored, upstream, make_args): + vendored_args = make_args(vendored) + upstream_args = make_args(upstream) + + assert asdict(vendored_args) == asdict(upstream_args) + assert parameter_shapes(vendored.Model(vendored_args)) == parameter_shapes( + upstream.Model(upstream_args) + ) diff --git a/tests/vendor_parity/test_parity_llama.py b/tests/vendor_parity/test_parity_llama.py new file mode 100644 index 000000000..06ea3d515 --- /dev/null +++ b/tests/vendor_parity/test_parity_llama.py @@ -0,0 +1,36 @@ +from dataclasses import asdict + +import pytest +from mlx.utils import tree_flatten + +upstream = pytest.importorskip("mlx_lm.models.llama") + +from mlx_audio.lm.models import llama as vendored + + +def make_args(module): + return module.ModelArgs( + model_type="llama", + hidden_size=16, + num_hidden_layers=1, + intermediate_size=32, + num_attention_heads=2, + rms_norm_eps=1e-5, + vocab_size=64, + ) + + +def parameter_shapes(model): + return { + name: parameter.shape for name, parameter in tree_flatten(model.parameters()) + } + + +def test_llama_config_and_parameter_keys_match_upstream(): + vendored_args = make_args(vendored) + upstream_args = make_args(upstream) + + assert asdict(vendored_args) == asdict(upstream_args) + assert parameter_shapes(vendored.Model(vendored_args)) == parameter_shapes( + upstream.Model(upstream_args) + ) diff --git a/tests/vendor_parity/test_parity_qwen.py b/tests/vendor_parity/test_parity_qwen.py new file mode 100644 index 000000000..ddbe024e9 --- /dev/null +++ b/tests/vendor_parity/test_parity_qwen.py @@ -0,0 +1,63 @@ +from dataclasses import asdict + +import pytest +from mlx.utils import tree_flatten + +upstream_qwen2 = pytest.importorskip("mlx_lm.models.qwen2") +upstream_qwen3 = pytest.importorskip("mlx_lm.models.qwen3") + +from mlx_audio.lm.models import qwen2 as vendored_qwen2 +from mlx_audio.lm.models import qwen3 as vendored_qwen3 + + +def parameter_shapes(model): + return { + name: parameter.shape for name, parameter in tree_flatten(model.parameters()) + } + + +def qwen2_args(module): + return module.ModelArgs( + model_type="qwen2", + hidden_size=16, + num_hidden_layers=1, + intermediate_size=32, + num_attention_heads=2, + num_key_value_heads=2, + rms_norm_eps=1e-5, + vocab_size=64, + ) + + +def qwen3_args(module): + return module.ModelArgs( + model_type="qwen3", + hidden_size=16, + num_hidden_layers=1, + intermediate_size=32, + num_attention_heads=2, + num_key_value_heads=2, + rms_norm_eps=1e-5, + vocab_size=64, + max_position_embeddings=128, + rope_theta=1000000.0, + head_dim=8, + tie_word_embeddings=True, + ) + + +@pytest.mark.parametrize( + ("vendored", "upstream", "make_args"), + [ + (vendored_qwen2, upstream_qwen2, qwen2_args), + (vendored_qwen3, upstream_qwen3, qwen3_args), + ], +) +def test_qwen_config_and_parameter_keys_match_upstream(vendored, upstream, make_args): + vendored_args = make_args(vendored) + upstream_args = make_args(upstream) + + assert asdict(vendored_args) == asdict(upstream_args) + assert parameter_shapes(vendored.Model(vendored_args)) == parameter_shapes( + upstream.Model(upstream_args) + ) diff --git a/tests/vendor_parity/test_parity_sampling.py b/tests/vendor_parity/test_parity_sampling.py new file mode 100644 index 000000000..d15bef9c5 --- /dev/null +++ b/tests/vendor_parity/test_parity_sampling.py @@ -0,0 +1,119 @@ +import mlx.core as mx +import pytest +from conftest import assert_exactly_equal, rng_fingerprint + +upstream = pytest.importorskip("mlx_lm.sample_utils") + +from mlx_audio.lm import sample_utils as vendored + +VOCAB = 128 + + +def logprobs(batch=1, seed=0, dtype=mx.float32): + mx.random.seed(seed) + x = mx.random.normal((batch, VOCAB)).astype(dtype) + return x - mx.logsumexp(x, axis=-1, keepdims=True) + + +def both(fn_name, x, *args, **kwargs): + """Call vendored and upstream under identical RNG state; compare result and state.""" + mx.random.seed(1234) + got = getattr(vendored, fn_name)(x, *args, **kwargs) + got_rng = rng_fingerprint() + mx.random.seed(1234) + want = getattr(upstream, fn_name)(x, *args, **kwargs) + want_rng = rng_fingerprint() + assert_exactly_equal(got, want) + assert_exactly_equal(got_rng, want_rng) + + +@pytest.mark.parametrize("k", [1, 5, 64, VOCAB - 1]) +@pytest.mark.parametrize("batch", [1, 3]) +@pytest.mark.parametrize("call_twice", [False, True]) +def test_apply_top_k_matches_upstream(k, batch, call_twice): + x = logprobs(batch) + both("apply_top_k", x, k) + if call_twice: + both("apply_top_k", x, k) + + +@pytest.mark.parametrize("p", [0.1, 0.5, 0.9, 1.0]) +@pytest.mark.parametrize("batch", [1, 3]) +@pytest.mark.parametrize("call_twice", [False, True]) +def test_apply_top_p_matches_upstream(p, batch, call_twice): + x = logprobs(batch) + both("apply_top_p", x, p) + if call_twice: + both("apply_top_p", x, p) + + +@pytest.mark.parametrize("p", [0.05, 0.5]) +def test_apply_min_p_matches_upstream(p): + both("apply_min_p", logprobs(2), p) + + +def test_apply_min_p_zero_fails_identically_to_upstream(): + """min_p=0 hits math.log(0) in both; kept as a deliberate non-divergence.""" + x = logprobs(2) + with pytest.raises(ValueError): + vendored.apply_min_p(x, 0.0) + with pytest.raises(ValueError): + upstream.apply_min_p(x, 0.0) + + +def test_apply_min_p_min_tokens_to_keep_diverges_from_upstream(): + """Deliberate fix: upstream passes a Python bool to mx.put_along_axis and + raises TypeError whenever min_p > 0 and min_tokens_to_keep > 1.""" + x = logprobs(2) + with pytest.raises(TypeError): + upstream.apply_min_p(x, 0.9, 5) + kept = (vendored.apply_min_p(x, 0.9, 5) != -mx.inf).sum(-1) + assert kept.tolist() == [5, 5] + + +SAMPLER_CASES = [ + dict(temp=0.0), + dict(temp=1.0), + dict(temp=0.7, top_p=0.9), + dict(temp=0.7, min_p=0.05), + dict(temp=0.7, top_k=10), + dict(temp=0.7, top_p=0.95, top_k=20, min_p=0.02), + dict(temp=0.8, xtc_probability=0.5, xtc_threshold=0.1), +] + + +@pytest.mark.parametrize("kwargs", SAMPLER_CASES) +def test_make_sampler_token_streams_match_upstream(kwargs): + v_sampler = vendored.make_sampler(**kwargs) + u_sampler = upstream.make_sampler(**kwargs) + for step in range(20): + x = logprobs(seed=step) + mx.random.seed(99) + got = v_sampler(x) + got_rng = rng_fingerprint() + mx.random.seed(99) + want = u_sampler(x) + want_rng = rng_fingerprint() + assert_exactly_equal(got, want) + assert_exactly_equal(got_rng, want_rng) + + +PROCESSOR_CASES = [ + dict(repetition_penalty=1.2), + dict(repetition_penalty=1.1, repetition_context_size=4), + dict(presence_penalty=0.5), + dict(frequency_penalty=0.3), + dict(logit_bias={5: 10.0, 7: -10.0}), + dict(repetition_penalty=1.15, presence_penalty=0.2, frequency_penalty=0.1), +] + + +@pytest.mark.parametrize("kwargs", PROCESSOR_CASES) +def test_make_logits_processors_matches_upstream(kwargs): + v_procs = vendored.make_logits_processors(**kwargs) + u_procs = upstream.make_logits_processors(**kwargs) + assert len(v_procs) == len(u_procs) + tokens = mx.array([3, 5, 5, 9, 3, 3, 1]) + x = logprobs() + for vp, up in zip(v_procs, u_procs): + assert_exactly_equal(vp(tokens, x), up(tokens, x)) diff --git a/tests/vendor_parity/test_parity_support.py b/tests/vendor_parity/test_parity_support.py new file mode 100644 index 000000000..df7971205 --- /dev/null +++ b/tests/vendor_parity/test_parity_support.py @@ -0,0 +1,73 @@ +import mlx.core as mx +import pytest +from conftest import assert_exactly_equal + +upstream_rope = pytest.importorskip("mlx_lm.models.rope_utils") +upstream_switch = pytest.importorskip("mlx_lm.models.switch_layers") +upstream_act = pytest.importorskip("mlx_lm.models.activations") + +from mlx_audio.lm.models import activations as vendored_act +from mlx_audio.lm.models import rope_utils as vendored_rope +from mlx_audio.lm.models import switch_layers as vendored_switch + +DIMS, HEADS, SEQ = 64, 2, 12 + +SCALINGS = [ + None, + {"rope_type": "linear", "factor": 2.0}, + { + "rope_type": "llama3", + "factor": 8.0, + "low_freq_factor": 1.0, + "high_freq_factor": 4.0, + "original_max_position_embeddings": 8192, + }, + {"rope_type": "yarn", "factor": 4.0, "original_max_position_embeddings": 4096}, +] + + +@pytest.mark.parametrize("scaling", SCALINGS) +@pytest.mark.parametrize("traditional", [False, True]) +@pytest.mark.parametrize("offset", [0, 5]) +def test_rope_matches_upstream(scaling, traditional, offset): + kwargs = dict( + dims=DIMS, base=10000.0, traditional=traditional, scaling_config=scaling + ) + v = vendored_rope.initialize_rope(**kwargs) + u = upstream_rope.initialize_rope(**kwargs) + mx.random.seed(0) + x = mx.random.normal((1, HEADS, SEQ, DIMS)) + assert_exactly_equal(v(x, offset=offset), u(x, offset=offset)) + + +def test_swiglu_matches_upstream(): + mx.random.seed(0) + gate, x = mx.random.normal((2, 16)), mx.random.normal((2, 16)) + assert_exactly_equal(vendored_act.swiglu(gate, x), upstream_act.swiglu(gate, x)) + + +@pytest.mark.parametrize("cls", ["SwitchLinear", "SwitchGLU"]) +def test_switch_layers_match_upstream(cls): + n_experts, in_dims, out_dims = 4, 16, 32 + if cls == "SwitchLinear": + v = vendored_switch.SwitchLinear(in_dims, out_dims, n_experts) + u = upstream_switch.SwitchLinear(in_dims, out_dims, n_experts) + else: + v = vendored_switch.SwitchGLU(in_dims, out_dims, n_experts) + u = upstream_switch.SwitchGLU(in_dims, out_dims, n_experts) + u.update(v.parameters()) + mx.random.seed(0) + x = mx.random.normal((2, 1, in_dims)) + idx = mx.array([[0, 2], [1, 3]]) + assert_exactly_equal(v(x, idx), u(x, idx)) + + +def test_quantized_switch_linear_matches_upstream(): + v = vendored_switch.SwitchLinear(64, 32, 4) + u = upstream_switch.SwitchLinear(64, 32, 4) + u.update(v.parameters()) + vq, uq = v.to_quantized(group_size=32), u.to_quantized(group_size=32) + mx.random.seed(0) + x = mx.random.normal((2, 1, 64)) + idx = mx.array([[0, 2], [1, 3]]) + assert_exactly_equal(vq(x, idx), uq(x, idx)) diff --git a/uv.lock b/uv.lock index ea66a9e5c..61c608357 100644 --- a/uv.lock +++ b/uv.lock @@ -968,7 +968,6 @@ dependencies = [ { name = "huggingface-hub" }, { name = "miniaudio" }, { name = "mlx" }, - { name = "mlx-lm" }, { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, { name = "numpy", version = "2.3.5", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, { name = "scipy", version = "1.15.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, @@ -998,6 +997,12 @@ docs = [ { name = "mkdocs-material" }, { name = "mkdocstrings", extra = ["python"] }, ] +llm = [ + { name = "mlx-lm" }, +] +parity = [ + { name = "mlx-lm" }, +] server = [ { name = "fastapi" }, { name = "python-multipart" }, @@ -1031,7 +1036,8 @@ requires-dist = [ { name = "mkdocs-material", marker = "extra == 'docs'", specifier = ">=9.6" }, { name = "mkdocstrings", extras = ["python"], marker = "extra == 'docs'", specifier = ">=0.29" }, { name = "mlx", specifier = ">=0.31.1" }, - { name = "mlx-lm", specifier = ">=0.31.1" }, + { name = "mlx-lm", marker = "extra == 'llm'", specifier = ">=0.31.1" }, + { name = "mlx-lm", marker = "extra == 'parity'", specifier = "==0.31.3" }, { name = "numpy", specifier = ">=1.26.4" }, { name = "pre-commit", marker = "extra == 'dev'", specifier = ">=3.7.0" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=7.0.0" }, @@ -1047,18 +1053,18 @@ requires-dist = [ { name = "setuptools", marker = "extra == 'sts'", specifier = "<81" }, { name = "sounddevice", specifier = ">=0.5.3" }, { name = "tqdm", specifier = ">=4.67.1" }, - { name = "transformers", specifier = ">=5.5.0,<5.13.0" }, + { name = "transformers", specifier = ">=5.14.0" }, { name = "uvicorn", extras = ["standard"], marker = "extra == 'all'", specifier = ">=0.22.0" }, { name = "uvicorn", extras = ["standard"], marker = "extra == 'server'", specifier = ">=0.22.0" }, { name = "webrtcvad", marker = "extra == 'all'", specifier = ">=2.0.10" }, { name = "webrtcvad", marker = "extra == 'server'", specifier = ">=2.0.10" }, { name = "webrtcvad", marker = "extra == 'sts'", specifier = ">=2.0.10" }, ] -provides-extras = ["stt", "tts", "server", "sts", "all", "dev", "docs"] +provides-extras = ["stt", "tts", "server", "sts", "all", "llm", "parity", "dev", "docs"] [[package]] name = "mlx-lm" -version = "0.31.1" +version = "0.31.3" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "jinja2" }, @@ -1070,9 +1076,9 @@ dependencies = [ { name = "sentencepiece" }, { name = "transformers" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/9e/f9/3f5597c62bd5733ebb3c9f96c33f2065db16353d743b8548bb05a01b7dd3/mlx_lm-0.31.1.tar.gz", hash = "sha256:1b2362ea301427004e5dda43b9241d751d4cb80eba641f6b85b29fc493affac5", size = 285473, upload-time = "2026-03-11T02:02:57.466Z" } +sdist = { url = "https://files.pythonhosted.org/packages/84/94/9a38d6b0c6fcca995b9136c94eb7da1e9c5165652edf228b96b29960fa7a/mlx_lm-0.31.3.tar.gz", hash = "sha256:61eb0e3ba09444f77f874aff295401d7ccd20b39495cbbce0c782a15474ce733", size = 304318, upload-time = "2026-04-22T07:37:27.922Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/3b/8e/2e40713fc673d7268e32222752919075c0024b02635af56531d9ee8317b5/mlx_lm-0.31.1-py3-none-any.whl", hash = "sha256:bfc3e08e919b87bebb6fe2dbea980ece8ae8ee7132b31421aa0923c4286ecfa0", size = 393971, upload-time = "2026-03-11T02:02:54.973Z" }, + { url = "https://files.pythonhosted.org/packages/90/02/9a67b8e4f87e3e2e5cd7b1ad79304b93c09a0db6af34bee75e6551c06c60/mlx_lm-0.31.3-py3-none-any.whl", hash = "sha256:758cfddf1180053b7613db76fad3d246a331a2a905808e1164a275621fc983b8", size = 408890, upload-time = "2026-04-22T07:37:25.965Z" }, ] [[package]] @@ -2082,28 +2088,26 @@ wheels = [ [[package]] name = "safetensors" -version = "0.7.0" +version = "0.8.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/29/9c/6e74567782559a63bd040a236edca26fd71bc7ba88de2ef35d75df3bca5e/safetensors-0.7.0.tar.gz", hash = "sha256:07663963b67e8bd9f0b8ad15bb9163606cd27cc5a1b96235a50d8369803b96b0", size = 200878, upload-time = "2025-11-19T15:18:43.199Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/fa/47/aef6c06649039accf914afef490268e1067ed82be62bcfa5b7e886ad15e8/safetensors-0.7.0-cp38-abi3-macosx_10_12_x86_64.whl", hash = "sha256:c82f4d474cf725255d9e6acf17252991c3c8aac038d6ef363a4bf8be2f6db517", size = 467781, upload-time = "2025-11-19T15:18:35.84Z" }, - { url = "https://files.pythonhosted.org/packages/e8/00/374c0c068e30cd31f1e1b46b4b5738168ec79e7689ca82ee93ddfea05109/safetensors-0.7.0-cp38-abi3-macosx_11_0_arm64.whl", hash = "sha256:94fd4858284736bb67a897a41608b5b0c2496c9bdb3bf2af1fa3409127f20d57", size = 447058, upload-time = "2025-11-19T15:18:34.416Z" }, - { url = "https://files.pythonhosted.org/packages/f1/06/578ffed52c2296f93d7fd2d844cabfa92be51a587c38c8afbb8ae449ca89/safetensors-0.7.0-cp38-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e07d91d0c92a31200f25351f4acb2bc6aff7f48094e13ebb1d0fb995b54b6542", size = 491748, upload-time = "2025-11-19T15:18:09.79Z" }, - { url = "https://files.pythonhosted.org/packages/ae/33/1debbbb70e4791dde185edb9413d1fe01619255abb64b300157d7f15dddd/safetensors-0.7.0-cp38-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:8469155f4cb518bafb4acf4865e8bb9d6804110d2d9bdcaa78564b9fd841e104", size = 503881, upload-time = "2025-11-19T15:18:16.145Z" }, - { url = "https://files.pythonhosted.org/packages/8e/1c/40c2ca924d60792c3be509833df711b553c60effbd91da6f5284a83f7122/safetensors-0.7.0-cp38-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:54bef08bf00a2bff599982f6b08e8770e09cc012d7bba00783fc7ea38f1fb37d", size = 623463, upload-time = "2025-11-19T15:18:21.11Z" }, - { url = "https://files.pythonhosted.org/packages/9b/3a/13784a9364bd43b0d61eef4bea2845039bc2030458b16594a1bd787ae26e/safetensors-0.7.0-cp38-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:42cb091236206bb2016d245c377ed383aa7f78691748f3bb6ee1bfa51ae2ce6a", size = 532855, upload-time = "2025-11-19T15:18:25.719Z" }, - { url = "https://files.pythonhosted.org/packages/a0/60/429e9b1cb3fc651937727befe258ea24122d9663e4d5709a48c9cbfceecb/safetensors-0.7.0-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:dac7252938f0696ddea46f5e855dd3138444e82236e3be475f54929f0c510d48", size = 507152, upload-time = "2025-11-19T15:18:33.023Z" }, - { url = "https://files.pythonhosted.org/packages/3c/a8/4b45e4e059270d17af60359713ffd83f97900d45a6afa73aaa0d737d48b6/safetensors-0.7.0-cp38-abi3-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:1d060c70284127fa805085d8f10fbd0962792aed71879d00864acda69dbab981", size = 541856, upload-time = "2025-11-19T15:18:31.075Z" }, - { url = "https://files.pythonhosted.org/packages/06/87/d26d8407c44175d8ae164a95b5a62707fcc445f3c0c56108e37d98070a3d/safetensors-0.7.0-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:cdab83a366799fa730f90a4ebb563e494f28e9e92c4819e556152ad55e43591b", size = 674060, upload-time = "2025-11-19T15:18:37.211Z" }, - { url = "https://files.pythonhosted.org/packages/11/f5/57644a2ff08dc6325816ba7217e5095f17269dada2554b658442c66aed51/safetensors-0.7.0-cp38-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:672132907fcad9f2aedcb705b2d7b3b93354a2aec1b2f706c4db852abe338f85", size = 771715, upload-time = "2025-11-19T15:18:38.689Z" }, - { url = "https://files.pythonhosted.org/packages/86/31/17883e13a814bd278ae6e266b13282a01049b0c81341da7fd0e3e71a80a3/safetensors-0.7.0-cp38-abi3-musllinux_1_2_i686.whl", hash = "sha256:5d72abdb8a4d56d4020713724ba81dac065fedb7f3667151c4a637f1d3fb26c0", size = 714377, upload-time = "2025-11-19T15:18:40.162Z" }, - { url = "https://files.pythonhosted.org/packages/4a/d8/0c8a7dc9b41dcac53c4cbf9df2b9c83e0e0097203de8b37a712b345c0be5/safetensors-0.7.0-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:b0f6d66c1c538d5a94a73aa9ddca8ccc4227e6c9ff555322ea40bdd142391dd4", size = 677368, upload-time = "2025-11-19T15:18:41.627Z" }, - { url = "https://files.pythonhosted.org/packages/05/e5/cb4b713c8a93469e3c5be7c3f8d77d307e65fe89673e731f5c2bfd0a9237/safetensors-0.7.0-cp38-abi3-win32.whl", hash = "sha256:c74af94bf3ac15ac4d0f2a7c7b4663a15f8c2ab15ed0fc7531ca61d0835eccba", size = 326423, upload-time = "2025-11-19T15:18:45.74Z" }, - { url = "https://files.pythonhosted.org/packages/5d/e6/ec8471c8072382cb91233ba7267fd931219753bb43814cbc71757bfd4dab/safetensors-0.7.0-cp38-abi3-win_amd64.whl", hash = "sha256:d1239932053f56f3456f32eb9625590cc7582e905021f94636202a864d470755", size = 341380, upload-time = "2025-11-19T15:18:44.427Z" }, - { url = "https://files.pythonhosted.org/packages/a7/6a/4d08d89a6fcbe905c5ae68b8b34f0791850882fc19782d0d02c65abbdf3b/safetensors-0.7.0-pp310-pypy310_pp73-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f4729811a6640d019a4b7ba8638ee2fd21fa5ca8c7e7bdf0fed62068fcaac737", size = 492430, upload-time = "2025-11-19T15:18:11.884Z" }, - { url = "https://files.pythonhosted.org/packages/dd/29/59ed8152b30f72c42d00d241e58eaca558ae9dbfa5695206e2e0f54c7063/safetensors-0.7.0-pp310-pypy310_pp73-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:12f49080303fa6bb424b362149a12949dfbbf1e06811a88f2307276b0c131afd", size = 503977, upload-time = "2025-11-19T15:18:17.523Z" }, - { url = "https://files.pythonhosted.org/packages/d3/0b/4811bfec67fa260e791369b16dab105e4bae82686120554cc484064e22b4/safetensors-0.7.0-pp310-pypy310_pp73-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:0071bffba4150c2f46cae1432d31995d77acfd9f8db598b5d1a2ce67e8440ad2", size = 623890, upload-time = "2025-11-19T15:18:22.666Z" }, - { url = "https://files.pythonhosted.org/packages/58/5b/632a58724221ef03d78ab65062e82a1010e1bef8e8e0b9d7c6d7b8044841/safetensors-0.7.0-pp310-pypy310_pp73-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:473b32699f4200e69801bf5abf93f1a4ecd432a70984df164fc22ccf39c4a6f3", size = 531885, upload-time = "2025-11-19T15:18:27.146Z" }, +sdist = { url = "https://files.pythonhosted.org/packages/45/06/f955dbbb1859e3bd23c8ac6141af5106e7ad5fedec4a3a6e3d60f94b7001/safetensors-0.8.0.tar.gz", hash = "sha256:fabaf3e0f18a6618d9b36560682562157f77c2b71fcffc7b432be2baed9d753d", size = 325846, upload-time = "2026-06-09T07:52:25.563Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/39/a0/f718cda65b05407d228f97602cf60dca269c979867aa5beb25410de26cd3/safetensors-0.8.0-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:c554f85858e05226d3c2828e32395e677434685d6d94594a41643361c5e837f0", size = 473568, upload-time = "2026-06-09T07:52:18.829Z" }, + { url = "https://files.pythonhosted.org/packages/f5/b1/fa7c600e7dceae12e9606c7578cbc9ff1e1ed55844883ee5c92205e86226/safetensors-0.8.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:c80201d22cbf405b80647a60ada77bba06c8fba2da2743ba1e89cdcc39a81f25", size = 484562, upload-time = "2026-06-09T07:52:17.518Z" }, + { url = "https://files.pythonhosted.org/packages/09/7d/65a7de0af421317bb36a067241e4235fff194eed60b961ed6d3f59a3fc60/safetensors-0.8.0-cp310-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7a46e5ff292c356d6991e60942ba7f79817682d3a2cef0702136448cb9c4d235", size = 502844, upload-time = "2026-06-09T07:52:07.624Z" }, + { url = "https://files.pythonhosted.org/packages/91/4f/3175c9d75634e0e0dda0082794193521035edd7c70a6f212bf33ca06ddf4/safetensors-0.8.0-cp310-abi3-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:4124502b78f03534117c848f87a39b8f31e577b15eff423bf8bfb95f2a8c30d0", size = 511823, upload-time = "2026-06-09T07:52:09.565Z" }, + { url = "https://files.pythonhosted.org/packages/20/87/846c289e7aa2299eff406335717cf43ce8777194ece8aad75772e0411615/safetensors-0.8.0-cp310-abi3-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:7bc0a787ba8a35be368ee3574edfa2b1ad389eebd0a72e482ae275490e3f6c98", size = 633461, upload-time = "2026-06-09T07:52:11.128Z" }, + { url = "https://files.pythonhosted.org/packages/76/22/8d64d9df2c45d5ded401df889d0ad90882804ca172d79ec4f0df8f727fe0/safetensors-0.8.0-cp310-abi3-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:040070828e36dc8e122178bbbd5830ff9e97920affb84cbe0f46442497bed358", size = 545148, upload-time = "2026-06-09T07:52:13.603Z" }, + { url = "https://files.pythonhosted.org/packages/28/50/f203ff3a3ddfe19308efc83c5a3a29ed02bf786732ec35e68bf9162f3365/safetensors-0.8.0-cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fd6f3f93c9a0a7cc2788ee63fb763353d4bd2e89b0751bc78fcf7dda00bea774", size = 516040, upload-time = "2026-06-09T07:52:16.29Z" }, + { url = "https://files.pythonhosted.org/packages/46/fb/cdaed17ceb2948784fd9c36b6fd3e951b608547cea81a48e8ee6f8cfdfcb/safetensors-0.8.0-cp310-abi3-manylinux_2_31_riscv64.whl", hash = "sha256:fcdd41ec4628fee5799f807c73c353629130fbd942aa23d83c623dd6c9d52d78", size = 513832, upload-time = "2026-06-09T07:52:12.37Z" }, + { url = "https://files.pythonhosted.org/packages/0d/49/1e15de264dcc3b77943d2d0c56a95809956883b1c2d6d585c792523f180b/safetensors-0.8.0-cp310-abi3-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:8e9f537aa183a38ace122d27303dcd986b26bd2a7591f9181d7f0c396f4677ca", size = 559930, upload-time = "2026-06-09T07:52:14.743Z" }, + { url = "https://files.pythonhosted.org/packages/2a/43/bf38443278eab4b1be1fce2931e2b012ad9cb7df52ada751d0aab8f7659a/safetensors-0.8.0-cp310-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:87eec7ffed2b809f05a398a8becb7d013f19f7837cd15d9748580d6cf30dbaf4", size = 678670, upload-time = "2026-06-09T07:52:20.032Z" }, + { url = "https://files.pythonhosted.org/packages/72/e3/68cd3fa5b48488e84add63e04cb12f3bc28ae4638c06d4508c6e88823d0e/safetensors-0.8.0-cp310-abi3-musllinux_1_2_armv7l.whl", hash = "sha256:4a95ae2b05d7726d751da4ebf626a2ca782b706e101bd894c95bc2450b1cffcc", size = 786679, upload-time = "2026-06-09T07:52:21.322Z" }, + { url = "https://files.pythonhosted.org/packages/29/4b/1c19c509d56e01f4fbb3d0a2e597450f6cc04d1d56cf52defb0a62dfd715/safetensors-0.8.0-cp310-abi3-musllinux_1_2_i686.whl", hash = "sha256:3ae091f16662658bdc019a4ff6cb4c085bb7d725eb5978b183ffd265863b6d2d", size = 765683, upload-time = "2026-06-09T07:52:22.594Z" }, + { url = "https://files.pythonhosted.org/packages/27/43/41c1621732edd934d868a00d1b891584c892a7b62a9aab82ea5a0a5623ee/safetensors-0.8.0-cp310-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:8e080062fcde23be189565e1c3305d16751a218ecf9412c8601e64204eb6f846", size = 722361, upload-time = "2026-06-09T07:52:23.924Z" }, + { url = "https://files.pythonhosted.org/packages/8e/3f/73ccf82579412b4a71c4ca673f10b5f1f888d7cf5af7fe24f27d30307be4/safetensors-0.8.0-cp310-abi3-win32.whl", hash = "sha256:2ddf52eac562eda224f99acfa7889d02968c1fd59a5b011ae7d8137c37e9c02d", size = 342401, upload-time = "2026-06-09T07:52:28.895Z" }, + { url = "https://files.pythonhosted.org/packages/1b/6d/3fba214c1e5e0f69991677ec3bc17023f0421776975e1de0c682dca475e2/safetensors-0.8.0-cp310-abi3-win_amd64.whl", hash = "sha256:096ec1a98435df7beb08853bb5aa9081a84f23d0adc67ed1a0a10550f608373f", size = 355540, upload-time = "2026-06-09T07:52:27.832Z" }, + { url = "https://files.pythonhosted.org/packages/8d/fc/7eedc3510d97878876e32774eebbeb61c43f148a96e915c84229a3e967aa/safetensors-0.8.0-cp310-abi3-win_arm64.whl", hash = "sha256:f7838e5135a406ad3e02efdcb8cf2e5397d368b0154537c4fec682dbc544d452", size = 340500, upload-time = "2026-06-09T07:52:26.745Z" }, ] [[package]] @@ -2570,7 +2574,7 @@ wheels = [ [[package]] name = "transformers" -version = "5.6.0" +version = "5.14.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "huggingface-hub" }, @@ -2584,9 +2588,9 @@ dependencies = [ { name = "tqdm" }, { name = "typer" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/88/14/e3eb58bd4c9df45af154f9de59a6e66a2ece30dd64434c710e40e5a4b1f1/transformers-5.6.0.tar.gz", hash = "sha256:291951976b79a6f93ec06d6ab14489a99aecdbc8f05aaabab538ea1d508c9a97", size = 8311711, upload-time = "2026-04-22T15:42:03.393Z" } +sdist = { url = "https://files.pythonhosted.org/packages/5a/fb/2a2ba88f325e68a921d8b69ff63b477830b2e73ade9a3c8c8cab2f06d741/transformers-5.14.1.tar.gz", hash = "sha256:60d196c27781eacf8637e2b533f517582907ad6f9ae142046d6b69431a5b2173", size = 9295927, upload-time = "2026-07-16T09:41:57.773Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/e8/80/ac700df26a83969d2cf8d09d3dd191bb9bbe8156fa4607e614438b1da8c8/transformers-5.6.0-py3-none-any.whl", hash = "sha256:def5aac47b28e2f1386fa64f2f458a5474ba7733ab74c1765e7c67a79d1f1cbd", size = 10364790, upload-time = "2026-04-22T15:41:58.723Z" }, + { url = "https://files.pythonhosted.org/packages/6f/67/8d85ca2323233ae3c0365a659c4e52ee1f587b440e4bc577e7d8e4416d0f/transformers-5.14.1-py3-none-any.whl", hash = "sha256:9db974c4079ede2d1a3ea7ca5a240df33f2cc26fc2b36ba64c5f2a4f43b6e725", size = 11625234, upload-time = "2026-07-16T09:41:54.143Z" }, ] [[package]]