diff --git a/tests/test_model_registry.py b/tests/test_model_registry.py index 1f9ddd922e..a767d42cdb 100644 --- a/tests/test_model_registry.py +++ b/tests/test_model_registry.py @@ -1,3 +1,13 @@ +""" + +Test model registration methods +Checks that model registration methods work for respective models as well as all models +The check is performed +- by registering the models +- checking that the instantiated models can be found on huggingface hub by querying for the model id + +""" + from dataclasses import dataclass import pytest @@ -10,7 +20,7 @@ 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, ModelInfo +from unsloth.registry.registry import MODEL_REGISTRY from unsloth.utils.hf_hub import get_model_info MODEL_NAMES = [ @@ -34,7 +44,7 @@ MODEL_REGISTRATION_METHODS = [ @dataclass class ModelTestParam: name: str - registration_models: callable + register_models: callable def _test_model_uploaded(model_ids: list[str]): @@ -52,13 +62,13 @@ TestParams = [ for name, models in zip(MODEL_NAMES, MODEL_REGISTRATION_METHODS) ] + # Test that model registration methods register respective models -@pytest.mark.parametrize( - "model_test_param", TestParams, ids=lambda param: param.name -) +@pytest.mark.parametrize("model_test_param", TestParams, ids=lambda param: param.name) def test_model_registration(model_test_param: ModelTestParam): MODEL_REGISTRY.clear() - model_test_param.registration_models() + registration_method = model_test_param.register_models + registration_method() registered_models = MODEL_REGISTRY.keys() missing_models = _test_model_uploaded(registered_models) assert not missing_models, ( @@ -66,23 +76,8 @@ def test_model_registration(model_test_param: ModelTestParam): ) -# if __name__ == "__main__": -# for method in [ -# get_llama_models, -# get_llama_vision_models, -# get_qwen_models, -# get_qwen_vl_models, -# get_phi_models, -# get_phi_instruct_models, -# ]: -# models = method() -# model_name = next(iter(models.values())).base_name -# print(f"{model_name}: {len(models)} registered") -# for model_info in models.values(): -# print(f" {model_info.model_path}") -# missing_models = test_model_uploaded(list(models.keys())) - -# if missing_models: -# print("--------------------------------") -# print(f"Missing models: {missing_models}") -# print("--------------------------------") +def test_all_model_registration(): + register_models() + registered_models = MODEL_REGISTRY.keys() + missing_models = _test_model_uploaded(registered_models) + assert not missing_models, f"Missing following models: {missing_models}" diff --git a/unsloth/registry/__init__.py b/unsloth/registry/__init__.py index 1b92fef74d..a46ab773d8 100644 --- a/unsloth/registry/__init__.py +++ b/unsloth/registry/__init__.py @@ -4,9 +4,15 @@ from ._llama import register_llama_models as _register_llama_models from ._mistral import register_mistral_models as _register_mistral_models from ._phi import register_phi_models as _register_phi_models from ._qwen import register_qwen_models as _register_qwen_models +from .registry import MODEL_REGISTRY, ModelInfo, QuantType +_ARE_MODELS_REGISTERED = False -def register_models(): +def register_models(): + global _ARE_MODELS_REGISTERED + + if _ARE_MODELS_REGISTERED: + return _register_deepseek_models() _register_gemma_models() _register_llama_models() @@ -14,3 +20,32 @@ def register_models(): _register_phi_models() _register_qwen_models() + _ARE_MODELS_REGISTERED = True + +def get_model_info(org: str = None, base_name: str = None, version: str = None, size: str = None, quant_types: list[QuantType] = None, search_pattern: str = None) -> list[ModelInfo]: + """ + Get model info from the registry. + + See registry.ModelInfo for more fields. + + If search_pattern is provided, the full model path will be matched against the pattern, where the model path is the model_id on huggingface hub. + + """ + if not _ARE_MODELS_REGISTERED: + register_models() + + model_infos = MODEL_REGISTRY.values() + if org: + model_infos = [model_info for model_info in model_infos if model_info.org == org] + if base_name: + model_infos = [model_info for model_info in model_infos if model_info.base_name == base_name] + if version: + model_infos = [model_info for model_info in model_infos if model_info.version == version] + if size: + model_infos = [model_info for model_info in model_infos if model_info.size == size] + if quant_types: + model_infos = [model_info for model_info in model_infos if any(model_info.quant_type == quant_type for quant_type in quant_types)] + if search_pattern: + model_infos = [model_info for model_info in model_infos if search_pattern in model_info.model_path] + + return model_infos \ No newline at end of file diff --git a/unsloth/registry/model_registry.py b/unsloth/registry/model_registry.py index de9609934c..b51644beb7 100644 --- a/unsloth/registry/model_registry.py +++ b/unsloth/registry/model_registry.py @@ -306,4 +306,4 @@ if __name__ == "__main__": if len(missing_models) == 0: # print unicode checkmark - print(f"\u2713 All models found!") \ No newline at end of file + print("\u2713 All models found!") \ No newline at end of file