diff --git a/mlx_lm/models/laguna.py b/mlx_lm/models/laguna.py new file mode 100644 index 000000000..515ec8230 --- /dev/null +++ b/mlx_lm/models/laguna.py @@ -0,0 +1,499 @@ +from dataclasses import dataclass, field +from typing import Any, Dict, List, 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 + vocab_size: int + hidden_size: int + intermediate_size: int + num_hidden_layers: int + num_attention_heads: int + num_key_value_heads: int + head_dim: int + max_position_embeddings: int + rms_norm_eps: float = 1e-6 + qkv_bias: bool = False + attention_bias: bool = False + gating: Union[bool, str] = True + tie_word_embeddings: bool = False + rope_theta: float = 500000.0 + rope_parameters: Optional[Dict[str, Any]] = None + rope_scaling: Optional[Dict[str, Any]] = None + partial_rotary_factor: Optional[float] = None + rope_style: str = "rotate-half" + sliding_window: Optional[int] = None + layer_types: Optional[List[str]] = None + num_attention_heads_per_layer: Optional[List[int]] = None + swa_rope_parameters: Optional[Dict[str, Any]] = None + swa_attention_sink_enabled: bool = False + num_experts: int = 0 + num_experts_per_tok: int = 0 + moe_intermediate_size: int = 0 + shared_expert_intermediate_size: int = 0 + norm_topk_prob: bool = True + decoder_sparse_step: int = 1 + mlp_only_layers: List[int] = field(default_factory=lambda: [0]) + moe_routed_scaling_factor: float = 1.0 + moe_apply_router_weight_on_input: bool = False + moe_router_logit_softcapping: float = 0.0 + moe_router_use_sigmoid: bool = True + + def __post_init__(self): + if self.gating is True: + self.gating = "per-head" + + if self.layer_types is None: + self.layer_types = ["full_attention"] * self.num_hidden_layers + if len(self.layer_types) != self.num_hidden_layers: + raise ValueError("layer_types must match num_hidden_layers.") + + if self.num_attention_heads_per_layer is None: + self.num_attention_heads_per_layer = [ + self.num_attention_heads + ] * self.num_hidden_layers + if len(self.num_attention_heads_per_layer) != self.num_hidden_layers: + raise ValueError( + "num_attention_heads_per_layer must match num_hidden_layers." + ) + if any( + h % self.num_key_value_heads for h in self.num_attention_heads_per_layer + ): + raise ValueError( + "Every query-head count must be divisible by num_key_value_heads." + ) + + rope_parameters = ( + dict(self.rope_parameters) + if self.rope_parameters is not None + else ( + dict(self.rope_scaling) + if self.rope_scaling is not None + else {"rope_type": "default", "rope_theta": self.rope_theta} + ) + ) + + layer_types = set(self.layer_types) + layer_rope_parameters = { + k: v + for k, v in rope_parameters.items() + if k in layer_types and isinstance(v, dict) + } + if layer_rope_parameters: + top_level_parameters = { + k: v + for k, v in rope_parameters.items() + if k not in layer_types and not isinstance(v, dict) + } + + def rope_parameters_for(layer_type: str) -> Dict[str, Any]: + params = dict(layer_rope_parameters.get(layer_type, {})) + for k, v in top_level_parameters.items(): + params.setdefault(k, v) + return params + + default_layer_type = ( + "full_attention" + if "full_attention" in layer_rope_parameters + else next(iter(layer_rope_parameters)) + ) + self.rope_parameters = rope_parameters_for(default_layer_type) + + if ( + self.swa_rope_parameters is None + and "sliding_attention" in layer_rope_parameters + ): + self.swa_rope_parameters = rope_parameters_for("sliding_attention") + else: + self.rope_parameters = rope_parameters + + if self.swa_rope_parameters is not None: + self.swa_rope_parameters = dict(self.swa_rope_parameters) + + self.rope_parameters.setdefault("rope_type", "default") + if self.swa_rope_parameters is not None: + self.swa_rope_parameters.setdefault("rope_type", "default") + + if self.partial_rotary_factor is not None: + self.rope_parameters.setdefault( + "partial_rotary_factor", self.partial_rotary_factor + ) + if self.swa_rope_parameters is not None: + self.swa_rope_parameters.setdefault( + "partial_rotary_factor", self.partial_rotary_factor + ) + + +def _rope_base(args: ModelArgs, rope_config: Dict[str, Union[float, str]]) -> float: + return float(rope_config.get("rope_theta", args.rope_theta)) + + +def _rope_dims(args: ModelArgs, rope_config: Dict[str, Union[float, str]]) -> int: + partial = float(rope_config.get("partial_rotary_factor", 1.0)) + return int(args.head_dim * partial) + + +class MLP(nn.Module): + def __init__(self, dim: int, hidden_dim: int): + 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 LagunaTopKRouter(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.top_k = args.num_experts_per_tok + self.norm_topk_prob = args.norm_topk_prob + self.use_sigmoid = args.moe_router_use_sigmoid + self.router_logit_softcapping = args.moe_router_logit_softcapping + self.proj = nn.Linear(args.hidden_size, args.num_experts, bias=False) + self.e_score_correction_bias = mx.zeros((args.num_experts,)) + + def __call__(self, x: mx.array) -> tuple[mx.array, mx.array]: + dtype = x.dtype + logits = self.proj(x).astype(mx.float32) + if self.router_logit_softcapping > 0.0: + c = self.router_logit_softcapping + logits = mx.tanh(logits / c) * c + + scores = mx.sigmoid(logits) if self.use_sigmoid else mx.softmax(logits, axis=-1) + corrected_scores = scores + self.e_score_correction_bias.astype(scores.dtype) + + k = self.top_k + inds = mx.stop_gradient( + mx.argpartition(-corrected_scores, kth=k - 1, axis=-1)[..., :k] + ) + weights = mx.take_along_axis(scores, inds, axis=-1) + if self.norm_topk_prob: + weights = weights / mx.sum(weights, axis=-1, keepdims=True) + return inds, weights.astype(dtype) + + +class LagunaSparseMoeBlock(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + if args.moe_apply_router_weight_on_input: + raise NotImplementedError( + "moe_apply_router_weight_on_input=True is not supported." + ) + self.routed_scaling_factor = args.moe_routed_scaling_factor + self.gate = LagunaTopKRouter(args) + self.switch_mlp = SwitchGLU( + args.hidden_size, args.moe_intermediate_size, args.num_experts + ) + self.shared_expert = MLP(args.hidden_size, args.shared_expert_intermediate_size) + + def __call__(self, x: mx.array) -> mx.array: + inds, scores = self.gate(x) + y = self.switch_mlp(x, inds) + y = mx.sum(y * scores[..., None], axis=-2) + if self.routed_scaling_factor != 1.0: + y = y * self.routed_scaling_factor + return y + self.shared_expert(x) + + +class Attention(nn.Module): + def __init__(self, args: ModelArgs, layer_idx: int): + super().__init__() + + self.n_heads = args.num_attention_heads_per_layer[layer_idx] + self.n_kv_heads = args.num_key_value_heads + self.head_dim = args.head_dim + self.scale = self.head_dim**-0.5 + self.gate_per_head = args.gating == "per-head" + self.gating = bool(args.gating) + self.is_sliding = args.layer_types[layer_idx] == "sliding_attention" + self.sliding_window = args.sliding_window if self.is_sliding else None + + dim = args.hidden_size + self.q_proj = nn.Linear(dim, self.n_heads * self.head_dim, bias=args.qkv_bias) + self.k_proj = nn.Linear( + dim, self.n_kv_heads * self.head_dim, bias=args.qkv_bias + ) + self.v_proj = nn.Linear( + dim, self.n_kv_heads * self.head_dim, bias=args.qkv_bias + ) + self.o_proj = nn.Linear( + self.n_heads * self.head_dim, dim, bias=args.attention_bias + ) + + if self.gating: + gate_dim = ( + self.n_heads if self.gate_per_head else self.n_heads * self.head_dim + ) + self.g_proj = nn.Linear(dim, gate_dim, bias=False) + + if self.is_sliding and args.swa_attention_sink_enabled: + self.sink = mx.zeros((self.n_heads,)) + else: + self.sink = None + + self.q_norm = nn.RMSNorm(self.head_dim, eps=args.rms_norm_eps) + self.k_norm = nn.RMSNorm(self.head_dim, eps=args.rms_norm_eps) + + rope_config = ( + args.swa_rope_parameters + if self.is_sliding and args.swa_rope_parameters is not None + else args.rope_parameters + ) + self.rope = initialize_rope( + _rope_dims(args, rope_config), + base=_rope_base(args, rope_config), + traditional=False, + scaling_config=rope_config, + 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, _ = 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, self.head_dim) + ).transpose(0, 2, 1, 3) + keys = self.k_norm( + keys.reshape(B, L, self.n_kv_heads, self.head_dim) + ).transpose(0, 2, 1, 3) + values = values.reshape(B, L, self.n_kv_heads, self.head_dim).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, + sinks=self.sink, + ) + output = output.transpose(0, 2, 1, 3).reshape(B, L, -1) + + if self.gating: + gate = nn.softplus(self.g_proj(x).astype(mx.float32)).astype(output.dtype) + if self.gate_per_head: + shape = output.shape + output = ( + output.reshape(B, L, self.n_heads, self.head_dim) * gate[..., None] + ).reshape(shape) + else: + output = output * gate + + return self.o_proj(output) + + +class DecoderLayer(nn.Module): + def __init__(self, args: ModelArgs, layer_idx: int): + super().__init__() + self.self_attn = Attention(args, layer_idx) + if (layer_idx not in args.mlp_only_layers) and ( + args.num_experts > 0 and (layer_idx + 1) % args.decoder_sparse_step == 0 + ): + self.mlp = LagunaSparseMoeBlock(args) + else: + 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.attention_type = args.layer_types[layer_idx] + + 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)) + return h + r + + +class LagunaModel(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 = [ + DecoderLayer(args, layer_idx) for layer_idx in range(args.num_hidden_layers) + ] + self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps) + self.fa_idx = args.layer_types.index("full_attention") + self.swa_idx = ( + args.layer_types.index("sliding_attention") + if "sliding_attention" in args.layer_types + else None + ) + + def __call__( + self, + inputs: mx.array, + cache=None, + input_embeddings: Optional[mx.array] = None, + ) -> mx.array: + if input_embeddings is not None: + h = input_embeddings + else: + h = self.embed_tokens(inputs) + + if cache is None: + cache = [None] * len(self.layers) + + full_mask = create_attention_mask(h, cache[self.fa_idx]) + if self.swa_idx is not None: + sliding_mask = create_attention_mask( + h, cache[self.swa_idx], window_size=self.args.sliding_window + ) + + for layer, c in zip(self.layers, cache): + mask = ( + sliding_mask + if layer.attention_type == "sliding_attention" + else full_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 = LagunaModel(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, + ) -> mx.array: + out = self.model(inputs, cache, input_embeddings) + if self.args.tie_word_embeddings: + return self.model.embed_tokens.as_linear(out) + return self.lm_head(out) + + def sanitize(self, weights): + if self.args.tie_word_embeddings: + weights.pop("lm_head.weight", None) + + weights = self._unpack_compressed_tensors(weights) + weights = self._remap_router_weights(weights) + weights = self._stack_experts(weights) + return { + k: v + for k, v in weights.items() + if "rotary_emb.inv_freq" not in k + and not k.endswith(".self_attn.k_scale") + and not k.endswith(".self_attn.v_scale") + } + + def _unpack_compressed_tensors(self, weights): + if not any(k.endswith(".weight_shape") for k in weights): + return weights + + new_weights = {} + for k, v in weights.items(): + if k.endswith(".weight_shape"): + base = k[: -len("weight_shape")] + if ( + f"{base}weight_packed" in weights + and f"{base}weight_scale" in weights + ): + scales = weights[f"{base}weight_scale"] + new_weights[f"{base}weight"] = weights[f"{base}weight_packed"].view( + mx.uint32 + ) + new_weights[f"{base}scales"] = scales + new_weights[f"{base}biases"] = (-8 * scales).astype(scales.dtype) + elif k.endswith(".weight_packed") or k.endswith(".weight_scale"): + base = k.rsplit(".", 1)[0] + "." + if f"{base}weight_shape" in weights: + continue + new_weights[k] = v + else: + new_weights[k] = v + return new_weights + + def _remap_router_weights(self, weights): + for layer_idx in range(self.args.num_hidden_layers): + prefix = f"model.layers.{layer_idx}.mlp" + gate_weight = f"{prefix}.gate.weight" + if gate_weight in weights: + weights[f"{prefix}.gate.proj.weight"] = weights.pop(gate_weight) + + legacy_bias = f"{prefix}.experts.e_score_correction_bias" + if legacy_bias in weights: + weights[f"{prefix}.gate.e_score_correction_bias"] = weights.pop( + legacy_bias + ) + return weights + + def _stack_experts(self, weights): + for layer_idx in range(self.args.num_hidden_layers): + prefix = f"model.layers.{layer_idx}.mlp" + for proj in ["gate_proj", "up_proj", "down_proj"]: + for suffix in ["weight", "scales", "biases"]: + first_key = f"{prefix}.experts.0.{proj}.{suffix}" + if first_key not in weights: + continue + weights[f"{prefix}.switch_mlp.{proj}.{suffix}"] = mx.stack( + [ + weights.pop(f"{prefix}.experts.{e}.{proj}.{suffix}") + for e in range(self.args.num_experts) + ] + ) + return weights + + @property + def quant_predicate(self): + def predicate(path, _): + if path.endswith("mlp.gate.proj"): + return {"group_size": 64, "bits": 8} + return True + + return predicate + + @property + def cast_predicate(self): + def predicate(k): + return "e_score_correction_bias" not in k + + return predicate + + @property + def layers(self): + return self.model.layers diff --git a/mlx_lm/tokenizer_utils.py b/mlx_lm/tokenizer_utils.py index c7e50fbe7..f475e9f9b 100644 --- a/mlx_lm/tokenizer_utils.py +++ b/mlx_lm/tokenizer_utils.py @@ -557,6 +557,12 @@ def _infer_tool_parser(chat_template): return "function_gemma" elif "" in chat_template: return "longcat" + elif ( + "function-name" in chat_template + and "" in chat_template + and "" in chat_template + ): + return "laguna" elif "" in chat_template: return "glm47" elif "<|tool_list_start|>" in chat_template: diff --git a/mlx_lm/tool_parsers/laguna.py b/mlx_lm/tool_parsers/laguna.py new file mode 100644 index 000000000..444b37f41 --- /dev/null +++ b/mlx_lm/tool_parsers/laguna.py @@ -0,0 +1,83 @@ +# Copyright © 2026 Apple Inc. + +""" +Tool parser for Poolside Laguna XML-like tool calls. + +Format: +function-name +argument-key +value-of-argument-key + +""" + +import ast +import json +from typing import Any + +import regex as re + +tool_call_start = "" +tool_call_end = "" + +_tool_call_regex = re.compile(r"(.*?)", re.DOTALL) +_func_name_regex = re.compile(r"^(.*?)", re.DOTALL) +_arg_pair_regex = re.compile( + r"(.*?)(?:\\n|\s)*(.*?)", + re.DOTALL, +) + + +def _is_string_type( + tool_name: str, + arg_name: str, + tools: list[Any] | None, +) -> bool: + if tools is None: + return False + for tool in tools: + func = tool.get("function", {}) + if func.get("name") != tool_name: + continue + params = func.get("parameters") or {} + arg_type = params.get("properties", {}).get(arg_name, {}).get("type") + return arg_type == "string" + return False + + +def _deserialize(value: str) -> Any: + try: + return json.loads(value) + except Exception: + pass + try: + return ast.literal_eval(value) + except Exception: + pass + return value + + +def _parse_single_call(text: str, tools: list[Any] | None): + text = text.strip() + match = _func_name_regex.search(text) + if not match: + func_name = text.split("\n", 1)[0].strip() + return dict(name=func_name, arguments={}) + + func_name = match.group(1).strip() + arguments = {} + for match in _arg_pair_regex.finditer(text): + arg_key = match.group(1).strip() + arg_val = match.group(2).strip() + if not _is_string_type(func_name, arg_key, tools): + arg_val = _deserialize(arg_val) + arguments[arg_key] = arg_val + return dict(name=func_name, arguments=arguments) + + +def parse_tool_call(text: str, tools: list[Any] | None = None): + matches = _tool_call_regex.findall(text) + if matches: + calls = [_parse_single_call(match, tools) for match in matches] + return calls[0] if len(calls) == 1 else calls + + return _parse_single_call(text, tools) diff --git a/mlx_lm/utils.py b/mlx_lm/utils.py index ef3d266b9..93b5758ef 100644 --- a/mlx_lm/utils.py +++ b/mlx_lm/utils.py @@ -172,6 +172,31 @@ def _transform_awq_weights( return new_weights, mlx_quantization +def _compressed_tensors_config(quantization_config: Dict[str, Any]) -> Dict[str, Any]: + config_groups = quantization_config.get("config_groups") + if not config_groups: + # Default quantization settings + return {"group_size": 32, "bits": 4, "mode": "affine"} + + # Use first config group for parameters + first_group = next(iter(config_groups.values())) + weights_config = first_group.get("weights", {}) + + format_type = quantization_config.get("format") + weights_type = weights_config.get("type") + + if format_type == "pack-quantized" and weights_type == "int": + return { + "group_size": weights_config.get("group_size", 32), + "bits": weights_config.get("num_bits", 4), + "mode": "affine", + } + + raise ValueError( + f"Unsupported compressed-tensors quantization format: {format_type}" + ) + + def _get_classes(config: dict): """ Retrieve the model and model args classes based on the configuration. @@ -378,7 +403,7 @@ def class_predicate(p, m): config["quantization_config"] = quantization _quantize(quantization) elif quant_method == "compressed-tensors": - quantization = {"group_size": 32, "bits": 4, "mode": "affine"} + quantization = _compressed_tensors_config(quantization_config) config["quantization"] = quantization config["quantization_config"] = quantization _quantize(quantization) diff --git a/tests/test_models.py b/tests/test_models.py index 6e1fcd96e..a7cc96a4d 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -842,6 +842,68 @@ def test_gemma4_moe_router_quantizes_to_8bit(self): ) self.assertEqual(config["quantization"]["bits"], 4) + def test_laguna_sanitize_hf_moe_weights(self): + from mlx_lm.models import laguna + + args = laguna.ModelArgs.from_dict( + { + "model_type": "laguna", + "vocab_size": 32, + "hidden_size": 8, + "intermediate_size": 16, + "num_hidden_layers": 2, + "num_attention_heads": 2, + "num_key_value_heads": 1, + "head_dim": 4, + "max_position_embeddings": 32, + "num_experts": 2, + "num_experts_per_tok": 1, + "decoder_sparse_step": 1, + "mlp_only_layers": [0], + "moe_intermediate_size": 3, + "shared_expert_intermediate_size": 3, + "layer_types": ["full_attention", "sliding_attention"], + "sliding_window": 4, + "swa_rope_parameters": { + "rope_theta": 10000.0, + "rope_type": "linear", + "factor": 1.0, + "partial_rotary_factor": 1.0, + }, + } + ) + model = laguna.Model(args) + + weights = { + "model.layers.1.mlp.gate.weight": mx.ones((2, 8)), + "model.layers.1.mlp.experts.e_score_correction_bias": mx.arange( + 2, dtype=mx.float32 + ), + } + for e in range(2): + weights[f"model.layers.1.mlp.experts.{e}.gate_proj.weight"] = mx.full( + (3, 8), e + 1, dtype=mx.float32 + ) + weights[f"model.layers.1.mlp.experts.{e}.up_proj.weight"] = mx.full( + (3, 8), e + 3, dtype=mx.float32 + ) + weights[f"model.layers.1.mlp.experts.{e}.down_proj.weight"] = mx.full( + (8, 3), e + 5, dtype=mx.float32 + ) + + converted = model.sanitize(weights) + + self.assertIn("model.layers.1.mlp.gate.proj.weight", converted) + self.assertIn("model.layers.1.mlp.gate.e_score_correction_bias", converted) + self.assertIn("model.layers.1.mlp.switch_mlp.gate_proj.weight", converted) + self.assertIn("model.layers.1.mlp.switch_mlp.up_proj.weight", converted) + self.assertIn("model.layers.1.mlp.switch_mlp.down_proj.weight", converted) + self.assertEqual( + converted["model.layers.1.mlp.switch_mlp.gate_proj.weight"].shape, + (2, 3, 8), + ) + self.assertFalse(any(".experts." in k for k in converted)) + def test_qwen2_moe(self): from mlx_lm.models import qwen2_moe @@ -1797,6 +1859,57 @@ def test_hunyuan(self): model, args.model_type, args.vocab_size, args.num_hidden_layers ) + def test_laguna(self): + from mlx_lm.models import laguna + + args = laguna.ModelArgs( + model_type="laguna", + vocab_size=1000, + hidden_size=128, + intermediate_size=256, + num_hidden_layers=4, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=16, + max_position_embeddings=1000, + rms_norm_eps=1e-5, + attention_bias=False, + mlp_only_layers=[0], + num_experts=4, + num_experts_per_tok=2, + decoder_sparse_step=1, + moe_intermediate_size=64, + shared_expert_intermediate_size=64, + norm_topk_prob=True, + moe_routed_scaling_factor=2.5, + gating=True, + sliding_window=4, + layer_types=[ + "full_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + ], + num_attention_heads_per_layer=[4, 8, 8, 4], + rope_parameters={ + "full_attention": { + "rope_theta": 10000.0, + "rope_type": "default", + "partial_rotary_factor": 0.5, + }, + "sliding_attention": { + "rope_theta": 10000.0, + "rope_type": "linear", + "factor": 1.0, + "partial_rotary_factor": 1.0, + }, + }, + ) + model = laguna.Model(args) + self.model_test_runner( + model, args.model_type, args.vocab_size, args.num_hidden_layers + ) + def test_hunyuan_v1_dense(self): from mlx_lm.models import hunyuan_v1_dense diff --git a/tests/test_tool_parsing.py b/tests/test_tool_parsing.py index 52892b7ff..7eff6fae4 100644 --- a/tests/test_tool_parsing.py +++ b/tests/test_tool_parsing.py @@ -7,6 +7,7 @@ glm47, json_tools, kimi_k2, + laguna, longcat, minimax_m2, mistral, @@ -29,6 +30,14 @@ def test_parsers(self): "multiplya12234585b48838483920", glm47, ), + ( + "multiply\n" + "a\n" + "12234585\n" + "b\n" + "48838483920", + laguna, + ), ( '{"name": "multiply", "arguments": {"a": 12234585, "b": 48838483920}}', json_tools, @@ -99,6 +108,12 @@ def test_parsers(self): 'get_current_temperaturelocation"London"', glm47, ), + ( + "get_current_temperature\n" + "location\n" + "London", + laguna, + ), ( '{"name": "get_current_temperature", "arguments": {"location": "London"}}', json_tools, @@ -329,6 +344,73 @@ def test_minimax_m2(self): tool_calls = minimax_m2.parse_tool_call(test_case, None) self.assertEqual(expected, tool_calls) + def test_laguna_strips_function_name_and_preserves_string_args(self): + tools = [ + { + "type": "function", + "function": { + "name": "get_weather", + "parameters": { + "type": "object", + "properties": { + "location": {"type": "string"}, + "days": {"type": "integer"}, + "postal_code": {"type": "string"}, + }, + }, + }, + } + ] + test_case = ( + "get_weather\n" + "location\n" + "Warsaw\n" + "days\n" + "3\n" + "postal_code\n" + "00123" + ) + + tool_call = laguna.parse_tool_call(test_case, tools) + + self.assertEqual(tool_call["name"], "get_weather") + self.assertEqual( + tool_call["arguments"], + {"location": "Warsaw", "days": 3, "postal_code": "00123"}, + ) + + def test_laguna_full_tool_call_tags(self): + test_case = ( + "search\n" + "query\n" + "weather\n" + "" + ) + tool_call = laguna.parse_tool_call(test_case, None) + self.assertEqual( + tool_call, + {"name": "search", "arguments": {"query": "weather"}}, + ) + + test_case = ( + "search\n" + "query\n" + "weather\n" + "" + "read_file\n" + "path\n" + "/tmp/test.txt\n" + "" + ) + tool_calls = laguna.parse_tool_call(test_case, None) + self.assertEqual( + tool_calls, + [ + {"name": "search", "arguments": {"query": "weather"}}, + {"name": "read_file", "arguments": {"path": "/tmp/test.txt"}}, + ], + ) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_utils.py b/tests/test_utils.py index 88b68fe33..e7f19abe3 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -79,6 +79,23 @@ def test_quantize(self): self.assertEqual(config["quantization"]["group_size"], 64) self.assertEqual(config["quantization"]["bits"], 4) + def test_compressed_tensors_config_uses_weight_group_size(self): + quantization = utils._compressed_tensors_config( + { + "format": "pack-quantized", + "config_groups": { + "group_0": { + "weights": { + "type": "int", + "num_bits": 4, + "group_size": 128, + } + } + }, + } + ) + self.assertEqual(quantization, {"group_size": 128, "bits": 4, "mode": "affine"}) + def test_convert(self): mlx_path = os.path.join(self.test_dir, "mlx_model")