wehub-resource-sync
1644 lines
66 KiB
Python
1644 lines
66 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Tests for model discovery functionality."""
|
|
|
|
import json
|
|
import logging
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from omlx.model_discovery import (
|
|
DiscoveredModel,
|
|
_is_adapter_dir,
|
|
_is_helper_checkpoint,
|
|
_is_unsupported_model,
|
|
_read_model_context_length,
|
|
_register_model,
|
|
_resolve_hf_cache_entry,
|
|
detect_model_type,
|
|
discover_models,
|
|
discover_models_from_dirs,
|
|
estimate_model_size,
|
|
format_size,
|
|
is_helper_config_model_type,
|
|
is_helper_model_config,
|
|
model_directory_access_error,
|
|
)
|
|
|
|
|
|
class TestIsHelperConfigModelType:
|
|
"""Tests for helper (dFlash / Assistant / Draft) model_type detection."""
|
|
|
|
@pytest.mark.parametrize(
|
|
"config_model_type",
|
|
["gemma4_assistant", "qwen3_5_mtp", "foo_mtp", "bar_assistant", "QWEN3_5_MTP"],
|
|
)
|
|
def test_helper_types(self, config_model_type):
|
|
assert is_helper_config_model_type(config_model_type) is True
|
|
|
|
@pytest.mark.parametrize(
|
|
"config_model_type",
|
|
["qwen3", "gemma4_text", "gemma4", "llama", "qwen2_5_vl", "", None, 123],
|
|
)
|
|
def test_non_helper_types(self, config_model_type):
|
|
assert is_helper_config_model_type(config_model_type) is False
|
|
|
|
|
|
class TestIsHelperModelConfig:
|
|
"""Tests for full-config drafter detection (model_type / architecture / config-block)."""
|
|
|
|
def test_dflash_draft_via_architecture(self):
|
|
# DFlash drafts declare a plain qwen3 model_type but a DFlashDraftModel arch.
|
|
config = {"model_type": "qwen3", "architectures": ["DFlashDraftModel"]}
|
|
assert is_helper_model_config(config) is True
|
|
|
|
def test_dflash_draft_via_config_block(self):
|
|
config = {"model_type": "qwen3", "dflash_config": {"block_size": 16}}
|
|
assert is_helper_model_config(config) is True
|
|
|
|
def test_assistant_via_model_type(self):
|
|
config = {
|
|
"model_type": "gemma4_assistant",
|
|
"architectures": ["Gemma4Assistant"],
|
|
}
|
|
assert is_helper_model_config(config) is True
|
|
|
|
def test_mtp_via_model_type(self):
|
|
config = {"model_type": "qwen3_5_mtp"}
|
|
assert is_helper_model_config(config) is True
|
|
|
|
@pytest.mark.parametrize(
|
|
"config",
|
|
[
|
|
{"model_type": "qwen3", "architectures": ["Qwen3ForCausalLM"]},
|
|
{
|
|
"model_type": "gemma4",
|
|
"architectures": ["Gemma4ForConditionalGeneration"],
|
|
},
|
|
{"model_type": "llama"},
|
|
{},
|
|
{"architectures": None},
|
|
{"model_type": 123, "architectures": 123},
|
|
],
|
|
)
|
|
def test_non_helper_configs(self, config):
|
|
assert is_helper_model_config(config) is False
|
|
|
|
|
|
class TestDetectModelType:
|
|
"""Tests for detect_model_type function."""
|
|
|
|
def test_detect_llm_model(self, tmp_path):
|
|
"""Test detection of LLM model."""
|
|
config = {
|
|
"model_type": "llama",
|
|
"architectures": ["LlamaForCausalLM"],
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(tmp_path) == "llm"
|
|
|
|
def test_detect_qwen_model(self, tmp_path):
|
|
"""Test detection of Qwen model as LLM."""
|
|
config = {
|
|
"model_type": "qwen2",
|
|
"architectures": ["Qwen2ForCausalLM"],
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(tmp_path) == "llm"
|
|
|
|
def test_detect_embedding_model_by_type(self, tmp_path):
|
|
"""Test detection of embedding model by model_type."""
|
|
config = {
|
|
"model_type": "bert",
|
|
"architectures": ["BertModel"],
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(tmp_path) == "embedding"
|
|
|
|
def test_detect_embedding_model_by_architecture(self, tmp_path):
|
|
"""Test detection of embedding model by architecture."""
|
|
config = {
|
|
"model_type": "unknown",
|
|
"architectures": ["XLMRobertaModel"],
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(tmp_path) == "embedding"
|
|
|
|
def test_detect_modernbert_embedding(self, tmp_path):
|
|
"""Test detection of ModernBERT as embedding model."""
|
|
config = {
|
|
"model_type": "modernbert",
|
|
"architectures": ["ModernBertModel"],
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(tmp_path) == "embedding"
|
|
|
|
def test_detect_reranker_model(self, tmp_path):
|
|
"""Test detection of reranker model by architecture."""
|
|
config = {
|
|
"model_type": "modernbert",
|
|
"architectures": ["ModernBertForSequenceClassification"],
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(tmp_path) == "reranker"
|
|
|
|
def test_detect_xlm_roberta_reranker(self, tmp_path):
|
|
"""Test detection of XLM-RoBERTa reranker."""
|
|
config = {
|
|
"model_type": "xlm-roberta",
|
|
"architectures": ["XLMRobertaForSequenceClassification"],
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(tmp_path) == "reranker"
|
|
|
|
def test_detect_jina_reranker_without_name_heuristic(self, tmp_path):
|
|
"""JinaForRanking should detect as reranker without requiring 'rerank' in directory name."""
|
|
model_dir = tmp_path / "jina-v3-mlx"
|
|
model_dir.mkdir()
|
|
config = {
|
|
"model_type": "qwen3",
|
|
"architectures": ["JinaForRanking"],
|
|
}
|
|
(model_dir / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(model_dir) == "reranker"
|
|
|
|
def test_detect_causal_lm_reranker(self, tmp_path):
|
|
"""Test detection of CausalLM-based reranker (e.g., Qwen3-Reranker)."""
|
|
reranker_dir = tmp_path / "Qwen3-Reranker-0.6B-mxfp8"
|
|
reranker_dir.mkdir()
|
|
config = {
|
|
"model_type": "qwen3",
|
|
"architectures": ["Qwen3ForCausalLM"],
|
|
}
|
|
(reranker_dir / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(reranker_dir) == "reranker"
|
|
|
|
def test_detect_causal_lm_embedding(self, tmp_path):
|
|
"""Test detection of CausalLM-based embedding (e.g., Qwen3-Embedding)."""
|
|
embed_dir = tmp_path / "Qwen3-Embedding-8B-mxfp8"
|
|
embed_dir.mkdir()
|
|
config = {
|
|
"model_type": "qwen3",
|
|
"architectures": ["Qwen3ForCausalLM"],
|
|
}
|
|
(embed_dir / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(embed_dir) == "embedding"
|
|
|
|
def test_causal_lm_without_reranker_or_embedding_name_is_llm(self, tmp_path):
|
|
"""Test that Qwen3ForCausalLM without 'reranker' or 'embedding' in name is LLM."""
|
|
llm_dir = tmp_path / "Qwen3-0.6B"
|
|
llm_dir.mkdir()
|
|
config = {
|
|
"model_type": "qwen3",
|
|
"architectures": ["Qwen3ForCausalLM"],
|
|
}
|
|
(llm_dir / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(llm_dir) == "llm"
|
|
|
|
def test_detect_qwen2_causal_lm_embedding(self, tmp_path):
|
|
"""Qwen2ForCausalLM with an embedding-named dir is an embedding (#686).
|
|
|
|
jina-code-embeddings ships a Qwen2ForCausalLM architecture without
|
|
lm_head weights, so it must classify as an embedding model rather than
|
|
a chat LLM.
|
|
"""
|
|
embed_dir = tmp_path / "jina-code-embeddings-1.5b"
|
|
embed_dir.mkdir()
|
|
config = {
|
|
"model_type": "qwen2",
|
|
"architectures": ["Qwen2ForCausalLM"],
|
|
}
|
|
(embed_dir / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(embed_dir) == "embedding"
|
|
|
|
def test_qwen2_causal_lm_without_embedding_name_is_llm(self, tmp_path):
|
|
"""Qwen2ForCausalLM without an embedding-named dir stays an LLM (#686).
|
|
|
|
Regression guard: adding Qwen2ForCausalLM to the embedding-arch set must
|
|
not reclassify ordinary Qwen2/Qwen2.5 chat models, which is gated by the
|
|
_is_causal_lm_embedding directory-name heuristic.
|
|
"""
|
|
llm_dir = tmp_path / "Qwen2.5-7B-Instruct"
|
|
llm_dir.mkdir()
|
|
config = {
|
|
"model_type": "qwen2",
|
|
"architectures": ["Qwen2ForCausalLM"],
|
|
}
|
|
(llm_dir / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(llm_dir) == "llm"
|
|
|
|
def test_detect_lfm2_text_model_is_llm(self, tmp_path):
|
|
"""LFM2 text checkpoints share model_type with non-text variants."""
|
|
llm_dir = tmp_path / "LFM2-1.2B"
|
|
llm_dir.mkdir()
|
|
config = {
|
|
"model_type": "lfm2",
|
|
"architectures": ["Lfm2ForCausalLM"],
|
|
}
|
|
(llm_dir / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(llm_dir) == "llm"
|
|
|
|
def test_detect_lfm2_5_moe_text_model_is_llm(self, tmp_path):
|
|
"""LiquidAI/LFM2.5-8B-A1B is a text-generation MoE LLM."""
|
|
llm_dir = tmp_path / "LFM2.5-8B-A1B"
|
|
llm_dir.mkdir()
|
|
config = {
|
|
"model_type": "lfm2_moe",
|
|
"architectures": ["Lfm2MoeForCausalLM"],
|
|
}
|
|
(llm_dir / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(llm_dir) == "llm"
|
|
|
|
def test_detect_qwen3_vl_reranker(self, tmp_path):
|
|
"""Qwen3VLForConditionalGeneration + 'reranker' in dir name → reranker."""
|
|
reranker_dir = tmp_path / "Qwen3-VL-Reranker-2B-4bit"
|
|
reranker_dir.mkdir()
|
|
config = {
|
|
"model_type": "qwen3_vl",
|
|
"architectures": ["Qwen3VLForConditionalGeneration"],
|
|
"vision_config": {"hidden_size": 1024},
|
|
}
|
|
(reranker_dir / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(reranker_dir) == "reranker"
|
|
|
|
def test_detect_qwen3_vl_embedding(self, tmp_path):
|
|
"""Qwen3VLForConditionalGeneration + 'embedding' in dir name → embedding."""
|
|
embed_dir = tmp_path / "Qwen3-VL-Embedding-2B"
|
|
embed_dir.mkdir()
|
|
config = {
|
|
"model_type": "qwen3_vl",
|
|
"architectures": ["Qwen3VLForConditionalGeneration"],
|
|
"vision_config": {"hidden_size": 1024},
|
|
}
|
|
(embed_dir / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(embed_dir) == "embedding"
|
|
|
|
def test_detect_qwen3_vl_without_rerank_or_embed_name_is_vlm(self, tmp_path):
|
|
"""Plain Qwen3-VL without rerank/embed hints still classifies as VLM."""
|
|
vlm_dir = tmp_path / "Qwen3-VL-2B-Instruct"
|
|
vlm_dir.mkdir()
|
|
config = {
|
|
"model_type": "qwen3_vl",
|
|
"architectures": ["Qwen3VLForConditionalGeneration"],
|
|
"vision_config": {"hidden_size": 1024},
|
|
}
|
|
(vlm_dir / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(vlm_dir) == "vlm"
|
|
|
|
def test_missing_config_defaults_to_llm(self, tmp_path):
|
|
"""Test that missing config.json defaults to LLM."""
|
|
assert detect_model_type(tmp_path) == "llm"
|
|
|
|
def test_invalid_json_defaults_to_llm(self, tmp_path):
|
|
"""Test that invalid JSON defaults to LLM."""
|
|
(tmp_path / "config.json").write_text("not valid json")
|
|
assert detect_model_type(tmp_path) == "llm"
|
|
|
|
def test_empty_config_defaults_to_llm(self, tmp_path):
|
|
"""Test that empty config defaults to LLM."""
|
|
(tmp_path / "config.json").write_text("{}")
|
|
assert detect_model_type(tmp_path) == "llm"
|
|
|
|
def test_detect_vlm_by_model_type(self, tmp_path):
|
|
"""Test detection of VLM model by model_type."""
|
|
config = {
|
|
"model_type": "qwen2_5_vl",
|
|
"architectures": ["Qwen2_5_VLForConditionalGeneration"],
|
|
"vision_config": {"hidden_size": 1152},
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(tmp_path) == "vlm"
|
|
|
|
def test_detect_diffusion_gemma_as_vlm(self, tmp_path):
|
|
"""DiffusionGemma is served by mlx-vlm even without vision_config."""
|
|
config = {
|
|
"model_type": "diffusion_gemma",
|
|
"canvas_length": 256,
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(tmp_path) == "vlm"
|
|
|
|
def test_detect_cohere2_moe_as_vlm_without_vision_config(self, tmp_path):
|
|
"""Cohere2 MoE is text-only but implemented by mlx-vlm."""
|
|
config = {
|
|
"model_type": "cohere2_moe",
|
|
"architectures": ["Cohere2MoeForCausalLM"],
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(tmp_path) == "vlm"
|
|
|
|
def test_detect_minimax_m3_vl_as_vlm(self, tmp_path):
|
|
"""MiniMax M3 VL is served by mlx-vlm."""
|
|
config = {
|
|
"model_type": "minimax_m3_vl",
|
|
"architectures": ["MiniMaxM3VLForConditionalGeneration"],
|
|
"vision_config": {"hidden_size": 1280},
|
|
"text_config": {"hidden_size": 6144},
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(tmp_path) == "vlm"
|
|
|
|
def test_detect_minimax_m3_text_as_vlm_native_text(self, tmp_path):
|
|
"""MiniMax M3 text-only checkpoints are implemented by mlx-vlm."""
|
|
config = {
|
|
"model_type": "minimax_m3",
|
|
"architectures": ["MiniMaxM3ForCausalLM"],
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(tmp_path) == "vlm"
|
|
|
|
def test_detect_vlm_by_architecture(self, tmp_path):
|
|
"""Test detection of VLM model by architecture name."""
|
|
config = {
|
|
"model_type": "unknown_vlm",
|
|
"architectures": ["LlavaForConditionalGeneration"],
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(tmp_path) == "vlm"
|
|
|
|
def test_detect_vlm_by_vision_config(self, tmp_path):
|
|
"""Test detection of VLM model by vision_config + text_config."""
|
|
config = {
|
|
"model_type": "some_new_vlm",
|
|
"vision_config": {"hidden_size": 1024},
|
|
"text_config": {"hidden_size": 2048},
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(tmp_path) == "vlm"
|
|
|
|
def test_detect_vlm_gemma3(self, tmp_path):
|
|
"""Test detection of Gemma3 as VLM."""
|
|
config = {
|
|
"model_type": "gemma3",
|
|
"architectures": ["Gemma3ForConditionalGeneration"],
|
|
"vision_config": {"hidden_size": 1152},
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(tmp_path) == "vlm"
|
|
|
|
def test_detect_vlm_gemma4(self, tmp_path):
|
|
"""Test detection of Gemma4 as VLM."""
|
|
config = {
|
|
"model_type": "gemma4",
|
|
"architectures": ["Gemma4ForConditionalGeneration"],
|
|
"vision_config": {"hidden_size": 1152},
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(tmp_path) == "vlm"
|
|
|
|
def test_detect_vlm_gemma4_unified(self, tmp_path):
|
|
"""Test detection of Gemma4 unified as VLM."""
|
|
config = {
|
|
"model_type": "gemma4_unified",
|
|
"architectures": ["Gemma4UnifiedForConditionalGeneration"],
|
|
"vision_config": {"hidden_size": 1152},
|
|
"audio_config": {"feature_size": 128},
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(tmp_path) == "vlm"
|
|
|
|
def test_detect_text_only_gemma3_as_llm(self, tmp_path):
|
|
"""Text-only quant of Gemma3 (no vision_config) should be LLM."""
|
|
config = {
|
|
"model_type": "gemma3",
|
|
"architectures": ["Gemma3ForConditionalGeneration"],
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(tmp_path) == "llm"
|
|
|
|
def test_detect_text_only_gemma4_as_llm(self, tmp_path):
|
|
"""Text-only quant of Gemma4 (no vision_config) should be LLM."""
|
|
config = {
|
|
"model_type": "gemma4",
|
|
"architectures": ["Gemma4ForConditionalGeneration"],
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(tmp_path) == "llm"
|
|
|
|
def test_detect_gemma4_unified_without_vision_config_as_vlm(self, tmp_path):
|
|
"""Gemma4 unified is always VLM even without vision_config."""
|
|
config = {
|
|
"model_type": "gemma4_unified",
|
|
"architectures": ["Gemma4UnifiedForConditionalGeneration"],
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(tmp_path) == "vlm"
|
|
|
|
def test_detect_vlm_qwen3_5_moe(self, tmp_path):
|
|
"""Test detection of Qwen3.5 MoE as VLM."""
|
|
config = {
|
|
"model_type": "qwen3_5_moe",
|
|
"vision_config": {"depth": 32, "hidden_size": 1280},
|
|
"text_config": {"hidden_size": 4096},
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(tmp_path) == "vlm"
|
|
|
|
def test_detect_text_only_qwen3_5_moe_as_llm(self, tmp_path):
|
|
"""Text-only quant of qwen3_5_moe (no vision_config) should be LLM."""
|
|
config = {
|
|
"model_type": "qwen3_5_moe",
|
|
"architectures": ["Qwen3_5MoeForConditionalGeneration"],
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(tmp_path) == "llm"
|
|
|
|
def test_detect_qwen3_causal_lm_is_llm(self, tmp_path):
|
|
"""Qwen3 with CausalLM architecture should be LLM, not embedding."""
|
|
config = {
|
|
"model_type": "qwen3",
|
|
"architectures": ["Qwen3ForCausalLM"],
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(tmp_path) == "llm"
|
|
|
|
def test_detect_qwen3_embedding_by_architecture(self, tmp_path):
|
|
"""Qwen3 with TextEmbedding architecture should be embedding."""
|
|
config = {
|
|
"model_type": "qwen3",
|
|
"architectures": ["Qwen3ForTextEmbedding"],
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(tmp_path) == "embedding"
|
|
|
|
def test_detect_qwen3_no_architecture_defaults_to_llm(self, tmp_path):
|
|
"""Qwen3 without architectures field should default to LLM."""
|
|
config = {
|
|
"model_type": "qwen3",
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(tmp_path) == "llm"
|
|
|
|
def test_detect_gemma3_text_without_embedding_arch_is_llm(self, tmp_path):
|
|
"""gemma3_text with non-embedding architecture should be LLM."""
|
|
config = {
|
|
"model_type": "gemma3_text",
|
|
"architectures": ["Gemma3TextForCausalLM"],
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(tmp_path) == "llm"
|
|
|
|
def test_detect_gemma3_text_model_without_sentence_transformers_modules_is_llm(self, tmp_path):
|
|
"""gemma3_text base transformer without sentence-transformers modules stays LLM."""
|
|
config = {
|
|
"model_type": "gemma3_text",
|
|
"architectures": ["Gemma3TextModel"],
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(tmp_path) == "llm"
|
|
|
|
def test_detect_sentence_transformers_gemma3_text_as_embedding(self, tmp_path):
|
|
"""gemma3_text sentence-transformers exports should be detected as embeddings."""
|
|
config = {
|
|
"model_type": "gemma3_text",
|
|
"architectures": ["Gemma3TextModel"],
|
|
}
|
|
modules = [
|
|
{
|
|
"idx": 0,
|
|
"name": "0",
|
|
"path": "",
|
|
"type": "sentence_transformers.models.Transformer",
|
|
},
|
|
{
|
|
"idx": 1,
|
|
"name": "1",
|
|
"path": "1_Pooling",
|
|
"type": "sentence_transformers.models.Pooling",
|
|
},
|
|
{
|
|
"idx": 2,
|
|
"name": "2",
|
|
"path": "2_Normalize",
|
|
"type": "sentence_transformers.models.Normalize",
|
|
},
|
|
]
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
(tmp_path / "modules.json").write_text(json.dumps(modules))
|
|
assert detect_model_type(tmp_path) == "embedding"
|
|
|
|
def test_transformer_only_modules_json_is_not_embedding(self, tmp_path):
|
|
"""modules.json with only Transformer (no Pooling/Normalize) should not be embedding."""
|
|
config = {
|
|
"model_type": "gemma3_text",
|
|
"architectures": ["Gemma3TextModel"],
|
|
}
|
|
modules = [
|
|
{
|
|
"idx": 0,
|
|
"name": "0",
|
|
"path": "",
|
|
"type": "sentence_transformers.models.Transformer",
|
|
},
|
|
]
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
(tmp_path / "modules.json").write_text(json.dumps(modules))
|
|
assert detect_model_type(tmp_path) == "llm"
|
|
|
|
def test_detect_vlm_model_type_requires_vision_config(self, tmp_path):
|
|
"""VLM_MODEL_TYPES match without vision_config should fall back to LLM."""
|
|
config = {
|
|
"model_type": "gemma3",
|
|
# No VLM architecture, no vision_config — text-only derivative
|
|
"architectures": ["SomeTextOnlyArch"],
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(tmp_path) == "llm"
|
|
|
|
def test_detect_vlm_molmo2_via_vit_config(self, tmp_path):
|
|
"""Molmo2 bf16 ships ``vit_config`` (not ``vision_config``) — must classify as VLM."""
|
|
config = {
|
|
"model_type": "molmo2",
|
|
"architectures": ["Molmo2ForConditionalGeneration"],
|
|
"vit_config": {"hidden_size": 1024, "num_layers": 24},
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(tmp_path) == "vlm"
|
|
|
|
def test_detect_vlm_fastvlm_via_mm_vision_tower(self, tmp_path):
|
|
"""FastVLM bf16 ships ``mm_vision_tower`` (older LLaVA style) — must classify as VLM."""
|
|
config = {
|
|
"model_type": "llava_qwen2",
|
|
"architectures": ["LlavaQwen2ForCausalLM"],
|
|
"mm_vision_tower": "openai/clip-vit-large-patch14-336",
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(tmp_path) == "vlm"
|
|
|
|
def test_detect_text_only_quant_with_empty_mm_vision_tower_as_llm(self, tmp_path):
|
|
"""``mm_vision_tower`` evidence check is non-empty-only — empty string falls back to LLM."""
|
|
config = {
|
|
"model_type": "llava_qwen2",
|
|
"architectures": ["LlavaQwen2ForCausalLM"],
|
|
"mm_vision_tower": "",
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(tmp_path) == "llm"
|
|
|
|
def test_detect_text_only_quant_no_vision_evidence_as_llm(self, tmp_path):
|
|
"""A VLM_ARCHITECTURE match without any vision sub-config evidence is a text-only quant."""
|
|
config = {
|
|
# qwen3_5_moe is in VLM_MODEL_TYPES, and the unsloth Qwen3.5-122B
|
|
# text-only quant ships only the chat backbone arch.
|
|
"model_type": "qwen3_5_moe",
|
|
"architectures": ["Qwen3_5_MoEForCausalLM"],
|
|
# No vision_config / vit_config / mm_vision_tower.
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(tmp_path) == "llm"
|
|
|
|
def test_detect_lfm_text_moe_family_as_llm_not_audio_sts(self, tmp_path):
|
|
"""LFM text MoE (lfm*_moe + *ForCausalLM) must not use mlx-audio STS."""
|
|
config = {
|
|
"model_type": "lfm2_moe",
|
|
"architectures": ["Lfm2MoeForCausalLM"],
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(tmp_path) == "llm"
|
|
|
|
def test_detect_lfm_future_moe_variant_as_llm_not_audio_sts(self, tmp_path):
|
|
"""Any lfm*_moe text type with *ForCausalLM uses LLM, not mlx-audio STS."""
|
|
config = {
|
|
"model_type": "lfm2_5_moe",
|
|
"architectures": ["Lfm2MoeForCausalLM"],
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(tmp_path) == "llm"
|
|
|
|
def test_detect_lfm_audio_architecture_as_sts(self, tmp_path):
|
|
"""mlx-audio LFM STS remains classified via LFM2AudioModel architecture."""
|
|
config = {
|
|
"model_type": "lfm_audio",
|
|
"architectures": ["LFM2AudioModel"],
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(tmp_path) == "audio_sts"
|
|
|
|
def test_detect_lfm_audio_model_type_as_sts(self, tmp_path):
|
|
"""lfm_audio model_type is STS even without a recognized architecture."""
|
|
config = {
|
|
"model_type": "lfm_audio",
|
|
"architectures": [],
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(tmp_path) == "audio_sts"
|
|
|
|
def test_detect_unknown_lfm_prefix_without_causal_lm_as_sts(self, tmp_path):
|
|
"""Unknown mlx-audio lfm* types without CausalLM still use prefix fallback."""
|
|
config = {
|
|
"model_type": "lfm_custom_audio",
|
|
"architectures": ["SomeLegacyLFMAudioWrapper"],
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert detect_model_type(tmp_path) == "audio_sts"
|
|
|
|
|
|
class TestEstimateModelSize:
|
|
"""Tests for estimate_model_size function."""
|
|
|
|
def test_estimate_from_safetensors(self, tmp_path):
|
|
"""Test size estimation from safetensors files."""
|
|
# Create dummy safetensors files
|
|
(tmp_path / "model-00001-of-00002.safetensors").write_bytes(b"0" * 1000)
|
|
(tmp_path / "model-00002-of-00002.safetensors").write_bytes(b"0" * 2000)
|
|
|
|
size = estimate_model_size(tmp_path)
|
|
# 3000 bytes + 5% overhead
|
|
assert size == int(3000 * 1.05)
|
|
|
|
def test_estimate_from_single_safetensors(self, tmp_path):
|
|
"""Test size estimation from single safetensors file."""
|
|
(tmp_path / "model.safetensors").write_bytes(b"0" * 5000)
|
|
|
|
size = estimate_model_size(tmp_path)
|
|
assert size == int(5000 * 1.05)
|
|
|
|
def test_estimate_from_bin_files(self, tmp_path):
|
|
"""Test size estimation from .bin files when no safetensors."""
|
|
# Create dummy bin files
|
|
(tmp_path / "pytorch_model.bin").write_bytes(b"0" * 5000)
|
|
|
|
size = estimate_model_size(tmp_path)
|
|
# 5000 bytes + 5% overhead
|
|
assert size == int(5000 * 1.05)
|
|
|
|
def test_skip_optimizer_files(self, tmp_path):
|
|
"""Test that optimizer files are skipped."""
|
|
(tmp_path / "model.bin").write_bytes(b"0" * 1000)
|
|
(tmp_path / "optimizer.bin").write_bytes(b"0" * 2000)
|
|
|
|
size = estimate_model_size(tmp_path)
|
|
# Only model.bin (1000 bytes) + 5% overhead
|
|
assert size == int(1000 * 1.05)
|
|
|
|
def test_skip_training_files(self, tmp_path):
|
|
"""Test that training files are skipped."""
|
|
(tmp_path / "model.bin").write_bytes(b"0" * 1000)
|
|
(tmp_path / "training_args.bin").write_bytes(b"0" * 500)
|
|
|
|
size = estimate_model_size(tmp_path)
|
|
assert size == int(1000 * 1.05)
|
|
|
|
def test_no_weights_raises_error(self, tmp_path):
|
|
"""Test that missing weights raises ValueError."""
|
|
with pytest.raises(ValueError, match="No model weights found"):
|
|
estimate_model_size(tmp_path)
|
|
|
|
def test_nested_safetensors(self, tmp_path):
|
|
"""Test size estimation from nested safetensors."""
|
|
# Create subdirectory with safetensors
|
|
subdir = tmp_path / "weights"
|
|
subdir.mkdir()
|
|
(subdir / "model.safetensors").write_bytes(b"0" * 3000)
|
|
|
|
size = estimate_model_size(tmp_path)
|
|
assert size == int(3000 * 1.05)
|
|
|
|
def test_prefers_safetensors_over_bin(self, tmp_path):
|
|
"""Test that safetensors are preferred over bin files."""
|
|
(tmp_path / "model.safetensors").write_bytes(b"0" * 1000)
|
|
(tmp_path / "model.bin").write_bytes(b"0" * 5000)
|
|
|
|
size = estimate_model_size(tmp_path)
|
|
# Should use safetensors size only
|
|
assert size == int(1000 * 1.05)
|
|
|
|
|
|
class TestDiscoverModels:
|
|
"""Tests for discover_models function."""
|
|
|
|
def test_discover_single_model(self, tmp_path):
|
|
"""Test discovery of a single model."""
|
|
model_dir = tmp_path / "llama-3b"
|
|
model_dir.mkdir()
|
|
(model_dir / "config.json").write_text(json.dumps({"model_type": "llama"}))
|
|
(model_dir / "model.safetensors").write_bytes(b"0" * 1000)
|
|
|
|
models = discover_models(tmp_path)
|
|
assert len(models) == 1
|
|
assert "llama-3b" in models
|
|
assert models["llama-3b"].model_type == "llm"
|
|
assert models["llama-3b"].engine_type == "batched"
|
|
|
|
def test_register_model_skips_duplicate_id(self, tmp_path, caplog):
|
|
"""Collision guard: a second model with an already-registered model_id is
|
|
skipped (first registration kept), not silently overwritten, and the
|
|
collision is logged naming both the kept and the skipped path."""
|
|
# A real, registerable model dir — without the guard this WOULD overwrite.
|
|
second = tmp_path / "dup"
|
|
second.mkdir()
|
|
(second / "config.json").write_text(json.dumps({"model_type": "llama"}))
|
|
(second / "model.safetensors").write_bytes(b"0" * 1000)
|
|
|
|
original = DiscoveredModel(
|
|
model_id="dup",
|
|
model_path="/first/dup",
|
|
model_type="llm",
|
|
engine_type="batched",
|
|
estimated_size=123,
|
|
)
|
|
models = {"dup": original}
|
|
with caplog.at_level(logging.WARNING):
|
|
_register_model(models, second, "dup")
|
|
|
|
# Guard kept the first registration rather than overwriting it.
|
|
assert models["dup"] is original
|
|
assert models["dup"].model_path == "/first/dup"
|
|
# ...and surfaced the collision, naming both the kept and the skipped path.
|
|
assert "Duplicate model_id 'dup'" in caplog.text
|
|
assert "/first/dup" in caplog.text
|
|
assert str(second) in caplog.text
|
|
|
|
def test_discover_model_dir_is_itself_a_model(self, tmp_path):
|
|
"""Test that pointing directly at a model directory works."""
|
|
model_dir = tmp_path / "Qwen3-5-9B-MLX-4bit"
|
|
model_dir.mkdir()
|
|
(model_dir / "config.json").write_text(json.dumps({"model_type": "qwen2"}))
|
|
(model_dir / "model.safetensors").write_bytes(b"0" * 1000)
|
|
|
|
models = discover_models(model_dir)
|
|
assert len(models) == 1
|
|
assert "Qwen3-5-9B-MLX-4bit" in models
|
|
assert models["Qwen3-5-9B-MLX-4bit"].model_type == "llm"
|
|
assert models["Qwen3-5-9B-MLX-4bit"].engine_type == "batched"
|
|
|
|
def test_discover_no_fallback_when_subdirs_have_models(self, tmp_path):
|
|
"""Fallback should not trigger when subdirectory models exist,
|
|
even if model_dir itself has config.json."""
|
|
# model_dir itself looks like a model
|
|
(tmp_path / "config.json").write_text(json.dumps({"model_type": "llama"}))
|
|
(tmp_path / "model.safetensors").write_bytes(b"0" * 1000)
|
|
|
|
# but it also has a valid model subdirectory
|
|
sub = tmp_path / "qwen-7b"
|
|
sub.mkdir()
|
|
(sub / "config.json").write_text(json.dumps({"model_type": "qwen2"}))
|
|
(sub / "model.safetensors").write_bytes(b"0" * 2000)
|
|
|
|
models = discover_models(tmp_path)
|
|
assert len(models) == 1
|
|
assert "qwen-7b" in models
|
|
assert tmp_path.name not in models
|
|
|
|
def test_discover_multiple_models(self, tmp_path):
|
|
"""Test discovery of multiple models."""
|
|
# Create first LLM model
|
|
llm_dir = tmp_path / "llama-3b"
|
|
llm_dir.mkdir()
|
|
(llm_dir / "config.json").write_text(json.dumps({"model_type": "llama"}))
|
|
(llm_dir / "model.safetensors").write_bytes(b"0" * 1000)
|
|
|
|
# Create second LLM model
|
|
llm2_dir = tmp_path / "qwen-7b"
|
|
llm2_dir.mkdir()
|
|
(llm2_dir / "config.json").write_text(json.dumps({"model_type": "qwen2"}))
|
|
(llm2_dir / "model.safetensors").write_bytes(b"0" * 2000)
|
|
|
|
models = discover_models(tmp_path)
|
|
assert len(models) == 2
|
|
assert models["llama-3b"].engine_type == "batched"
|
|
assert models["qwen-7b"].engine_type == "batched"
|
|
|
|
def test_discover_embedding_model(self, tmp_path):
|
|
"""Test discovery of embedding model with correct engine type."""
|
|
emb_dir = tmp_path / "bge-small"
|
|
emb_dir.mkdir()
|
|
(emb_dir / "config.json").write_text(
|
|
json.dumps({"model_type": "bert", "architectures": ["BertModel"]})
|
|
)
|
|
(emb_dir / "model.safetensors").write_bytes(b"0" * 500)
|
|
|
|
models = discover_models(tmp_path)
|
|
assert len(models) == 1
|
|
assert models["bge-small"].model_type == "embedding"
|
|
assert models["bge-small"].engine_type == "embedding"
|
|
|
|
def test_discover_reranker_model(self, tmp_path):
|
|
"""Test discovery of reranker model with correct engine type."""
|
|
reranker_dir = tmp_path / "bge-reranker"
|
|
reranker_dir.mkdir()
|
|
(reranker_dir / "config.json").write_text(
|
|
json.dumps({
|
|
"model_type": "modernbert",
|
|
"architectures": ["ModernBertForSequenceClassification"]
|
|
})
|
|
)
|
|
(reranker_dir / "model.safetensors").write_bytes(b"0" * 500)
|
|
|
|
models = discover_models(tmp_path)
|
|
assert len(models) == 1
|
|
assert models["bge-reranker"].model_type == "reranker"
|
|
assert models["bge-reranker"].engine_type == "reranker"
|
|
|
|
def test_skip_invalid_directories(self, tmp_path):
|
|
"""Test that directories without config.json are skipped."""
|
|
# Valid model
|
|
valid_dir = tmp_path / "valid-model"
|
|
valid_dir.mkdir()
|
|
(valid_dir / "config.json").write_text(json.dumps({"model_type": "llama"}))
|
|
(valid_dir / "model.safetensors").write_bytes(b"0" * 1000)
|
|
|
|
# Invalid model (no config.json)
|
|
invalid_dir = tmp_path / "invalid-model"
|
|
invalid_dir.mkdir()
|
|
(invalid_dir / "model.safetensors").write_bytes(b"0" * 1000)
|
|
|
|
models = discover_models(tmp_path)
|
|
assert len(models) == 1
|
|
assert "valid-model" in models
|
|
assert "invalid-model" not in models
|
|
|
|
def test_skip_hidden_directories(self, tmp_path):
|
|
"""Test that hidden directories are skipped."""
|
|
# Hidden directory
|
|
hidden_dir = tmp_path / ".hidden"
|
|
hidden_dir.mkdir()
|
|
(hidden_dir / "config.json").write_text(json.dumps({"model_type": "llama"}))
|
|
(hidden_dir / "model.safetensors").write_bytes(b"0" * 1000)
|
|
|
|
models = discover_models(tmp_path)
|
|
assert len(models) == 0
|
|
|
|
def test_skip_files(self, tmp_path):
|
|
"""Test that files are skipped (only directories processed)."""
|
|
# Create a file at top level
|
|
(tmp_path / "README.md").write_text("readme")
|
|
|
|
# Create valid model
|
|
model_dir = tmp_path / "llama-3b"
|
|
model_dir.mkdir()
|
|
(model_dir / "config.json").write_text(json.dumps({"model_type": "llama"}))
|
|
(model_dir / "model.safetensors").write_bytes(b"0" * 1000)
|
|
|
|
models = discover_models(tmp_path)
|
|
assert len(models) == 1
|
|
|
|
def test_nonexistent_directory_raises_error(self, tmp_path):
|
|
"""Test that nonexistent directory raises ValueError."""
|
|
with pytest.raises(ValueError, match="does not exist"):
|
|
discover_models(tmp_path / "nonexistent")
|
|
|
|
def test_file_instead_of_directory_raises_error(self, tmp_path):
|
|
"""Test that file path raises ValueError."""
|
|
file_path = tmp_path / "file.txt"
|
|
file_path.write_text("content")
|
|
with pytest.raises(ValueError, match="not a directory"):
|
|
discover_models(file_path)
|
|
|
|
def test_unreadable_directory_is_skipped(self, tmp_path, monkeypatch):
|
|
"""Test that unreadable model roots do not crash discovery."""
|
|
original_iterdir = Path.iterdir
|
|
|
|
def fake_iterdir(path):
|
|
if path == tmp_path:
|
|
raise PermissionError("Operation not permitted")
|
|
return original_iterdir(path)
|
|
|
|
monkeypatch.setattr(Path, "iterdir", fake_iterdir)
|
|
|
|
assert discover_models(tmp_path) == {}
|
|
assert "not readable" in model_directory_access_error(tmp_path)
|
|
|
|
def test_unreadable_directory_skipped_in_multi_dir_discovery(
|
|
self, tmp_path, monkeypatch
|
|
):
|
|
"""Test that a readable directory still wins when another root is unreadable."""
|
|
unreadable = tmp_path / "unreadable"
|
|
unreadable.mkdir()
|
|
readable = tmp_path / "readable"
|
|
readable.mkdir()
|
|
model_dir = readable / "valid-model"
|
|
model_dir.mkdir()
|
|
(model_dir / "config.json").write_text(json.dumps({"model_type": "llama"}))
|
|
(model_dir / "model.safetensors").write_bytes(b"0" * 1000)
|
|
|
|
original_iterdir = Path.iterdir
|
|
|
|
def fake_iterdir(path):
|
|
if path == unreadable:
|
|
raise PermissionError("Operation not permitted")
|
|
return original_iterdir(path)
|
|
|
|
monkeypatch.setattr(Path, "iterdir", fake_iterdir)
|
|
|
|
models = discover_models_from_dirs([unreadable, readable])
|
|
assert list(models) == ["valid-model"]
|
|
|
|
def test_model_with_weight_error_skipped(self, tmp_path):
|
|
"""Test that models with no weights are skipped."""
|
|
# Model with config but no weights
|
|
model_dir = tmp_path / "no-weights"
|
|
model_dir.mkdir()
|
|
(model_dir / "config.json").write_text(json.dumps({"model_type": "llama"}))
|
|
|
|
models = discover_models(tmp_path)
|
|
assert len(models) == 0
|
|
|
|
def test_discovered_model_fields(self, tmp_path):
|
|
"""Test that DiscoveredModel has all expected fields."""
|
|
model_dir = tmp_path / "test-model"
|
|
model_dir.mkdir()
|
|
(model_dir / "config.json").write_text(json.dumps({"model_type": "llama"}))
|
|
(model_dir / "model.safetensors").write_bytes(b"0" * 1000)
|
|
|
|
models = discover_models(tmp_path)
|
|
model = models["test-model"]
|
|
|
|
assert model.model_id == "test-model"
|
|
assert model.model_path == str(model_dir)
|
|
assert model.model_type == "llm"
|
|
assert model.engine_type == "batched"
|
|
assert model.estimated_size == int(1000 * 1.05)
|
|
|
|
|
|
class TestFormatSize:
|
|
"""Tests for format_size function."""
|
|
|
|
def test_format_bytes(self):
|
|
"""Test formatting bytes."""
|
|
assert format_size(500) == "500.00B"
|
|
|
|
def test_format_kilobytes(self):
|
|
"""Test formatting kilobytes."""
|
|
assert format_size(1024) == "1.00KB"
|
|
assert format_size(2048) == "2.00KB"
|
|
|
|
def test_format_megabytes(self):
|
|
"""Test formatting megabytes."""
|
|
assert format_size(1024 * 1024) == "1.00MB"
|
|
assert format_size(5 * 1024 * 1024) == "5.00MB"
|
|
|
|
def test_format_gigabytes(self):
|
|
"""Test formatting gigabytes."""
|
|
assert format_size(1024 * 1024 * 1024) == "1.00GB"
|
|
assert format_size(32 * 1024 * 1024 * 1024) == "32.00GB"
|
|
|
|
def test_format_terabytes(self):
|
|
"""Test formatting terabytes."""
|
|
assert format_size(1024 * 1024 * 1024 * 1024) == "1.00TB"
|
|
|
|
def test_format_petabytes(self):
|
|
"""Test formatting petabytes."""
|
|
assert format_size(1024 * 1024 * 1024 * 1024 * 1024) == "1.00PB"
|
|
|
|
|
|
class TestAdapterDetection:
|
|
"""Tests for LoRA/PEFT adapter detection."""
|
|
|
|
def test_adapter_dir_detected(self, tmp_path):
|
|
"""Directory with adapter_config.json is detected as adapter."""
|
|
(tmp_path / "adapter_config.json").write_text("{}")
|
|
assert _is_adapter_dir(tmp_path) is True
|
|
|
|
def test_normal_model_not_adapter(self, tmp_path):
|
|
"""Normal model directory is not detected as adapter."""
|
|
(tmp_path / "config.json").write_text('{"model_type": "llama"}')
|
|
(tmp_path / "model.safetensors").write_bytes(b"0" * 1000)
|
|
assert _is_adapter_dir(tmp_path) is False
|
|
|
|
def test_discover_skips_lora_adapter(self, tmp_path):
|
|
"""discover_models should skip LoRA adapter directories."""
|
|
# Normal model
|
|
model_dir = tmp_path / "llama-3b"
|
|
model_dir.mkdir()
|
|
(model_dir / "config.json").write_text(json.dumps({"model_type": "llama"}))
|
|
(model_dir / "model.safetensors").write_bytes(b"0" * 1000)
|
|
|
|
# LoRA adapter (has both config.json and adapter_config.json)
|
|
adapter_dir = tmp_path / "my-lora"
|
|
adapter_dir.mkdir()
|
|
(adapter_dir / "config.json").write_text(json.dumps({"model_type": "qwen2"}))
|
|
(adapter_dir / "adapter_config.json").write_text("{}")
|
|
(adapter_dir / "adapters.safetensors").write_bytes(b"0" * 100)
|
|
|
|
models = discover_models(tmp_path)
|
|
assert "llama-3b" in models
|
|
assert "my-lora" not in models
|
|
|
|
def test_discover_skips_nested_lora_adapter(self, tmp_path):
|
|
"""discover_models should skip LoRA adapters in org folders."""
|
|
org_dir = tmp_path / "my-org"
|
|
org_dir.mkdir()
|
|
|
|
# Normal model under org
|
|
model_dir = org_dir / "llama-3b"
|
|
model_dir.mkdir()
|
|
(model_dir / "config.json").write_text(json.dumps({"model_type": "llama"}))
|
|
(model_dir / "model.safetensors").write_bytes(b"0" * 1000)
|
|
|
|
# LoRA adapter under org
|
|
adapter_dir = org_dir / "my-lora"
|
|
adapter_dir.mkdir()
|
|
(adapter_dir / "config.json").write_text(json.dumps({"model_type": "qwen2"}))
|
|
(adapter_dir / "adapter_config.json").write_text("{}")
|
|
|
|
models = discover_models(tmp_path)
|
|
assert "llama-3b" in models
|
|
assert "my-lora" not in models
|
|
|
|
|
|
class TestDiscoveredModel:
|
|
"""Tests for DiscoveredModel dataclass."""
|
|
|
|
def test_create_discovered_model(self):
|
|
"""Test creating a DiscoveredModel."""
|
|
model = DiscoveredModel(
|
|
model_id="test-model",
|
|
model_path="/path/to/model",
|
|
model_type="llm",
|
|
engine_type="batched",
|
|
estimated_size=1024 * 1024 * 1024, # 1GB
|
|
)
|
|
|
|
assert model.model_id == "test-model"
|
|
assert model.model_path == "/path/to/model"
|
|
assert model.model_type == "llm"
|
|
assert model.engine_type == "batched"
|
|
assert model.estimated_size == 1024 * 1024 * 1024
|
|
|
|
def test_discovered_model_embedding(self):
|
|
"""Test DiscoveredModel for embedding type."""
|
|
model = DiscoveredModel(
|
|
model_id="emb-model",
|
|
model_path="/path/to/emb",
|
|
model_type="embedding",
|
|
engine_type="embedding",
|
|
estimated_size=500 * 1024 * 1024,
|
|
)
|
|
|
|
assert model.model_type == "embedding"
|
|
assert model.engine_type == "embedding"
|
|
|
|
|
|
class TestReadModelContextLength:
|
|
"""Tests for _read_model_context_length — the discovery-time
|
|
config.json reader that backs #1308's API exposure fix."""
|
|
|
|
def _write(self, tmp_path, config=None, tokenizer_config=None):
|
|
if config is not None:
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
if tokenizer_config is not None:
|
|
(tmp_path / "tokenizer_config.json").write_text(json.dumps(tokenizer_config))
|
|
|
|
def test_max_position_embeddings_wins(self, tmp_path):
|
|
self._write(tmp_path, config={"max_position_embeddings": 262144})
|
|
assert _read_model_context_length(tmp_path) == 262144
|
|
|
|
def test_alternate_keys(self, tmp_path):
|
|
for key in ("max_seq_len", "max_seq_length", "seq_length", "n_positions"):
|
|
sub = tmp_path / key
|
|
sub.mkdir()
|
|
self._write(sub, config={key: 8192})
|
|
assert _read_model_context_length(sub) == 8192, key
|
|
|
|
def test_top_level_takes_precedence_over_nested(self, tmp_path):
|
|
self._write(tmp_path, config={
|
|
"max_position_embeddings": 200000,
|
|
"text_config": {"max_position_embeddings": 32768},
|
|
})
|
|
assert _read_model_context_length(tmp_path) == 200000
|
|
|
|
def test_text_config_fallback(self, tmp_path):
|
|
"""VLM wrappers stash the language head's ctx in text_config."""
|
|
self._write(tmp_path, config={"text_config": {"max_position_embeddings": 131072}})
|
|
assert _read_model_context_length(tmp_path) == 131072
|
|
|
|
def test_language_config_fallback(self, tmp_path):
|
|
"""Some Qwen-style configs use language_config instead."""
|
|
self._write(tmp_path, config={"language_config": {"max_position_embeddings": 65536}})
|
|
assert _read_model_context_length(tmp_path) == 65536
|
|
|
|
def test_tokenizer_max_length_fallback(self, tmp_path):
|
|
"""When config.json doesn't expose a length, tokenizer_config wins."""
|
|
self._write(tmp_path, config={"model_type": "llama"}, tokenizer_config={"model_max_length": 4096})
|
|
assert _read_model_context_length(tmp_path) == 4096
|
|
|
|
def test_tokenizer_sentinel_rejected(self, tmp_path):
|
|
"""Transformers' int(1e30) "no cap" sentinel must NOT leak through."""
|
|
self._write(
|
|
tmp_path,
|
|
config={"model_type": "llama"},
|
|
tokenizer_config={"model_max_length": int(1e30)},
|
|
)
|
|
assert _read_model_context_length(tmp_path) is None
|
|
|
|
def test_no_config_returns_none(self, tmp_path):
|
|
assert _read_model_context_length(tmp_path) is None
|
|
|
|
def test_invalid_types_rejected(self, tmp_path):
|
|
"""A string or zero in the config must not be returned as a length."""
|
|
self._write(tmp_path, config={"max_position_embeddings": "long"})
|
|
assert _read_model_context_length(tmp_path) is None
|
|
sub = tmp_path / "zero"
|
|
sub.mkdir()
|
|
self._write(sub, config={"max_position_embeddings": 0})
|
|
assert _read_model_context_length(sub) is None
|
|
|
|
def test_malformed_config_is_silent(self, tmp_path):
|
|
"""A broken config.json must not raise — discovery should keep going."""
|
|
(tmp_path / "config.json").write_text("{not valid json")
|
|
assert _read_model_context_length(tmp_path) is None
|
|
|
|
|
|
class TestTwoLevelDiscovery:
|
|
"""Tests for two-level model discovery."""
|
|
|
|
def _make_model(self, path: Path, model_type: str = "llama"):
|
|
"""Helper to create a valid model directory."""
|
|
path.mkdir(parents=True, exist_ok=True)
|
|
(path / "config.json").write_text(json.dumps({"model_type": model_type}))
|
|
(path / "model.safetensors").write_bytes(b"0" * 1000)
|
|
|
|
def test_two_level_org_folder(self, tmp_path):
|
|
"""Test discovery through organization folders."""
|
|
self._make_model(tmp_path / "mlx-community" / "llama-3b")
|
|
self._make_model(tmp_path / "mlx-community" / "qwen-7b")
|
|
|
|
models = discover_models(tmp_path)
|
|
assert len(models) == 2
|
|
assert "llama-3b" in models
|
|
assert "qwen-7b" in models
|
|
|
|
def test_mixed_flat_and_org(self, tmp_path):
|
|
"""Test mix of flat models and organization folders."""
|
|
# Flat model at level 1
|
|
self._make_model(tmp_path / "mistral-7b")
|
|
# Org folder with models at level 2
|
|
self._make_model(tmp_path / "Qwen" / "Qwen3-8B")
|
|
|
|
models = discover_models(tmp_path)
|
|
assert len(models) == 2
|
|
assert "mistral-7b" in models
|
|
assert "Qwen3-8B" in models
|
|
|
|
def test_multiple_org_folders(self, tmp_path):
|
|
"""Test multiple organization folders."""
|
|
self._make_model(tmp_path / "mlx-community" / "llama-3b")
|
|
self._make_model(tmp_path / "Qwen" / "Qwen3-8B")
|
|
self._make_model(tmp_path / "GLM" / "glm-4")
|
|
|
|
models = discover_models(tmp_path)
|
|
assert len(models) == 3
|
|
|
|
def test_empty_org_folder_skipped(self, tmp_path):
|
|
"""Test that empty org folders are silently skipped."""
|
|
self._make_model(tmp_path / "valid-model")
|
|
(tmp_path / "empty-org").mkdir()
|
|
|
|
models = discover_models(tmp_path)
|
|
assert len(models) == 1
|
|
assert "valid-model" in models
|
|
|
|
def test_org_folder_hidden_children_skipped(self, tmp_path):
|
|
"""Test that hidden subdirs inside org folders are skipped."""
|
|
org = tmp_path / "mlx-community"
|
|
self._make_model(org / "llama-3b")
|
|
self._make_model(org / ".hidden-model")
|
|
|
|
models = discover_models(tmp_path)
|
|
assert len(models) == 1
|
|
assert "llama-3b" in models
|
|
|
|
def test_org_folder_invalid_children_skipped(self, tmp_path):
|
|
"""Test that children without config.json in org folders are skipped."""
|
|
org = tmp_path / "mlx-community"
|
|
self._make_model(org / "llama-3b")
|
|
# Child without config.json
|
|
no_config = org / "broken-model"
|
|
no_config.mkdir(parents=True)
|
|
(no_config / "model.safetensors").write_bytes(b"0" * 1000)
|
|
|
|
models = discover_models(tmp_path)
|
|
assert len(models) == 1
|
|
|
|
def test_two_level_model_path_is_correct(self, tmp_path):
|
|
"""Test that model_path points to the actual model dir, not the org."""
|
|
self._make_model(tmp_path / "mlx-community" / "llama-3b")
|
|
|
|
models = discover_models(tmp_path)
|
|
assert models["llama-3b"].model_path == str(
|
|
tmp_path / "mlx-community" / "llama-3b"
|
|
)
|
|
|
|
|
|
class TestDiscoverModelsFromDirs:
|
|
"""Tests for discover_models_from_dirs function."""
|
|
|
|
def _make_model(self, path, model_type="llama"):
|
|
"""Helper to create a mock model directory."""
|
|
path.mkdir(parents=True, exist_ok=True)
|
|
config = {
|
|
"model_type": model_type,
|
|
"architectures": ["LlamaForCausalLM"],
|
|
}
|
|
(path / "config.json").write_text(json.dumps(config))
|
|
# Create a small safetensors file
|
|
(path / "model.safetensors").write_bytes(b"\x00" * 100)
|
|
|
|
def test_multiple_dirs(self, tmp_path):
|
|
"""Test discovering models from multiple directories."""
|
|
dir_a = tmp_path / "dir_a"
|
|
dir_b = tmp_path / "dir_b"
|
|
self._make_model(dir_a / "model-1")
|
|
self._make_model(dir_b / "model-2")
|
|
|
|
models = discover_models_from_dirs([dir_a, dir_b])
|
|
assert "model-1" in models
|
|
assert "model-2" in models
|
|
assert len(models) == 2
|
|
|
|
def test_first_directory_wins_on_conflict(self, tmp_path):
|
|
"""Test that first directory takes priority on model_id conflicts."""
|
|
dir_a = tmp_path / "dir_a"
|
|
dir_b = tmp_path / "dir_b"
|
|
self._make_model(dir_a / "same-model")
|
|
self._make_model(dir_b / "same-model")
|
|
|
|
models = discover_models_from_dirs([dir_a, dir_b])
|
|
assert len(models) == 1
|
|
assert models["same-model"].model_path == str(dir_a / "same-model")
|
|
|
|
def test_empty_list(self, tmp_path):
|
|
"""Test with empty directory list."""
|
|
models = discover_models_from_dirs([])
|
|
assert models == {}
|
|
|
|
def test_nonexistent_directory_skipped(self, tmp_path):
|
|
"""Test that non-existent directories are skipped with warning."""
|
|
dir_a = tmp_path / "dir_a"
|
|
self._make_model(dir_a / "model-1")
|
|
nonexistent = tmp_path / "does_not_exist"
|
|
|
|
models = discover_models_from_dirs([dir_a, nonexistent])
|
|
assert "model-1" in models
|
|
assert len(models) == 1
|
|
|
|
def test_mixed_valid_invalid_dirs(self, tmp_path):
|
|
"""Test with a mix of valid, empty, and non-existent directories."""
|
|
dir_valid = tmp_path / "valid"
|
|
dir_empty = tmp_path / "empty"
|
|
dir_empty.mkdir(parents=True)
|
|
self._make_model(dir_valid / "model-1")
|
|
|
|
models = discover_models_from_dirs(
|
|
[dir_valid, dir_empty, tmp_path / "nonexistent"]
|
|
)
|
|
assert "model-1" in models
|
|
assert len(models) == 1
|
|
|
|
|
|
class TestUnsupportedModels:
|
|
"""Tests for _is_unsupported_model() — audio models are now supported."""
|
|
|
|
def test_whisper_not_unsupported(self, tmp_path):
|
|
"""Whisper is now an audio_stt model, not unsupported."""
|
|
config = {
|
|
"model_type": "whisper",
|
|
"architectures": ["WhisperForConditionalGeneration"],
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert _is_unsupported_model(tmp_path) is False
|
|
|
|
def test_whisper_model_type_not_unsupported(self, tmp_path):
|
|
"""Whisper by model_type alone is not unsupported."""
|
|
config = {
|
|
"model_type": "whisper",
|
|
"architectures": ["SomeCustomWhisperArch"],
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert _is_unsupported_model(tmp_path) is False
|
|
|
|
def test_tts_not_unsupported(self, tmp_path):
|
|
"""qwen3_tts is now an audio_tts model, not unsupported."""
|
|
config = {
|
|
"model_type": "qwen3_tts",
|
|
"architectures": ["Qwen3TTSForConditionalGeneration"],
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert _is_unsupported_model(tmp_path) is False
|
|
|
|
def test_multimodal_with_audio_not_unsupported(self, tmp_path):
|
|
"""Multimodal model with nested audio_config is NOT unsupported."""
|
|
config = {
|
|
"model_type": "minicpmo",
|
|
"architectures": ["MiniCPMO"],
|
|
"vision_config": {"model_type": "siglip_vision_model"},
|
|
"audio_config": {"model_type": "whisper"},
|
|
"tts_config": {"model_type": "minicpmtts"},
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert _is_unsupported_model(tmp_path) is False
|
|
|
|
def test_llm_not_unsupported(self, tmp_path):
|
|
"""Regular LLM is not unsupported."""
|
|
config = {
|
|
"model_type": "llama",
|
|
"architectures": ["LlamaForCausalLM"],
|
|
}
|
|
(tmp_path / "config.json").write_text(json.dumps(config))
|
|
assert _is_unsupported_model(tmp_path) is False
|
|
|
|
def test_audio_models_included_in_discovery(self, tmp_path):
|
|
"""Audio models are now discovered (not skipped) with correct types."""
|
|
# Create a normal LLM model
|
|
llm_dir = tmp_path / "llama-3b"
|
|
llm_dir.mkdir()
|
|
(llm_dir / "config.json").write_text(
|
|
json.dumps({"model_type": "llama", "architectures": ["LlamaForCausalLM"]})
|
|
)
|
|
(llm_dir / "model.safetensors").write_bytes(b"0" * 1000)
|
|
|
|
# Create a whisper ASR model
|
|
asr_dir = tmp_path / "whisper-large-v3"
|
|
asr_dir.mkdir()
|
|
(asr_dir / "config.json").write_text(
|
|
json.dumps(
|
|
{
|
|
"model_type": "whisper",
|
|
"architectures": ["WhisperForConditionalGeneration"],
|
|
}
|
|
)
|
|
)
|
|
(asr_dir / "model.safetensors").write_bytes(b"0" * 2000)
|
|
|
|
# Create a TTS model
|
|
tts_dir = tmp_path / "Qwen3-TTS"
|
|
tts_dir.mkdir()
|
|
(tts_dir / "config.json").write_text(
|
|
json.dumps({"model_type": "qwen3_tts"})
|
|
)
|
|
(tts_dir / "model.safetensors").write_bytes(b"0" * 1500)
|
|
|
|
models = discover_models(tmp_path)
|
|
assert len(models) == 3
|
|
assert "llama-3b" in models
|
|
assert "whisper-large-v3" in models
|
|
assert models["whisper-large-v3"].model_type == "audio_stt"
|
|
assert "Qwen3-TTS" in models
|
|
assert models["Qwen3-TTS"].model_type == "audio_tts"
|
|
|
|
|
|
class TestHfCacheDiscovery:
|
|
"""Tests for HF Hub cache entry resolution and discovery."""
|
|
|
|
FAKE_COMMIT = "abc123def456"
|
|
|
|
def _make_model(self, path: Path, model_type: str = "llama"):
|
|
"""Helper to create a valid model directory."""
|
|
path.mkdir(parents=True, exist_ok=True)
|
|
(path / "config.json").write_text(json.dumps({"model_type": model_type}))
|
|
(path / "model.safetensors").write_bytes(b"0" * 1000)
|
|
|
|
def _make_hf_cache_entry(self, parent: Path, org: str, name: str):
|
|
"""Helper to create a bare HF Hub cache directory layout (no model files)."""
|
|
entry = parent / f"models--{org}--{name}"
|
|
refs = entry / "refs"
|
|
refs.mkdir(parents=True)
|
|
(refs / "main").write_text(self.FAKE_COMMIT)
|
|
snapshot = entry / "snapshots" / self.FAKE_COMMIT
|
|
snapshot.mkdir(parents=True)
|
|
return entry, snapshot
|
|
|
|
def _make_hf_cache_model(self, parent: Path, org: str, name: str, model_type: str = "llama"):
|
|
"""Helper to create an HF cache entry with a valid model in the snapshot."""
|
|
_, snapshot = self._make_hf_cache_entry(parent, org, name)
|
|
(snapshot / "config.json").write_text(json.dumps({"model_type": model_type}))
|
|
(snapshot / "model.safetensors").write_bytes(b"0" * 1000)
|
|
|
|
def test_resolve_valid_entry(self, tmp_path):
|
|
"""Valid HF cache entry resolves to snapshot path and repo metadata."""
|
|
entry, snapshot = self._make_hf_cache_entry(tmp_path, "mlx-community", "Qwen3-8B-4bit")
|
|
|
|
result = _resolve_hf_cache_entry(entry)
|
|
assert result is not None
|
|
assert result.snapshot_path == snapshot
|
|
assert result.model_id == "mlx-community--Qwen3-8B-4bit"
|
|
assert result.source_repo_id == "mlx-community/Qwen3-8B-4bit"
|
|
|
|
def test_resolve_regular_dir_returns_none(self, tmp_path):
|
|
"""Regular directory without models-- prefix returns None."""
|
|
regular = tmp_path / "mlx-community"
|
|
regular.mkdir()
|
|
assert _resolve_hf_cache_entry(regular) is None
|
|
|
|
def test_resolve_unnamespaced_repo(self, tmp_path):
|
|
"""models--Name (no org separator) is supported."""
|
|
entry, _ = self._make_hf_cache_entry(tmp_path, "", "gpt2")
|
|
entry.rename(tmp_path / "models--gpt2")
|
|
entry = tmp_path / "models--gpt2"
|
|
snapshot = entry / "snapshots" / self.FAKE_COMMIT
|
|
|
|
result = _resolve_hf_cache_entry(entry)
|
|
assert result is not None
|
|
assert result.snapshot_path == snapshot
|
|
assert result.model_id == "gpt2"
|
|
assert result.source_repo_id == "gpt2"
|
|
|
|
def test_resolve_missing_refs_main_returns_none(self, tmp_path):
|
|
"""Missing refs and snapshots returns None."""
|
|
entry = tmp_path / "models--mlx-community--Qwen3-8B"
|
|
entry.mkdir(parents=True)
|
|
assert _resolve_hf_cache_entry(entry) is None
|
|
|
|
def test_resolve_missing_snapshot_returns_none(self, tmp_path):
|
|
"""Valid refs/main but missing snapshot directory returns None."""
|
|
entry = tmp_path / "models--mlx-community--Qwen3-8B"
|
|
refs = entry / "refs"
|
|
refs.mkdir(parents=True)
|
|
(refs / "main").write_text("deadbeef")
|
|
assert _resolve_hf_cache_entry(entry) is None
|
|
|
|
def test_resolve_strips_whitespace_from_refs(self, tmp_path):
|
|
"""Trailing newline in refs/main is stripped (matches real HF cache)."""
|
|
entry, snapshot = self._make_hf_cache_entry(tmp_path, "mlx-community", "Qwen3-8B")
|
|
# Overwrite with trailing newline (like real HF cache)
|
|
(entry / "refs" / "main").write_text(self.FAKE_COMMIT + "\n")
|
|
|
|
result = _resolve_hf_cache_entry(entry)
|
|
assert result is not None
|
|
assert result.snapshot_path == snapshot
|
|
|
|
def test_discover_hf_cache_model(self, tmp_path):
|
|
"""HF cache entries are discovered as models."""
|
|
self._make_hf_cache_model(tmp_path, "mlx-community", "Qwen3-8B-4bit")
|
|
|
|
models = discover_models(tmp_path)
|
|
assert len(models) == 1
|
|
assert "mlx-community--Qwen3-8B-4bit" in models
|
|
assert models["mlx-community--Qwen3-8B-4bit"].model_type == "llm"
|
|
assert models["mlx-community--Qwen3-8B-4bit"].source_type == "hf_cache"
|
|
assert (
|
|
models["mlx-community--Qwen3-8B-4bit"].source_repo_id
|
|
== "mlx-community/Qwen3-8B-4bit"
|
|
)
|
|
|
|
def test_discover_multiple_hf_cache_models(self, tmp_path):
|
|
"""Multiple HF cache entries are all discovered."""
|
|
self._make_hf_cache_model(tmp_path, "mlx-community", "Qwen3-8B-4bit")
|
|
self._make_hf_cache_model(tmp_path, "mlx-community", "Mistral-7B-v0.3")
|
|
|
|
models = discover_models(tmp_path)
|
|
assert len(models) == 2
|
|
assert "mlx-community--Qwen3-8B-4bit" in models
|
|
assert "mlx-community--Mistral-7B-v0.3" in models
|
|
|
|
def test_hf_cache_model_path_points_to_snapshot(self, tmp_path):
|
|
"""model_path points to the snapshot dir, not the cache entry."""
|
|
self._make_hf_cache_model(tmp_path, "mlx-community", "Qwen3-8B-4bit")
|
|
|
|
models = discover_models(tmp_path)
|
|
assert models["mlx-community--Qwen3-8B-4bit"].model_path == str(
|
|
tmp_path / "models--mlx-community--Qwen3-8B-4bit" / "snapshots" / self.FAKE_COMMIT
|
|
)
|
|
|
|
def test_hf_cache_without_config_json_skipped(self, tmp_path):
|
|
"""HF cache entries without config.json in snapshot are skipped."""
|
|
self._make_hf_cache_entry(tmp_path, "mlx-community", "NoConfig")
|
|
|
|
models = discover_models(tmp_path)
|
|
assert len(models) == 0
|
|
|
|
def test_hf_cache_non_mlx_repo_skipped(self, tmp_path):
|
|
"""HF cache entries without MLX metadata or repo-name hints are skipped."""
|
|
self._make_hf_cache_model(tmp_path, "acme", "PlainModel")
|
|
|
|
models = discover_models(tmp_path)
|
|
assert len(models) == 0
|
|
|
|
def test_hf_cache_bad_helper_schema_is_skipped(self, tmp_path):
|
|
"""A malformed helper-looking schema must not abort HF cache discovery."""
|
|
_, snapshot = self._make_hf_cache_entry(tmp_path, "acme", "BadSchema")
|
|
(snapshot / "config.json").write_text(
|
|
json.dumps({"model_type": "qwen3", "architectures": 123})
|
|
)
|
|
(snapshot / "model.safetensors").write_bytes(b"0" * 1000)
|
|
|
|
models = discover_models(tmp_path)
|
|
assert models == {}
|
|
|
|
def test_hf_cache_mlx_metadata_is_discovered(self, tmp_path):
|
|
"""HF cache entries with safetensors format=mlx metadata are discovered."""
|
|
np = pytest.importorskip("numpy")
|
|
safetensors_numpy = pytest.importorskip("safetensors.numpy")
|
|
|
|
entry, snapshot = self._make_hf_cache_entry(tmp_path, "acme", "PlainModel")
|
|
(snapshot / "config.json").write_text(json.dumps({"model_type": "llama"}))
|
|
safetensors_numpy.save_file(
|
|
{"weight": np.zeros((1,), dtype=np.float32)},
|
|
snapshot / "model.safetensors",
|
|
metadata={"format": "mlx"},
|
|
)
|
|
|
|
models = discover_models(tmp_path)
|
|
assert len(models) == 1
|
|
assert "acme--PlainModel" in models
|
|
assert models["acme--PlainModel"].source_repo_id == "acme/PlainModel"
|
|
|
|
def test_hf_cache_dflash_draft_is_discovered(self, tmp_path):
|
|
"""HF cache DFlash draft checkpoints are discovered despite pt-format
|
|
safetensors and a repo name with no 'mlx' token (#1643)."""
|
|
np = pytest.importorskip("numpy")
|
|
safetensors_numpy = pytest.importorskip("safetensors.numpy")
|
|
|
|
entry, snapshot = self._make_hf_cache_entry(
|
|
tmp_path, "z-lab", "Qwen3.6-27B-DFlash"
|
|
)
|
|
(snapshot / "config.json").write_text(
|
|
json.dumps(
|
|
{
|
|
"model_type": "qwen3",
|
|
"architectures": ["DFlashDraftModel"],
|
|
"dflash_config": {},
|
|
}
|
|
)
|
|
)
|
|
# pt-format safetensors (NOT mlx) so only the DFlash architecture — not
|
|
# safetensors metadata or the repo name — can qualify this entry.
|
|
safetensors_numpy.save_file(
|
|
{"weight": np.zeros((1,), dtype=np.float32)},
|
|
snapshot / "model.safetensors",
|
|
metadata={"format": "pt"},
|
|
)
|
|
|
|
models = discover_models(tmp_path)
|
|
assert "z-lab--Qwen3.6-27B-DFlash" in models
|
|
# The shared DiscoveredModel construction path must flag the draft as
|
|
# a helper so it lands in the draft picker, not the chat-model list.
|
|
assert models["z-lab--Qwen3.6-27B-DFlash"].is_helper is True
|
|
|
|
def test_is_helper_checkpoint_on_disk_edge_cases(self, tmp_path):
|
|
"""The on-disk wrapper qualifies a drafter config and never raises on
|
|
unreadable entries (it runs outside discover_models' per-model guard).
|
|
|
|
Marker-family coverage lives in TestIsHelperModelConfig; this pins
|
|
only the wrapper's own file-handling behavior.
|
|
"""
|
|
draft = tmp_path / "draft"
|
|
draft.mkdir()
|
|
(draft / "config.json").write_text(
|
|
json.dumps({"model_type": "qwen3", "architectures": ["DFlashDraftModel"]})
|
|
)
|
|
assert _is_helper_checkpoint(draft) is True
|
|
|
|
# plain (non-draft) model config stays excluded
|
|
plain = tmp_path / "plain"
|
|
plain.mkdir()
|
|
(plain / "config.json").write_text(
|
|
json.dumps({"model_type": "llama", "architectures": ["LlamaForCausalLM"]})
|
|
)
|
|
assert _is_helper_checkpoint(plain) is False
|
|
|
|
# no config.json at all
|
|
empty = tmp_path / "empty"
|
|
empty.mkdir()
|
|
assert _is_helper_checkpoint(empty) is False
|
|
|
|
# malformed JSON must not raise
|
|
bad_json = tmp_path / "bad-json"
|
|
bad_json.mkdir()
|
|
(bad_json / "config.json").write_text("{not json")
|
|
assert _is_helper_checkpoint(bad_json) is False
|
|
|
|
# non-UTF-8 bytes must not raise: UnicodeDecodeError is a ValueError,
|
|
# not a JSONDecodeError, and an escape here aborts the whole scan
|
|
bad_bytes = tmp_path / "bad-bytes"
|
|
bad_bytes.mkdir()
|
|
(bad_bytes / "config.json").write_bytes(b'\xff\xfe{"a": 1}')
|
|
assert _is_helper_checkpoint(bad_bytes) is False
|
|
|
|
# valid JSON with an unexpected schema must not raise either
|
|
bad_schema = tmp_path / "bad-schema"
|
|
bad_schema.mkdir()
|
|
(bad_schema / "config.json").write_text(
|
|
json.dumps({"model_type": "qwen3", "architectures": 123})
|
|
)
|
|
assert _is_helper_checkpoint(bad_schema) is False
|
|
|
|
def test_mixed_flat_and_hf_cache(self, tmp_path):
|
|
"""Mix of flat models and HF cache entries."""
|
|
self._make_model(tmp_path / "mistral-7b")
|
|
self._make_hf_cache_model(tmp_path, "mlx-community", "Qwen3-8B-4bit")
|
|
|
|
models = discover_models(tmp_path)
|
|
assert len(models) == 2
|
|
assert "mistral-7b" in models
|
|
assert "mlx-community--Qwen3-8B-4bit" in models
|
|
|
|
def test_mixed_org_and_hf_cache(self, tmp_path):
|
|
"""Mix of org folders and HF cache entries."""
|
|
self._make_model(tmp_path / "Qwen" / "Qwen3-8B", model_type="qwen2")
|
|
self._make_hf_cache_model(tmp_path, "mlx-community", "Mistral-7B")
|
|
|
|
models = discover_models(tmp_path)
|
|
assert len(models) == 2
|
|
assert "Qwen3-8B" in models
|
|
assert "mlx-community--Mistral-7B" in models
|
|
|
|
def test_hf_cache_does_not_fall_through_to_org_scan(self, tmp_path):
|
|
"""HF cache entries don't get scanned as org folders."""
|
|
self._make_hf_cache_model(tmp_path, "mlx-community", "Qwen3-8B-4bit")
|
|
|
|
models = discover_models(tmp_path)
|
|
assert len(models) == 1
|
|
assert "mlx-community--Qwen3-8B-4bit" in models
|