From 0236e053077e1827fc3133491fd0de29a0ba99af Mon Sep 17 00:00:00 2001
From: jeromeku <jerome.ku@gmail.com>
Date: Mon, 31 Mar 2025 15:31:01 -0700
Subject: [PATCH] add mistral small to registry

---
 unsloth/registry/_deepseek.py |  9 ++---
 unsloth/registry/_mistral.py  | 66 +++++++++++++++++++++++++++++++++++
 2 files changed, 71 insertions(+), 4 deletions(-)
 create mode 100644 unsloth/registry/_mistral.py

diff --git a/unsloth/registry/_deepseek.py b/unsloth/registry/_deepseek.py
index 148093155c..35cbc17484 100644
--- a/unsloth/registry/_deepseek.py
+++ b/unsloth/registry/_deepseek.py
@@ -138,10 +138,6 @@ def register_deepseek_r1_distill_models(include_original_model: bool = False):
     register_deepseek_r1_distill_qwen_models(include_original_model=include_original_model)
     register_deepseek_r1_distill_llama_models(include_original_model=include_original_model)
 
-register_deepseek_v3_models(include_original_model=True)
-register_deepseek_r1_models(include_original_model=True)
-register_deepseek_r1_distill_models(include_original_model=True)
-
 def _list_deepseek_r1_distill_models():
     from unsloth.utils.hf_hub import ModelInfo as HfModelInfo
     from unsloth.utils.hf_hub import list_models
@@ -156,6 +152,11 @@ def _list_deepseek_r1_distill_models():
 
     return distill_models
 
+
+register_deepseek_v3_models(include_original_model=True)
+register_deepseek_r1_models(include_original_model=True)
+register_deepseek_r1_distill_models(include_original_model=True)
+
 if __name__ == "__main__":
     from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info
     for model_id, model_info in MODEL_REGISTRY.items():
diff --git a/unsloth/registry/_mistral.py b/unsloth/registry/_mistral.py
new file mode 100644
index 0000000000..65f1256708
--- /dev/null
+++ b/unsloth/registry/_mistral.py
@@ -0,0 +1,66 @@
+import copy
+
+from unsloth.registry.registry import ModelInfo, ModelMeta, QuantType, _register_models
+
+_IS_MISTRAL_SMALL_REGISTERED = False
+
+_MISTRAL_SMALL_03_25_VERSION = "2503"
+_MISTRAL_SMALL_01_25_VERSION = "2501"
+_MISTRAL_SMALL_09_24_VERSION = "2409" # Not uploaded to unsloth
+
+class MistralSmallModelInfo(ModelInfo):
+    @classmethod
+    def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag):
+        if version == _MISTRAL_SMALL_03_25_VERSION:
+            key = f"{base_name}-3.1-{size}B-{instruct_tag}"
+        else:
+            key = f"{base_name}-{size}B-{instruct_tag}"
+        key += f"-{version}"
+        key = cls.append_quant_type(key, quant_type)
+        
+        return key
+
+
+MistralSmall_2503_Base_Meta = ModelMeta(
+    org="mistralai",
+    base_name="Mistral-Small",
+    instruct_tags=["Base"],
+    model_version=_MISTRAL_SMALL_03_25_VERSION,
+    model_sizes=["24"],
+    model_info_cls=MistralSmallModelInfo,
+    is_multimodal=False,
+    quant_types=[QuantType.NONE, QuantType.UNSLOTH, QuantType.BNB],
+)
+
+MistralSmall_2503_Instruct_Meta = copy.deepcopy(MistralSmall_2503_Base_Meta)
+MistralSmall_2503_Instruct_Meta.instruct_tags = ["Instruct"]
+MistralSmall_2503_Instruct_Meta.quant_types = [QuantType.NONE, QuantType.UNSLOTH, QuantType.BNB, QuantType.GGUF]
+
+MistralSmall_2501_Base_Meta = copy.deepcopy(MistralSmall_2503_Base_Meta)
+MistralSmall_2501_Base_Meta.model_version = _MISTRAL_SMALL_01_25_VERSION
+
+MistralSmall_2501_Instruct_Meta = copy.deepcopy(MistralSmall_2503_Instruct_Meta)
+MistralSmall_2501_Instruct_Meta.model_version = _MISTRAL_SMALL_01_25_VERSION
+
+def register_mistral_small_models():
+    global _IS_MISTRAL_SMALL_REGISTERED
+    if _IS_MISTRAL_SMALL_REGISTERED:
+        return
+    _register_models(MistralSmall_2503_Base_Meta)
+    _register_models(MistralSmall_2503_Instruct_Meta)
+    _register_models(MistralSmall_2501_Base_Meta)
+    _register_models(MistralSmall_2501_Instruct_Meta)
+
+    _IS_MISTRAL_SMALL_REGISTERED = True
+
+register_mistral_small_models()
+
+
+if __name__ == "__main__":
+    from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info
+    for model_id, model_info in MODEL_REGISTRY.items():
+        model_info = _check_model_info(model_id)
+        if model_info is None:
+            print(f"\u2718 {model_id}")
+        else:
+            print(f"\u2713 {model_id}")    
\ No newline at end of file