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

add mistral small to registry

Browse source
This commit is contained in:
jeromeku 2025-03-31 15:31:01 -07:00
commit 0236e05307
2 changed files with 71 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

9
unsloth/registry/_deepseek.py
View file

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

66
unsloth/registry/_mistral.py Normal file
View file

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