diff --git a/unsloth/registry/_deepseek.py b/unsloth/registry/_deepseek.py index 8bdcd3c2e7..0346520060 100644 --- a/unsloth/registry/_deepseek.py +++ b/unsloth/registry/_deepseek.py @@ -1,6 +1,7 @@ from unsloth.registry.registry import ModelInfo, ModelMeta, QuantType, _register_models _IS_DEEPSEEKV3_REGISTERED = False +_IS_DEEPSEEKR1_REGISTERED = False class DeepseekV3ModelInfo(ModelInfo): @classmethod @@ -10,6 +11,14 @@ class DeepseekV3ModelInfo(ModelInfo): key = cls.append_quant_type(key, quant_type) return key +class DeepseekR1ModelInfo(ModelInfo): + @classmethod + def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag): + key = f"{base_name}-{version}" if version else base_name + key = cls.append_instruct_tag(key, instruct_tag) + key = cls.append_quant_type(key, quant_type) + return key + # Deepseek V3 Model Meta DeepseekV3Meta = ModelMeta( org="deepseek-ai", @@ -33,6 +42,17 @@ DeepseekV3_0324Meta = ModelMeta( quant_types=[QuantType.NONE, QuantType.GGUF], ) +DeepseekR1Meta = ModelMeta( + org="deepseek-ai", + base_name="DeepSeek-R1", + instruct_tags=[None], + model_version="", + model_sizes=[""], + model_info_cls=DeepseekR1ModelInfo, + is_multimodal=False, + quant_types=[QuantType.NONE, QuantType.BF16, QuantType.GGUF], +) + def register_deepseek_v3_models(include_original_model: bool = False): global _IS_DEEPSEEKV3_REGISTERED if _IS_DEEPSEEKV3_REGISTERED: @@ -41,7 +61,17 @@ def register_deepseek_v3_models(include_original_model: bool = False): _register_models(DeepseekV3_0324Meta, include_original_model=include_original_model) _IS_DEEPSEEKV3_REGISTERED = True -register_deepseek_v3_models(include_original_model=True) + +def register_deepseek_r1_models(include_original_model: bool = False): + global _IS_DEEPSEEKR1_REGISTERED + if _IS_DEEPSEEKR1_REGISTERED: + return + _register_models(DeepseekR1Meta, include_original_model=include_original_model) + _IS_DEEPSEEKR1_REGISTERED = True + +#register_deepseek_v3_models(include_original_model=True) +register_deepseek_r1_models(include_original_model=True) + if __name__ == "__main__": from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info diff --git a/unsloth/registry/registry.py b/unsloth/registry/registry.py index 1eee884259..1e2c667e13 100644 --- a/unsloth/registry/registry.py +++ b/unsloth/registry/registry.py @@ -35,7 +35,8 @@ class ModelInfo: is_multimodal: bool = False instruct_tag: str = None quant_type: QuantType = None - + description: str = None + def __post_init__(self): self.name = self.name or self.construct_model_name( self.base_name,