diff --git a/unsloth/registry/_deepseek.py b/unsloth/registry/_deepseek.py index bd0ea31cbf..b3bf398cf1 100644 --- a/unsloth/registry/_deepseek.py +++ b/unsloth/registry/_deepseek.py @@ -2,7 +2,8 @@ from unsloth.registry.registry import ModelInfo, ModelMeta, QuantType, _register _IS_DEEPSEEKV3_REGISTERED = False _IS_DEEPSEEKR1_REGISTERED = False - +_IS_DEEPSEEKR1_ZERO_REGISTERED = False +_IS_DEEPSEEKR1_DISTILL_REGISTERED = False class DeepseekV3ModelInfo(ModelInfo): @classmethod def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag): @@ -15,6 +16,8 @@ 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 + if size: + key = f"{key}-{size}B" key = cls.append_instruct_tag(key, instruct_tag) key = cls.append_quant_type(key, quant_type) return key @@ -63,6 +66,28 @@ DeepseekR1ZeroMeta = ModelMeta( is_multimodal=False, quant_types=[QuantType.NONE, QuantType.GGUF], ) + +DeepseekR1DistillMeta = ModelMeta( + org="deepseek-ai", + base_name="DeepSeek-R1-Distill", + instruct_tags=[None], + model_version="Llama", + model_sizes=["8", "70"], + model_info_cls=DeepseekR1ModelInfo, + is_multimodal=False, + quant_types={"8": [QuantType.UNSLOTH, QuantType.GGUF], "70": [QuantType.GGUF]}, +) + + # "Qwen-7B-unsloth-bnb-4bit", + # "Qwen-1.5B-unsloth-bnb-4bit", + # "Qwen-32B-GGUF", + # "Llama-8B-GGUF", + # "Qwen-14B-GGUF", + # "Qwen-32B-bnb-4bit", + # "Qwen-1.5B-GGUF", + # "Qwen-14B-unsloth-bnb-4bit", + # "Llama-70B-GGUF" + def register_deepseek_v3_models(include_original_model: bool = False): global _IS_DEEPSEEKV3_REGISTERED if _IS_DEEPSEEKV3_REGISTERED: @@ -78,11 +103,22 @@ def register_deepseek_r1_models(include_original_model: bool = False): return _register_models(DeepseekR1Meta, include_original_model=include_original_model) _register_models(DeepseekR1ZeroMeta, include_original_model=include_original_model) + _register_models(DeepseekR1DistillMeta, 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) +def _list_deepseek_r1_distill_models(): + from unsloth.utils.hf_hub import ModelInfo as HfModelInfo + from unsloth.utils.hf_hub import list_models + models: list[HfModelInfo] = list_models(author="unsloth", search="Distill") + for model in models: + model_id = model.id + model_name = model_id.split("/")[-1] + # parse out only the version + version = model_name.removeprefix("DeepSeek-R1-Distill-") + print(version) if __name__ == "__main__": from unsloth.registry.registry import MODEL_REGISTRY, _check_model_info diff --git a/unsloth/utils/hf_hub.py b/unsloth/utils/hf_hub.py index da3f72a18e..30255b8636 100644 --- a/unsloth/utils/hf_hub.py +++ b/unsloth/utils/hf_hub.py @@ -1,6 +1,6 @@ from huggingface_hub import HfApi, ModelInfo -api: HfApi +_HFAPI: HfApi = None POPULARITY_PROPERTIES = [ "downloads", @@ -32,27 +32,27 @@ def get_model_info( Default properties: ["safetensors", "lastModified"], only retrieves minimal information. Set to None to retrieve the full model information. """ - global api - if api is None: - api = HfApi() + global _HFAPI + if _HFAPI is None: + _HFAPI = HfApi() try: - model_info: ModelInfo = api.model_info(model_id, expand=properties) + model_info: ModelInfo = _HFAPI.model_info(model_id, expand=properties) except Exception as e: print(f"Error getting model info for {model_id}: {e}") model_info = None return model_info -def retrieve_models( +def list_models( properties: list[str] = None, full: bool = False, sort: str = "downloads", author: str = "unsloth", search: str = None, limit: int = 10, -) -> ModelInfo: +) -> list[ModelInfo]: """ - Retrieve models from the Hugging Face Hub. + Retrieve model information from the Hugging Face Hub. properties: list[str] = See https://huggingface.co/docs/huggingface_hub/api-ref/hf_hub/hf_api/list_models full: bool = Whether to retrieve the full model information, if True properties will be ignored. @@ -61,13 +61,13 @@ def retrieve_models( search: str = The search query for filtering models. """ - global api - if api is None: - api = HfApi() + global _HFAPI + if _HFAPI is None: + _HFAPI = HfApi() if full: properties = None - models: list[ModelInfo] = api.list_models( + models: list[ModelInfo] = _HFAPI.list_models( author=author, search=search, sort=sort,