| """Native MLX implementation of LiquidAI/LFM2.5-Encoder-350M. |
| |
| This is an encoder/masked-language model: attention and short convolutions are |
| bidirectional. It intentionally does not use mlx-lm's causal generation API. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| from dataclasses import dataclass |
| from pathlib import Path |
| from typing import Optional |
|
|
| import mlx.core as mx |
| import mlx.nn as nn |
|
|
|
|
| @dataclass |
| class ModelArgs: |
| vocab_size: int |
| hidden_size: int |
| num_hidden_layers: int |
| num_attention_heads: int |
| num_key_value_heads: int |
| norm_eps: float |
| conv_bias: bool |
| conv_L_cache: int |
| block_dim: int |
| intermediate_size: int |
| block_multiple_of: int |
| block_ffn_dim_multiplier: float |
| block_auto_adjust_ff_dim: bool |
| layer_types: list[str] |
| rope_theta: float = 1_000_000.0 |
|
|
| @classmethod |
| def from_dict(cls, config: dict) -> "ModelArgs": |
| rope = config.get("rope_parameters") or {} |
| fields = { |
| "vocab_size", |
| "hidden_size", |
| "num_hidden_layers", |
| "num_attention_heads", |
| "num_key_value_heads", |
| "norm_eps", |
| "conv_bias", |
| "conv_L_cache", |
| "block_dim", |
| "intermediate_size", |
| "block_multiple_of", |
| "block_ffn_dim_multiplier", |
| "block_auto_adjust_ff_dim", |
| "layer_types", |
| } |
| values = {key: config[key] for key in fields} |
| values["rope_theta"] = rope.get("rope_theta", config.get("rope_theta", 1_000_000.0)) |
| return cls(**values) |
|
|
|
|
| class Attention(nn.Module): |
| def __init__(self, args: ModelArgs): |
| super().__init__() |
| dim = args.hidden_size |
| self.n_heads = args.num_attention_heads |
| self.n_kv_heads = args.num_key_value_heads |
| self.head_dim = dim // self.n_heads |
| self.scale = self.head_dim**-0.5 |
|
|
| self.q_layernorm = nn.RMSNorm(self.head_dim, eps=args.norm_eps) |
| self.k_layernorm = nn.RMSNorm(self.head_dim, eps=args.norm_eps) |
| self.q_proj = nn.Linear(dim, self.n_heads * self.head_dim, bias=False) |
| self.k_proj = nn.Linear(dim, self.n_kv_heads * self.head_dim, bias=False) |
| self.v_proj = nn.Linear(dim, self.n_kv_heads * self.head_dim, bias=False) |
| self.out_proj = nn.Linear(self.n_heads * self.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) -> mx.array: |
| batch, length, _ = x.shape |
| queries = self.q_proj(x).reshape(batch, length, self.n_heads, -1) |
| keys = self.k_proj(x).reshape(batch, length, self.n_kv_heads, -1) |
| values = self.v_proj(x).reshape(batch, length, self.n_kv_heads, -1) |
|
|
| queries = self.q_layernorm(queries).transpose(0, 2, 1, 3) |
| keys = self.k_layernorm(keys).transpose(0, 2, 1, 3) |
| values = values.transpose(0, 2, 1, 3) |
| queries = self.rope(queries) |
| keys = self.rope(keys) |
|
|
| output = mx.fast.scaled_dot_product_attention( |
| queries, keys, values, scale=self.scale, mask=mask |
| ) |
| output = output.transpose(0, 2, 1, 3).reshape(batch, length, -1) |
| return self.out_proj(output) |
|
|
|
|
| class ShortConv(nn.Module): |
| def __init__(self, args: ModelArgs): |
| super().__init__() |
| self.kernel_size = args.conv_L_cache |
| self.conv = nn.Conv1d( |
| in_channels=args.hidden_size, |
| out_channels=args.hidden_size, |
| kernel_size=self.kernel_size, |
| groups=args.hidden_size, |
| bias=args.conv_bias, |
| ) |
| self.in_proj = nn.Linear(args.hidden_size, 3 * args.hidden_size, bias=args.conv_bias) |
| self.out_proj = nn.Linear(args.hidden_size, args.hidden_size, bias=args.conv_bias) |
|
|
| def __call__(self, x: mx.array, padding_mask: Optional[mx.array] = None) -> mx.array: |
| if padding_mask is not None: |
| x = mx.where(padding_mask[..., None], x, 0) |
| gate_b, gate_c, value = mx.split(self.in_proj(x), 3, axis=-1) |
| gated = gate_b * value |
| pad = self.kernel_size // 2 |
| gated = mx.pad(gated, [(0, 0), (pad, pad), (0, 0)]) |
| conv_out = self.conv(gated) |
| if conv_out.shape[1] > x.shape[1]: |
| conv_out = conv_out[:, : x.shape[1], :] |
| elif conv_out.shape[1] < x.shape[1]: |
| conv_out = mx.pad(conv_out, [(0, 0), (0, x.shape[1] - conv_out.shape[1]), (0, 0)]) |
| return self.out_proj(gate_c * conv_out) |
|
|
|
|
| class MLP(nn.Module): |
| def __init__(self, args: ModelArgs): |
| super().__init__() |
| ff_dim = args.intermediate_size |
| if args.block_auto_adjust_ff_dim: |
| ff_dim = int(2 * ff_dim / 3) |
| ff_dim = int(args.block_ffn_dim_multiplier * ff_dim) |
| multiple = args.block_multiple_of |
| ff_dim = multiple * ((ff_dim + multiple - 1) // multiple) |
| self.w1 = nn.Linear(args.block_dim, ff_dim, bias=False) |
| self.w3 = nn.Linear(args.block_dim, ff_dim, bias=False) |
| self.w2 = nn.Linear(ff_dim, args.block_dim, bias=False) |
|
|
| def __call__(self, x: mx.array) -> mx.array: |
| return self.w2(nn.silu(self.w1(x)) * self.w3(x)) |
|
|
|
|
| class EncoderLayer(nn.Module): |
| def __init__(self, args: ModelArgs, layer_type: str): |
| super().__init__() |
| self.is_attention_layer = layer_type == "full_attention" |
| if self.is_attention_layer: |
| self.self_attn = Attention(args) |
| else: |
| self.conv = ShortConv(args) |
| self.feed_forward = MLP(args) |
| 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, |
| attention_mask: Optional[mx.array], |
| padding_mask: Optional[mx.array], |
| ) -> mx.array: |
| normalized = self.operator_norm(x) |
| if self.is_attention_layer: |
| residual = self.self_attn(normalized, mask=attention_mask) |
| else: |
| residual = self.conv(normalized, padding_mask=padding_mask) |
| hidden = x + residual |
| return hidden + self.feed_forward(self.ffn_norm(hidden)) |
|
|
|
|
| class Lfm2BidirectionalModel(nn.Module): |
| def __init__(self, args: ModelArgs): |
| super().__init__() |
| self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size) |
| self.layers = [EncoderLayer(args, layer_type) for layer_type in args.layer_types] |
| self.embedding_norm = nn.RMSNorm(args.hidden_size, eps=args.norm_eps) |
|
|
| def __call__( |
| self, input_ids: mx.array, attention_mask: Optional[mx.array] = None |
| ) -> mx.array: |
| hidden = self.embed_tokens(input_ids) |
| padding_mask = None |
| additive_mask = None |
| if attention_mask is not None: |
| padding_mask = attention_mask.astype(mx.bool_) |
| additive_mask = mx.where( |
| padding_mask[:, None, None, :], |
| mx.array(0.0, dtype=hidden.dtype), |
| mx.array(-1e9, dtype=hidden.dtype), |
| ) |
| for layer in self.layers: |
| hidden = layer(hidden, additive_mask, padding_mask) |
| return self.embedding_norm(hidden) |
|
|
|
|
| class Model(nn.Module): |
| """Bidirectional encoder with a tied masked-language-model head.""" |
|
|
| def __init__(self, args: ModelArgs): |
| super().__init__() |
| self.args = args |
| self.lfm2 = Lfm2BidirectionalModel(args) |
|
|
| def __call__( |
| self, |
| input_ids: mx.array, |
| attention_mask: Optional[mx.array] = None, |
| *, |
| return_hidden: bool = False, |
| ) -> mx.array | tuple[mx.array, mx.array]: |
| hidden = self.lfm2(input_ids, attention_mask=attention_mask) |
| logits = self.lfm2.embed_tokens.as_linear(hidden) |
| return (logits, hidden) if return_hidden else logits |
|
|
|
|
| def load_model(path: str | Path) -> Model: |
| path = Path(path) |
| with (path / "config.json").open() as handle: |
| args = ModelArgs.from_dict(json.load(handle)) |
| model = Model(args) |
| weights = mx.load(str(path / "model.safetensors")) |
| model.load_weights(list(weights.items()), strict=True) |
| mx.eval(model.parameters()) |
| return model |
|
|
