add model search method

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}"

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

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!")
+        print("\u2713 All models found!")