ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

add quant type test

Browse source
This commit is contained in:
jeromeku 2025-03-31 17:58:44 -07:00
commit 1e6b5f02b3
1 changed files with 10 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

13
tests/test_model_registry.py
View file

@ -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)