import gzip import json from pathlib import Path import pytest from app.catalogue import build_record from app.parsing import ( detect_arch, detect_family, detect_params, detect_quant, param_bucket, quant_bucket, ) QWEN3_8B_CONFIG = { "architectures": ["Qwen3ForCausalLM"], "model_type": "qwen3", "head_dim": 128, "hidden_size": 4096, "max_position_embeddings": 40960, "num_attention_heads": 32, "num_hidden_layers": 36, "num_key_value_heads": 8, "quantization": {"group_size": 64, "bits": 4}, "sliding_window": None, "use_sliding_window": False, } @pytest.mark.parametrize("mid,tags,family", [ ("mlx-community/Qwen3.8-27B-4bit", ["base_model:Qwen/Qwen3.8-27B"], "Qwen"), ("mlx-community/Josiefied-Qwen3-8B-abliterated-v1-4bit", [], "Qwen"), ("mlx-community/Meta-Llama-3.1-8B-Instruct-4bit", [], "Llama"), ("mlx-community/gemma-4-31B-it-qat-4bit", [], "Gemma"), ("mlx-community/Devstral-Small-2507-4bit", [], "Mistral"), ("mlx-community/DeepSeek-V4-Flash-2bit", [], "DeepSeek"), ("mlx-community/Kimi-K2-Instruct-4bit", [], "Kimi"), ("mlx-community/NVIDIA-Nemotron-3-Nano-4bit", [], "Nemotron"), ("mlx-community/Zorblax-7B-4bit", [], "Zorblax"), # never seen before -> still a family ("mlx-community/7B", [], "Other"), ("mlx-community/functiongemma-270m-it-4bit", [], "Gemma"), ("mlx-community/OLMoE-1B-7B-0924-Instruct-4bit", [], "OLMo"), ("mlx-community/Instruct-Zorblax-7B", [], "Zorblax"), ("mlx-community/Starling-LM-7B-beta", [], "Starling"), ]) def test_family(mid, tags, family): assert detect_family(mid, tags) == family def test_family_ignores_junk_tags(): assert detect_family("mlx-community/Qwen3-8B-4bit", [None, 5, "base_model:"]) == "Qwen" def test_params_prefers_safetensors_then_name(): p = detect_params("mlx-community/Qwen3-8B-4bit", {"total": 8_190_735_360}) assert p.source == "safetensors" and param_bucket(p.total) == "8-15B" p = detect_params("mlx-community/Qwen3-8B-4bit", None) assert p.source == "name" and p.total == 8e9 and param_bucket(p.total) == "8-15B" def test_params_moe(): p = detect_params("mlx-community/Qwen3-30B-A3B-4bit", {"total": 30.5e9}) assert p.moe and p.active == 3e9 and param_bucket(p.total) == "15-35B" p = detect_params("mlx-community/Mixtral-8x7B-Instruct-v0.1", None) assert p.moe and p.total == 56e9 p = detect_params("mlx-community/whatever-4bit", None, {"num_local_experts": 8}) assert p.moe @pytest.mark.parametrize("st", [None, {}, {"total": "big"}, {"total": -1}, "junk"]) def test_params_unknown(st): p = detect_params("mlx-community/mystery-model-4bit", st) assert p.total is None and p.source == "unknown" and param_bucket(p.total) is None def test_quant_tags(): q = detect_quant("mlx-community/Qwen3-8B", ["4-bit"]) assert (q.bits, q.source, q.mixed, quant_bucket(q)) == (4, "tags", False, "4-bit") def test_quant_mixed_from_tags_and_config(): q = detect_quant("mlx-community/Qwen3.5-9B-OptiQ-4bit", ["4bit", "8bit", "4-bit"]) assert q.bits == 4 and q.mixed and q.label == "4-bit mixed" and quant_bucket(q) == "4-bit" cfg = {"model_type": "x", "quantization": {"bits": 4, "group_size": 64, "model.layers.0.mlp": {"bits": 8, "group_size": 64}}} q = detect_quant("mlx-community/foo", [], cfg) assert q.bits == 4 and q.mixed and q.source == "config" def test_quant_config_beats_name(): q = detect_quant("mlx-community/foo-8bit", ["8-bit"], QWEN3_8B_CONFIG) assert q.bits == 4 and q.source == "config" @pytest.mark.parametrize("name,bits,mode", [ ("gpt-oss-20b-MXFP4-Q8", 4, "mxfp4"), ("Qwen3-30B-A3B-nvfp4", 4, "nvfp4"), ("Qwen3-14B-4bit-DWQ", 4, "dwq"), ("Llama-3.2-3B-Instruct-bf16", 16, "float"), ("Mistral-7B-Instruct-v0.2-8-bit-mlx", 8, "affine"), ("llama2-13b-qnt4bit", 4, "affine"), ("Yi-9B-q", None, None), ]) def test_quant_names(name, bits, mode): q = detect_quant(f"mlx-community/{name}", []) assert q.bits == bits and q.mode == mode def test_quant_unquantized_config(): q = detect_quant("mlx-community/phi-2", [], {"model_type": "phi", "torch_dtype": "float16"}) assert q.bits == 16 and q.label == "16-bit" @pytest.mark.parametrize("cfg", [None, "junk", {}, {"quantization": "4"}, {"quantization": {"bits": "x"}}]) def test_quant_junk_config(cfg): q = detect_quant("mlx-community/mystery", [], cfg) assert q.bits is None and q.label == "unknown" def test_arch_standard(): a = detect_arch(QWEN3_8B_CONFIG) assert a.known and (a.layers, a.full_attention_layers, a.kv_heads, a.head_dim) == (36, 36, 8, 128) assert a.sliding_layers == 0 and a.max_context == 40960 def test_arch_hybrid_linear_attention(): cfg = {"text_config": {"num_hidden_layers": 8, "num_attention_heads": 16, "num_key_value_heads": 4, "head_dim": 256, "layer_types": ["linear_attention"] * 3 + ["full_attention"] + ["linear_attention"] * 3 + ["full_attention"]}} a = detect_arch(cfg) assert a.full_attention_layers == 2 and a.sliding_layers == 0 def test_arch_sliding(): cfg = {"num_hidden_layers": 4, "num_attention_heads": 8, "num_key_value_heads": 8, "head_dim": 64, "sliding_window": 128, "layer_types": ["sliding_attention", "full_attention"] * 2} a = detect_arch(cfg) assert (a.full_attention_layers, a.sliding_layers, a.sliding_window) == (2, 2, 128) @pytest.mark.parametrize("cfg", [None, {}, {"num_hidden_layers": "x"}, {"num_hidden_layers": True}]) def test_arch_unknown(cfg): assert not detect_arch(cfg).known @pytest.mark.parametrize("raw", [{}, {"id": 5}, {"id": "noslash"}, {"id": "a/b", "tags": "notalist"}, {"id": "a/b", "downloads": "many", "safetensors": [1]}]) def test_build_record_never_raises(raw): rec = build_record(raw) assert rec is None or rec.id == "a/b" def test_snapshot_parses_fully(): path = Path(__file__).resolve().parent.parent / "data" / "catalogue_snapshot.json.gz" raw = json.load(gzip.open(path, "rt")) recs = [r for r in map(build_record, raw) if r] assert len(recs) == len(raw) > 1000 llm = [r for r in recs if r.is_llm] known_quant = sum(r.quant_bucket != "unknown" for r in llm) / len(llm) known_size = sum(r.size_bucket is not None for r in llm) / len(llm) assert known_quant > 0.9 and known_size > 0.95 def test_partial_repo_detected(): p = detect_params("mlx-community/Qwen3.8-27B-MTP-4bit", {"total": 424_699_392}) assert p.partial and p.total == 424_699_392 and p.name_total == 27e9 assert not detect_params("mlx-community/Qwen3-8B-4bit", {"total": 8.19e9}).partial assert not detect_params("mlx-community/Qwen3-30B-A3B-4bit", {"total": 30.5e9}).partial def test_packed_param_counts_are_unpacked(): st = {"parameters": {"F16": 244_584_448, "U32": 977_272_832}, "total": 1_221_857_280} p = detect_params("mlx-community/EXAONE-3.5-7.8B-Instruct-4bit", st, bits=4) assert abs(p.total / 1e9 - 7.82) < 0.05 and not p.partial and param_bucket(p.total) == "3-8B" st6 = {"parameters": {"F16": 244_584_448, "U32": 1_465_909_248}, "total": 1_710_493_696} assert abs(detect_params("mlx-community/x-6bit", st6, bits=6).total / 1e9 - 7.82) < 0.05 big = {"parameters": {"F16": 3_468_959_744, "U32": 8_879_865_856}, "total": 12_348_825_600} assert 71e9 < detect_params("mlx-community/Liberated-Qwen1.5-72B-4bit", big, bits=4).total < 74e9 new_style = {"parameters": {"U32": 8_190_427_136, "BF16": 308_224}, "total": 8_190_735_360} assert detect_params("mlx-community/Qwen3-8B-4bit", new_style, bits=4).total == 8_190_735_360 # packed counts but no bit width: keep the Hub total, the partial check flags the mismatch p = detect_params("mlx-community/EXAONE-3.5-7.8B-Instruct", st, bits=None) assert p.total == st["total"] and p.partial # vision-language models have lots of unquantized float weights but are unpacked vl = {"parameters": {"U32": 7_615_283_200, "F16": 676_883_456}, "total": 8_292_166_656} assert detect_params("mlx-community/Qwen2.5-VL-7B-Instruct-4bit", vl, bits=4).total == vl["total"] vl2 = {"parameters": {"U32": 4_022_272_000, "BF16": 415_543_808}, "total": 4_437_815_808} assert detect_params("mlx-community/some-vlm-4bit", vl2, bits=4).total == vl2["total"] assert detect_params("mlx-community/some-vlm-8bit", vl2, bits=8).total == vl2["total"] assert abs(detect_params("mlx-community/unnamed-4bit", st, bits=4).total / 1e9 - 7.82) < 0.05 def test_unknown_quantization_format_does_not_mean_unquantized(): cfg = {"model_type": "gemma4", "quantization_config": {"quant_method": "gemma", "module_quant_configs": {"x": {"num_bits": 2}}}} q = detect_quant("mlx-community/gemma-4-E4B-it-qat-mobile", ["8-bit"], cfg) assert q.bits == 8 and q.source == "tags" def test_mixed_quant_uses_named_width(): q = detect_quant("mlx-community/Qwen3.5-122B-A10B-OptiQ-2bit", ["2bit", "mixed-precision", "4-bit"]) assert q.bits == 2 and q.mixed and quant_bucket(q) == "2-bit" cfg = {"model_type": "x", "quantization": {"bits": 4, "group_size": 64, "model.layers.0.mlp": {"bits": 2}}} q = detect_quant("mlx-community/Qwen3.5-122B-A10B-OptiQ-2bit", [], cfg) assert q.bits == 2 and q.mixed # uniform repos keep the config width even if a name says otherwise q = detect_quant("mlx-community/foo-8bit", [], {"model_type": "x", "quantization": {"bits": 4}}) assert q.bits == 4 @pytest.mark.parametrize("a,b,same", [ ("mlx-community/Qwen3.5-9B-OptiQ-4bit", "mlx-community/Qwen3.5-9B-MLX-4bit", True), ("mlx-community/Qwen3.5-35B-A3B-OptiQ-4bit", "mlx-community/Qwen3.5-35B-A3B-GPTQ-Int4", True), ("mlx-community/gemma-4-26B-A4B-it-OptiQ-4bit", "mlx-community/gemma-4-26b-a4b-it-nvfp4", True), ("mlx-community/Nex-N2.5-mini-OptiQ-4bit", "mlx-community/Nex-N2.5-mini-oQ4", True), ("mlx-community/gpt-oss-20b-OptiQ-4bit", "mlx-community/gpt-oss-20b-MXFP4-Q8", True), ("mlx-community/MiniCPM5-2B-OptiQ-4bit", "mlx-community/MiniCPM5-2B-mlx-4Bit", True), ("mlx-community/gemma-4-12B-it-OptiQ-4bit", "mlx-community/gemma-4-12B-coder-fable5-composer2.5-v1-4bit-msq", False), ("mlx-community/gemma-4-12B-it-qat-OptiQ-4bit", "mlx-community/Huihui-gemma-4-12B-it-qat-q4_0-unquantized-abliterated-5bit-msq", False), ]) def test_weights_key_groups_quants_not_finetunes(a, b, same): from app.parsing import weights_key assert (weights_key(a) == weights_key(b)) is same, (weights_key(a), weights_key(b))