tests/test_parsing.py

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"