diff --git a/tests/test_model_registry.py b/tests/test_model_registry.py index a767d42cdb..3d570af230 100644 --- a/tests/test_model_registry.py +++ b/tests/test_model_registry.py @@ -13,15 +13,14 @@ from dataclasses import dataclass import pytest from huggingface_hub import ModelInfo as HfModelInfo -from unsloth.registry import register_models +from unsloth.registry import get_model_info, register_models from unsloth.registry._deepseek import register_deepseek_models from unsloth.registry._gemma import register_gemma_models from unsloth.registry._llama import register_llama_models from unsloth.registry._mistral import register_mistral_models from unsloth.registry._phi import register_phi_models from unsloth.registry._qwen import register_qwen_models -from unsloth.registry.registry import MODEL_REGISTRY -from unsloth.utils.hf_hub import get_model_info +from unsloth.registry.registry import MODEL_REGISTRY, QUANT_TAG_MAP, QuantType MODEL_NAMES = [ "llama", @@ -81,3 +80,11 @@ def test_all_model_registration(): registered_models = MODEL_REGISTRY.keys() missing_models = _test_model_uploaded(registered_models) assert not missing_models, f"Missing following models: {missing_models}" + +def test_quant_type(): + # Test that the quant_type is correctly set for model paths + # NOTE: for models registered under org="unsloth" with QuantType.NONE aliases QuantType.UNSLOTH + dynamic_quant_models = get_model_info(quant_types=[QuantType.UNSLOTH]) + assert all(m.quant_type == QuantType.UNSLOTH for m in dynamic_quant_models) + quant_tag = QUANT_TAG_MAP[QuantType.UNSLOTH] + assert all(quant_tag in m.model_path for m in dynamic_quant_models) \ No newline at end of file