From 83eef36689cf8cf1a6679907b409da431f2a68e9 Mon Sep 17 00:00:00 2001
From: jeromeku <jerome.ku@gmail.com>
Date: Mon, 31 Mar 2025 09:09:34 -0700
Subject: [PATCH] add option to include original model in registry

---
 unsloth/registry/_qwen.py    | 44 +++++++++++++++++++++++++++++-------
 unsloth/registry/registry.py | 16 +++++++++++--
 2 files changed, 50 insertions(+), 10 deletions(-)

diff --git a/unsloth/registry/_qwen.py b/unsloth/registry/_qwen.py
index 92f366bb76..2ea340b813 100644
--- a/unsloth/registry/_qwen.py
+++ b/unsloth/registry/_qwen.py
@@ -2,7 +2,7 @@ from unsloth.registry.registry import ModelInfo, ModelMeta, QuantType, _register
 
 _IS_QWEN_REGISTERED = False
 _IS_QWEN_VL_REGISTERED = False
-
+_IS_QWEN_QWQ_REGISTERED = False
 class QwenModelInfo(ModelInfo):
     @classmethod
     def construct_model_name(
@@ -24,7 +24,16 @@ class QwenVLModelInfo(ModelInfo):
         key = cls.append_quant_type(key, quant_type)
         return key
 
-
+class QwenQwQModelInfo(ModelInfo):
+    @classmethod
+    def construct_model_name(
+        cls, base_name, version, size, quant_type, instruct_tag
+    ):
+        key = f"{base_name}-{size}B"
+        key = cls.append_instruct_tag(key, instruct_tag)
+        key = cls.append_quant_type(key, quant_type)
+        return key
+    
 # Qwen Model Meta
 QwenMeta = ModelMeta(
     org="Qwen",
@@ -49,23 +58,42 @@ QwenVLMeta = ModelMeta(
     quant_types=[QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH],
 )
 
-def register_qwen_models():
+# Qwen QwQ Model Meta
+QwenQwQMeta = ModelMeta(
+    org="Qwen",
+    base_name="QwQ",
+    instruct_tags=[None],
+    model_version="",
+    model_sizes=[32],
+    model_info_cls=QwenQwQModelInfo,
+    is_multimodal=False,
+    quant_types=[QuantType.NONE, QuantType.BNB, QuantType.UNSLOTH, QuantType.GGUF],
+)
+
+def register_qwen_models(include_original_model: bool = False):
     global _IS_QWEN_REGISTERED
     if _IS_QWEN_REGISTERED:
         return
-    _register_models(QwenMeta)
+    _register_models(QwenMeta, include_original_model)
     _IS_QWEN_REGISTERED = True
 
-def register_qwen_vl_models():
+def register_qwen_vl_models(include_original_model: bool = False):
     global _IS_QWEN_VL_REGISTERED
     if _IS_QWEN_VL_REGISTERED:
         return
-    _register_models(QwenVLMeta)
+    _register_models(QwenVLMeta, include_original_model)
     _IS_QWEN_VL_REGISTERED = True
 
-register_qwen_models()
-register_qwen_vl_models()
+def register_qwen_qwq_models(include_original_model: bool = False):
+    global _IS_QWEN_QWQ_REGISTERED
+    if _IS_QWEN_QWQ_REGISTERED:
+        return
+    _register_models(QwenQwQMeta, include_original_model)
+    _IS_QWEN_QWQ_REGISTERED = True
 
+# register_qwen_models()
+# register_qwen_vl_models()
+register_qwen_qwq_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 d045f5bd55..3ca7c20f8f 100644
--- a/unsloth/registry/registry.py
+++ b/unsloth/registry/registry.py
@@ -134,7 +134,7 @@ def _check_model_info(model_id: str, properties: list[str] = ["lastModified"]):
     return model_info
 
 
-def _register_models(model_meta: ModelMeta):
+def _register_models(model_meta: ModelMeta, include_original_model: bool = False):
     org = model_meta.org
     base_name = model_meta.base_name
     instruct_tags = model_meta.instruct_tags
@@ -147,7 +147,7 @@ def _register_models(model_meta: ModelMeta):
     for size in model_sizes:
         for instruct_tag in instruct_tags:
             for quant_type in quant_types:
-                _org = "unsloth" if quant_type is not None else org
+                _org = "unsloth" # unsloth models -- these are all quantized versions of the original model
                 register_model(
                     model_info_cls=model_info_cls,
                     org=_org,
@@ -158,3 +158,15 @@ def _register_models(model_meta: ModelMeta):
                     quant_type=quant_type,
                     is_multimodal=is_multimodal,
                 )
+            # include original model from releasing organization
+            if include_original_model:
+                register_model(
+                    model_info_cls=model_info_cls,
+                    org=org,
+                    base_name=base_name,
+                    version=model_version,
+                    size=size,
+                    instruct_tag=instruct_tag,
+                    quant_type=QuantType.NONE,
+                    is_multimodal=is_multimodal,
+                )