add model registry
This commit is contained in:
jeromeku
parent
7afd6afe2c
commit
6f6a7a5e9b
4 changed files with 548 additions and 0 deletions
0
tests/__init__.py
Normal file
0
tests/__init__.py
Normal file
86
tests/test_model_registry.py
Normal file
86
tests/test_model_registry.py
Normal file
|
|
@ -0,0 +1,86 @@
|
|||
from dataclasses import dataclass
|
||||
|
||||
import pytest
|
||||
from huggingface_hub import ModelInfo as HfModelInfo
|
||||
from unsloth.model_registry import (
|
||||
ModelInfo,
|
||||
get_llama_models,
|
||||
get_llama_vision_models,
|
||||
get_phi_instruct_models,
|
||||
get_phi_models,
|
||||
get_qwen_models,
|
||||
get_qwen_vl_models,
|
||||
)
|
||||
|
||||
from .utils.hf_hub import get_model_info
|
||||
|
||||
MODEL_NAMES = [
|
||||
"llama",
|
||||
"llama_vision",
|
||||
"qwen",
|
||||
"qwen_vl",
|
||||
"phi",
|
||||
"phi_instruct",
|
||||
]
|
||||
REGISTERED_MODELS = [
|
||||
get_llama_models(),
|
||||
get_llama_vision_models(),
|
||||
get_qwen_models(),
|
||||
get_qwen_vl_models(),
|
||||
get_phi_models(),
|
||||
get_phi_instruct_models(),
|
||||
]
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelTestParam:
|
||||
name: str
|
||||
models: dict[str, ModelInfo]
|
||||
|
||||
|
||||
def _test_model_uploaded(model_ids: list[str]):
|
||||
missing_models = []
|
||||
for _id in model_ids:
|
||||
model_info: HfModelInfo = get_model_info(_id)
|
||||
if not model_info:
|
||||
missing_models.append(_id)
|
||||
|
||||
return missing_models
|
||||
|
||||
|
||||
TestParams = [
|
||||
ModelTestParam(name, models)
|
||||
for name, models in zip(MODEL_NAMES, REGISTERED_MODELS)
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_test_param", TestParams, ids=lambda param: param.name
|
||||
)
|
||||
def test_model_uploaded(model_test_param: ModelTestParam):
|
||||
missing_models = _test_model_uploaded(model_test_param.models)
|
||||
assert not missing_models, (
|
||||
f"{model_test_param.name} missing following models: {missing_models}"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
for method in [
|
||||
get_llama_models,
|
||||
get_llama_vision_models,
|
||||
get_qwen_models,
|
||||
get_qwen_vl_models,
|
||||
get_phi_models,
|
||||
get_phi_instruct_models,
|
||||
]:
|
||||
models = method()
|
||||
model_name = next(iter(models.values())).base_name
|
||||
print(f"{model_name}: {len(models)} registered")
|
||||
for model_info in models.values():
|
||||
print(f" {model_info.model_path}")
|
||||
missing_models = test_model_uploaded(list(models.keys()))
|
||||
|
||||
if missing_models:
|
||||
print("--------------------------------")
|
||||
print(f"Missing models: {missing_models}")
|
||||
print("--------------------------------")
|
||||
72
tests/utils/hf_hub.py
Normal file
72
tests/utils/hf_hub.py
Normal file
|
|
@ -0,0 +1,72 @@
|
|||
from huggingface_hub import HfApi, ModelInfo
|
||||
|
||||
api = HfApi()
|
||||
|
||||
POPULARITY_PROPERTIES = [
|
||||
"downloads",
|
||||
"downloadsAllTime",
|
||||
"trendingScore",
|
||||
"likes",
|
||||
]
|
||||
THOUSAND = 1000
|
||||
MILLION = 1000000
|
||||
BILLION = 1000000000
|
||||
|
||||
|
||||
def formatted_int(value: int) -> str:
|
||||
if value < THOUSAND:
|
||||
return str(value)
|
||||
elif value < MILLION:
|
||||
return f"{float(value) / 1000:,.1f}K"
|
||||
elif value < BILLION:
|
||||
return f"{float(value) // 1000000:,.1f}M"
|
||||
|
||||
|
||||
def get_model_info(
|
||||
model_id: str, properties: list[str] = ["safetensors", "lastModified"]
|
||||
) -> ModelInfo:
|
||||
"""
|
||||
Get the model info for a specific model.
|
||||
|
||||
properties: list[str] = See https://huggingface.co/docs/huggingface_hub/api-ref/hf_hub/hf_api/model_info
|
||||
Default properties: ["safetensors", "lastModified"], only retrieves minimal information.
|
||||
Set to None to retrieve the full model information.
|
||||
"""
|
||||
try:
|
||||
model_info: ModelInfo = api.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(
|
||||
properties: list[str] = None,
|
||||
full: bool = False,
|
||||
sort: str = "downloads",
|
||||
author: str = "unsloth",
|
||||
search: str = None,
|
||||
limit: int = 10,
|
||||
) -> ModelInfo:
|
||||
"""
|
||||
Retrieve models 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.
|
||||
sort: str = The sort order.
|
||||
author: str = The author of the model.
|
||||
search: str = The search query for filtering models.
|
||||
|
||||
"""
|
||||
if full:
|
||||
properties = None
|
||||
|
||||
models: list[ModelInfo] = api.list_models(
|
||||
author=author,
|
||||
search=search,
|
||||
sort=sort,
|
||||
limit=limit,
|
||||
expand=properties,
|
||||
full=full,
|
||||
)
|
||||
return models
|
||||
390
unsloth/model_registry.py
Normal file
390
unsloth/model_registry.py
Normal file
|
|
@ -0,0 +1,390 @@
|
|||
from dataclasses import dataclass, field
|
||||
from functools import partial
|
||||
from typing import Callable, Literal
|
||||
|
||||
BNB_QUANTIZED_TAG = "bnb-4bit"
|
||||
UNSLOTH_DYNAMIC_QUANT_TAG = "unsloth" + "-" + BNB_QUANTIZED_TAG
|
||||
INSTRUCT_TAG = "Instruct"
|
||||
QUANT_TYPES = [None, "bnb", "unsloth"]
|
||||
|
||||
_IS_LLAMA_REGISTERED = False
|
||||
_IS_LLAMA_VISION_REGISTERED = False
|
||||
|
||||
_IS_QWEN_REGISTERED = False
|
||||
_IS_QWEN_VL_REGISTERED = False
|
||||
|
||||
_IS_GEMMA_REGISTERED = False
|
||||
|
||||
_IS_PHI_REGISTERED = False
|
||||
_IS_PHI_INSTRUCT_REGISTERED = False
|
||||
|
||||
|
||||
def construct_model_key(org, base_name, version, size, quant_type, instruct_tag):
|
||||
key = f"{org}/{base_name}-{version}-{size}B"
|
||||
if instruct_tag:
|
||||
key = "-".join([key, instruct_tag])
|
||||
if quant_type:
|
||||
if quant_type == "bnb":
|
||||
key = "-".join([key, BNB_QUANTIZED_TAG])
|
||||
elif quant_type == "unsloth":
|
||||
key = "-".join([key, UNSLOTH_DYNAMIC_QUANT_TAG])
|
||||
return key
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelInfo:
|
||||
org: str
|
||||
base_name: str
|
||||
version: str
|
||||
size: int
|
||||
name: str = None # full model name, constructed from base_name, version, and size unless provided
|
||||
is_multimodal: bool = False
|
||||
instruct_tag: str = None
|
||||
quant_type: Literal["bnb", "unsloth"] = None
|
||||
|
||||
def __post_init__(self):
|
||||
self.name = self.name or self.construct_model_name(
|
||||
self.base_name,
|
||||
self.version,
|
||||
self.size,
|
||||
self.quant_type,
|
||||
self.instruct_tag,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def append_instruct_tag(key: str, instruct_tag: str = None):
|
||||
if instruct_tag:
|
||||
key = "-".join([key, instruct_tag])
|
||||
return key
|
||||
|
||||
@staticmethod
|
||||
def append_quant_type(key: str, quant_type: Literal["bnb", "unsloth"] = None):
|
||||
if quant_type:
|
||||
if quant_type == "bnb":
|
||||
key = "-".join([key, BNB_QUANTIZED_TAG])
|
||||
elif quant_type == "unsloth":
|
||||
key = "-".join([key, UNSLOTH_DYNAMIC_QUANT_TAG])
|
||||
return key
|
||||
|
||||
@classmethod
|
||||
def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag):
|
||||
raise NotImplementedError("Subclass must implement this method")
|
||||
|
||||
@property
|
||||
def model_path(
|
||||
self,
|
||||
) -> str:
|
||||
return f"{self.org}/{self.name}"
|
||||
|
||||
|
||||
class LlamaModelInfo(ModelInfo):
|
||||
@classmethod
|
||||
def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag):
|
||||
key = f"{base_name}-{version}-{size}B"
|
||||
key = cls.append_instruct_tag(key, instruct_tag)
|
||||
key = cls.append_quant_type(key, quant_type)
|
||||
return key
|
||||
|
||||
|
||||
class LlamaVisionModelInfo(ModelInfo):
|
||||
@classmethod
|
||||
def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag):
|
||||
key = f"{base_name}-{version}-{size}B-Vision"
|
||||
key = cls.append_instruct_tag(key, instruct_tag)
|
||||
key = cls.append_quant_type(key, quant_type)
|
||||
return key
|
||||
|
||||
|
||||
class QwenModelInfo(ModelInfo):
|
||||
@classmethod
|
||||
def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag):
|
||||
key = f"{base_name}{version}-{size}B"
|
||||
key = cls.append_instruct_tag(key, instruct_tag)
|
||||
key = cls.append_quant_type(key, quant_type)
|
||||
return key
|
||||
|
||||
|
||||
class QwenVLModelInfo(ModelInfo):
|
||||
@classmethod
|
||||
def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag):
|
||||
key = f"{base_name}{version}-VL-{size}B"
|
||||
key = cls.append_instruct_tag(key, instruct_tag)
|
||||
key = cls.append_quant_type(key, quant_type)
|
||||
return key
|
||||
|
||||
|
||||
class PhiModelInfo(ModelInfo):
|
||||
@classmethod
|
||||
def construct_model_name(cls, base_name, version, size, quant_type, instruct_tag):
|
||||
key = f"{base_name}-{version}"
|
||||
key = cls.append_instruct_tag(key, instruct_tag)
|
||||
key = cls.append_quant_type(key, quant_type)
|
||||
return key
|
||||
|
||||
|
||||
# Llama text only models
|
||||
_LLAMA_INFO = {
|
||||
"org": "meta-llama",
|
||||
"base_name": "Llama",
|
||||
"instruct_tags": [None, "Instruct"],
|
||||
"model_versions": ["3.2", "3.1"],
|
||||
"model_sizes": {"3.2": [1, 3], "3.1": [8]},
|
||||
"is_multimodal": False,
|
||||
"model_info_cls": LlamaModelInfo,
|
||||
}
|
||||
|
||||
_LLAMA_VISION_INFO = {
|
||||
"org": "meta-llama",
|
||||
"base_name": "Llama",
|
||||
"instruct_tags": [None, "Instruct"],
|
||||
"model_versions": ["3.2"],
|
||||
"model_sizes": {"3.2": [11, 90]},
|
||||
"is_multimodal": True,
|
||||
"model_info_cls": LlamaVisionModelInfo,
|
||||
}
|
||||
# Qwen text only models
|
||||
# NOTE: Qwen vision models will be registered separately
|
||||
_QWEN_INFO = {
|
||||
"org": "Qwen",
|
||||
"base_name": "Qwen",
|
||||
"instruct_tags": [None, "Instruct"],
|
||||
"model_versions": ["2.5"],
|
||||
"model_sizes": {"2.5": [3, 7]},
|
||||
"is_multimodal": False,
|
||||
"model_info_cls": QwenModelInfo,
|
||||
}
|
||||
|
||||
_QWEN_VL_INFO = {
|
||||
"org": "Qwen",
|
||||
"base_name": "Qwen",
|
||||
"instruct_tags": ["Instruct"], # No base, only instruction tuned
|
||||
"model_versions": ["2.5"],
|
||||
"model_sizes": {"2.5": [3, 7, 32, 72]},
|
||||
"is_multimodal": True,
|
||||
"instruction_tuned_only": True,
|
||||
"model_info_cls": QwenVLModelInfo,
|
||||
}
|
||||
|
||||
_GEMMA_INFO = {
|
||||
"org": "google",
|
||||
"base_name": "gemma",
|
||||
"instruct_tags": ["pt", "it"], # pt = base, it = instruction tuned
|
||||
"model_versions": ["3"],
|
||||
"model_sizes": {"3": [1, 4, 12, 27]},
|
||||
"is_multimodal": True,
|
||||
}
|
||||
|
||||
_PHI_INFO = {
|
||||
"org": "microsoft",
|
||||
"base_name": "phi",
|
||||
"model_versions": ["4"],
|
||||
"model_sizes": {"4": [None]}, # -1 means only 1 size
|
||||
"instruct_tags": [None],
|
||||
"is_multimodal": False,
|
||||
"model_info_cls": PhiModelInfo,
|
||||
}
|
||||
|
||||
_PHI_INSTRUCT_INFO = {
|
||||
"org": "microsoft",
|
||||
"base_name": "Phi",
|
||||
"model_versions": ["4"],
|
||||
"model_sizes": {"4": [None]}, # -1 means only 1 size
|
||||
"instruct_tags": ["mini-instruct"],
|
||||
"is_multimodal": False,
|
||||
"model_info_cls": PhiModelInfo,
|
||||
}
|
||||
|
||||
|
||||
MODEL_REGISTRY = {}
|
||||
|
||||
|
||||
def register_model(
|
||||
model_info_cls: ModelInfo,
|
||||
org: str,
|
||||
base_name: str,
|
||||
version: str,
|
||||
size: int,
|
||||
quant_type: Literal["bnb", "unsloth"] = None,
|
||||
is_multimodal: bool = False,
|
||||
instruct_tag: str = INSTRUCT_TAG,
|
||||
name: str = None,
|
||||
):
|
||||
name = name or model_info_cls.construct_model_name(
|
||||
base_name=base_name,
|
||||
version=version,
|
||||
size=size,
|
||||
quant_type=quant_type,
|
||||
instruct_tag=instruct_tag,
|
||||
)
|
||||
key = f"{org}/{name}"
|
||||
|
||||
if key in MODEL_REGISTRY:
|
||||
raise ValueError(f"Model {key} already registered")
|
||||
|
||||
MODEL_REGISTRY[key] = model_info_cls(
|
||||
org=org,
|
||||
base_name=base_name,
|
||||
version=version,
|
||||
size=size,
|
||||
is_multimodal=is_multimodal,
|
||||
instruct_tag=instruct_tag,
|
||||
quant_type=quant_type,
|
||||
name=name,
|
||||
)
|
||||
|
||||
|
||||
def _register_models(model_info: dict):
|
||||
org = model_info["org"]
|
||||
base_name = model_info["base_name"]
|
||||
instruct_tags = model_info["instruct_tags"]
|
||||
model_versions = model_info["model_versions"]
|
||||
model_sizes = model_info["model_sizes"]
|
||||
is_multimodal = model_info["is_multimodal"]
|
||||
model_info_cls = model_info["model_info_cls"]
|
||||
|
||||
for version in model_versions:
|
||||
for size in model_sizes[version]:
|
||||
for instruct_tag in instruct_tags:
|
||||
for quant_type in QUANT_TYPES:
|
||||
_org = "unsloth" if quant_type is not None else org
|
||||
register_model(
|
||||
model_info_cls=model_info_cls,
|
||||
org=_org,
|
||||
base_name=base_name,
|
||||
version=version,
|
||||
size=size,
|
||||
instruct_tag=instruct_tag,
|
||||
quant_type=quant_type,
|
||||
is_multimodal=is_multimodal,
|
||||
)
|
||||
|
||||
|
||||
def register_llama_models():
|
||||
global _IS_LLAMA_REGISTERED
|
||||
if _IS_LLAMA_REGISTERED:
|
||||
return
|
||||
_register_models(_LLAMA_INFO)
|
||||
_IS_LLAMA_REGISTERED = True
|
||||
|
||||
|
||||
def register_llama_vision_models():
|
||||
global _IS_LLAMA_VISION_REGISTERED
|
||||
if _IS_LLAMA_VISION_REGISTERED:
|
||||
return
|
||||
_register_models(_LLAMA_VISION_INFO)
|
||||
_IS_LLAMA_VISION_REGISTERED = True
|
||||
|
||||
|
||||
def register_qwen_models():
|
||||
global _IS_QWEN_REGISTERED
|
||||
if _IS_QWEN_REGISTERED:
|
||||
return
|
||||
|
||||
_register_models(_QWEN_INFO)
|
||||
_IS_QWEN_REGISTERED = True
|
||||
|
||||
|
||||
def register_qwen_vl_models():
|
||||
global _IS_QWEN_VL_REGISTERED
|
||||
if _IS_QWEN_VL_REGISTERED:
|
||||
return
|
||||
|
||||
_register_models(_QWEN_VL_INFO)
|
||||
_IS_QWEN_VL_REGISTERED = True
|
||||
|
||||
|
||||
def register_gemma_models():
|
||||
global _IS_GEMMA_REGISTERED
|
||||
_register_models(_GEMMA_INFO)
|
||||
_IS_GEMMA_REGISTERED = True
|
||||
|
||||
|
||||
def register_phi_models():
|
||||
global _IS_PHI_REGISTERED
|
||||
if _IS_PHI_REGISTERED:
|
||||
return
|
||||
_register_models(_PHI_INFO)
|
||||
_IS_PHI_REGISTERED = True
|
||||
|
||||
|
||||
def register_phi_instruct_models():
|
||||
global _IS_PHI_INSTRUCT_REGISTERED
|
||||
if _IS_PHI_INSTRUCT_REGISTERED:
|
||||
return
|
||||
|
||||
_register_models(_PHI_INSTRUCT_INFO)
|
||||
_IS_PHI_INSTRUCT_REGISTERED = True
|
||||
|
||||
|
||||
def _base_name_filter(model_info: ModelInfo, base_name: str):
|
||||
return model_info.base_name == base_name
|
||||
|
||||
|
||||
def _get_models(filter_func: Callable[[ModelInfo], bool] = _base_name_filter):
|
||||
return {k: v for k, v in MODEL_REGISTRY.items() if filter_func(v)}
|
||||
|
||||
|
||||
def get_llama_models():
|
||||
if not _IS_LLAMA_REGISTERED:
|
||||
register_llama_models()
|
||||
|
||||
return _get_models(partial(_base_name_filter, base_name=_LLAMA_INFO["base_name"]))
|
||||
|
||||
|
||||
def get_llama_vision_models():
|
||||
if not _IS_LLAMA_VISION_REGISTERED:
|
||||
register_llama_vision_models()
|
||||
|
||||
return _get_models(
|
||||
lambda model_info: model_info.base_name == _LLAMA_VISION_INFO["base_name"]
|
||||
and model_info.is_multimodal
|
||||
)
|
||||
|
||||
|
||||
def get_qwen_models():
|
||||
if not _IS_QWEN_REGISTERED:
|
||||
register_qwen_models()
|
||||
|
||||
return _get_models(
|
||||
lambda model_info: model_info.base_name == _QWEN_INFO["base_name"]
|
||||
)
|
||||
|
||||
|
||||
def get_qwen_vl_models():
|
||||
if not _IS_QWEN_VL_REGISTERED:
|
||||
register_qwen_vl_models()
|
||||
return _get_models(
|
||||
lambda model_info: model_info.base_name == _QWEN_VL_INFO["base_name"]
|
||||
)
|
||||
|
||||
|
||||
def get_gemma_models():
|
||||
if not _IS_GEMMA_REGISTERED:
|
||||
register_gemma_models()
|
||||
|
||||
return _get_models(
|
||||
lambda model_info: model_info.base_name == _GEMMA_INFO["base_name"]
|
||||
)
|
||||
|
||||
|
||||
def get_phi_models():
|
||||
if not _IS_PHI_REGISTERED:
|
||||
register_phi_models()
|
||||
return _get_models(
|
||||
lambda model_info: model_info.base_name == _PHI_INFO["base_name"]
|
||||
)
|
||||
|
||||
|
||||
def get_phi_instruct_models():
|
||||
if not _IS_PHI_INSTRUCT_REGISTERED:
|
||||
register_phi_instruct_models()
|
||||
return _get_models(
|
||||
lambda model_info: model_info.base_name == _PHI_INSTRUCT_INFO["base_name"]
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
register_llama_models()
|
||||
for k, v in MODEL_REGISTRY.items():
|
||||
print(f"{k}: {v}")
|
||||
print(v.model_path)
|
||||
Loading…
Add table
Add a link
Reference in a new issue