clear registry when executing individual model registration file

tests/test_model_registry.py

@@ -13,7 +13,7 @@ from dataclasses import dataclass
 import pytest
 from huggingface_hub import ModelInfo as HfModelInfo
 
-from unsloth.registry import get_model_info, register_models
+from unsloth.registry import register_models, search_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
@@ -21,6 +21,7 @@ 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, QUANT_TAG_MAP, QuantType
+from unsloth.utils.hf_hub import get_model_info
 
 MODEL_NAMES = [
     "llama",
@@ -84,7 +85,7 @@ def test_all_model_registration():
 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])
+    dynamic_quant_models = search_models(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)

unsloth/registry/REGISTRY.md

@@ -1,6 +1,16 @@
 ## Model Registry
 
 ### Structure
+```
+unsloth
+    -registry
+        __init__.py
+        registry.py
+        _llama.py
+        _mistral.py
+        _phi.py
+        ...
+```
 
 Each model is registered in a separate file within the `registry` module (e.g. `registry/_llama.py`).
 
@@ -41,5 +51,43 @@ class LlamaModelInfo(ModelInfo):
         return super().construct_model_name(base_name, version, size, quant_type, instruct_tag, key)
 ```
 
-Once these constructs are defined, the model is registered in the `registry` module by calling `register_models` with the `ModelMeta` and `ModelInfo` classes.
+Once these constructs are defined, the model is registered by writing a register_xx_models function.
+```python
+def register_llama_3_1_models(include_original_model: bool = False):
+    global _IS_LLAMA_3_1_REGISTERED
+    if _IS_LLAMA_3_1_REGISTERED:
+        return
+    _register_models(LlamaMeta_3_1, include_original_model=include_original_model)
+    _IS_LLAMA_3_1_REGISTERED = True
+```
+
+`_register_models` is a helper function that registers the model with the registry.  The global `_IS_XX_REGISTERED` is used to prevent duplicate registration.
+
+Once a model is registered, registry.registry.MODEL_REGISTRY is updated with the model info and can be searched with `registry.search_models`.
+
+### Tests
+
+The `tests/test_model_registry.py` file contains tests for the model registry.
+
+Also, each model registration file is an executable module that checks that all registered models are available on `huggingface_hub`.
+```python
+python unsloth.registry._llama.py
+```
+
+Prints the following (abridged) output:
+```bash
+✓ unsloth/Llama-3.1-8B
+✓ unsloth/Llama-3.1-8B-bnb-4bit
+✓ unsloth/Llama-3.1-8B-unsloth-bnb-4bit
+✓ meta-llama/Llama-3.1-8B
+✓ unsloth/Llama-3.1-8B-Instruct
+✓ unsloth/Llama-3.1-8B-Instruct-bnb-4bit
+✓ unsloth/Llama-3.1-8B-Instruct-unsloth-bnb-4bit
+✓ meta-llama/Llama-3.1-8B-Instruct
+✓ unsloth/Llama-3.2-1B
+✓ unsloth/Llama-3.2-1B-bnb-4bit
+✓ unsloth/Llama-3.2-1B-unsloth-bnb-4bit
+✓ meta-llama/Llama-3.2-1B
+...
+```
 

unsloth/registry/__init__.py

@@ -22,7 +22,7 @@ def register_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]:
+def search_models(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.
 

unsloth/registry/_deepseek.py

@@ -163,6 +163,10 @@ register_deepseek_models(include_original_model=True)
 
 if __name__ == "__main__":
     from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info
+    MODEL_REGISTRY.clear()
+    
+    register_deepseek_models(include_original_model=True)
+    
     for model_id, model_info in MODEL_REGISTRY.items():
         model_info = _check_model_info(model_id)
         if model_info is None:

unsloth/registry/_gemma.py

@@ -53,8 +53,11 @@ def register_gemma_models(include_original_model: bool = False):
 
 
 if __name__ == "__main__":
-    register_gemma_models(include_original_model=True)
     from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info
+    MODEL_REGISTRY.clear()
+    
+    register_gemma_models(include_original_model=True)
+    
     for model_id, model_info in MODEL_REGISTRY.items():
         model_info = _check_model_info(model_id)
         if model_info is None:

unsloth/registry/_llama.py

@@ -100,8 +100,10 @@ def register_llama_models(include_original_model: bool = False):
     register_llama_3_2_vision_models(include_original_model=include_original_model)
 
 if __name__ == "__main__":
-    register_llama_models(include_original_model=True)
     from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info
+    MODEL_REGISTRY.clear()
+    
+    register_llama_models(include_original_model=True)
 
     for model_id, model_info in MODEL_REGISTRY.items():
         model_info = _check_model_info(model_id)

unsloth/registry/_mistral.py

@@ -57,8 +57,11 @@ def register_mistral_models(include_original_model: bool = False):
     register_mistral_small_models(include_original_model=include_original_model)
 
 if __name__ == "__main__":
-    register_mistral_models(include_original_model=True)
     from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info
+    MODEL_REGISTRY.clear()
+    
+    register_mistral_models(include_original_model=True)
+    
     for model_id, model_info in MODEL_REGISTRY.items():
         model_info = _check_model_info(model_id)
         if model_info is None:

unsloth/registry/_phi.py

@@ -52,8 +52,11 @@ def register_phi_models(include_original_model: bool = False):
     register_phi_4_instruct_models(include_original_model=include_original_model)
 
 if __name__ == "__main__":
-    register_phi_models(include_original_model=True)
     from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info
+    MODEL_REGISTRY.clear()
+    
+    register_phi_models(include_original_model=True)
+    
     for model_id, model_info in MODEL_REGISTRY.items():
         model_info = _check_model_info(model_id)
         if model_info is None:

unsloth/registry/_qwen.py

@@ -104,8 +104,11 @@ def register_qwen_models(include_original_model: bool = False):
     register_qwen_qwq_models(include_original_model=include_original_model)
 
 if __name__ == "__main__":
-    register_qwen_models(include_original_model=True)
     from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info
+    MODEL_REGISTRY.clear()
+    
+    register_qwen_models(include_original_model=True)
+    
     for model_id, model_info in MODEL_REGISTRY.items():
         model_info = _check_model_info(model_id)
         if model_info is None: