start registry reog
This commit is contained in:
jeromeku
parent
ab77cc7c0d
commit
73ac5bf9fb
5 changed files with 493 additions and 219 deletions
177
.gitignore
vendored
Normal file
177
.gitignore
vendored
Normal file
|
|
@ -0,0 +1,177 @@
|
|||
# Byte-compiled / optimized / DLL files
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
*$py.class
|
||||
|
||||
# C extensions
|
||||
*.so
|
||||
|
||||
# Distribution / packaging
|
||||
.Python
|
||||
build/
|
||||
develop-eggs/
|
||||
dist/
|
||||
downloads/
|
||||
eggs/
|
||||
.eggs/
|
||||
lib/
|
||||
lib64/
|
||||
parts/
|
||||
sdist/
|
||||
var/
|
||||
wheels/
|
||||
share/python-wheels/
|
||||
*.egg-info/
|
||||
.installed.cfg
|
||||
*.egg
|
||||
MANIFEST
|
||||
|
||||
# PyInstaller
|
||||
# Usually these files are written by a python script from a template
|
||||
# before PyInstaller builds the exe, so as to inject date/other infos into it.
|
||||
*.manifest
|
||||
*.spec
|
||||
|
||||
# Installer logs
|
||||
pip-log.txt
|
||||
pip-delete-this-directory.txt
|
||||
|
||||
# Unit test / coverage reports
|
||||
htmlcov/
|
||||
.tox/
|
||||
.nox/
|
||||
.coverage
|
||||
.coverage.*
|
||||
.cache
|
||||
nosetests.xml
|
||||
coverage.xml
|
||||
*.cover
|
||||
*.py,cover
|
||||
.hypothesis/
|
||||
.pytest_cache/
|
||||
cover/
|
||||
|
||||
# Translations
|
||||
*.mo
|
||||
*.pot
|
||||
|
||||
# Django stuff:
|
||||
*.log
|
||||
local_settings.py
|
||||
db.sqlite3
|
||||
db.sqlite3-journal
|
||||
|
||||
# Flask stuff:
|
||||
instance/
|
||||
.webassets-cache
|
||||
|
||||
# Scrapy stuff:
|
||||
.scrapy
|
||||
|
||||
# Sphinx documentation
|
||||
docs/_build/
|
||||
|
||||
# PyBuilder
|
||||
.pybuilder/
|
||||
target/
|
||||
|
||||
# Jupyter Notebook
|
||||
.ipynb_checkpoints
|
||||
|
||||
# IPython
|
||||
profile_default/
|
||||
ipython_config.py
|
||||
|
||||
# pyenv
|
||||
# For a library or package, you might want to ignore these files since the code is
|
||||
# intended to run in multiple environments; otherwise, check them in:
|
||||
# .python-version
|
||||
|
||||
# pipenv
|
||||
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
|
||||
# However, in case of collaboration, if having platform-specific dependencies or dependencies
|
||||
# having no cross-platform support, pipenv may install dependencies that don't work, or not
|
||||
# install all needed dependencies.
|
||||
#Pipfile.lock
|
||||
|
||||
# UV
|
||||
# Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
|
||||
# This is especially recommended for binary packages to ensure reproducibility, and is more
|
||||
# commonly ignored for libraries.
|
||||
#uv.lock
|
||||
|
||||
# poetry
|
||||
# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
|
||||
# This is especially recommended for binary packages to ensure reproducibility, and is more
|
||||
# commonly ignored for libraries.
|
||||
# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
|
||||
#poetry.lock
|
||||
|
||||
# pdm
|
||||
# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
|
||||
#pdm.lock
|
||||
# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
|
||||
# in version control.
|
||||
# https://pdm.fming.dev/latest/usage/project/#working-with-version-control
|
||||
.pdm.toml
|
||||
.pdm-python
|
||||
.pdm-build/
|
||||
|
||||
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
|
||||
__pypackages__/
|
||||
|
||||
# Celery stuff
|
||||
celerybeat-schedule
|
||||
celerybeat.pid
|
||||
|
||||
# SageMath parsed files
|
||||
*.sage.py
|
||||
|
||||
# Environments
|
||||
.env
|
||||
.venv
|
||||
env/
|
||||
venv/
|
||||
ENV/
|
||||
env.bak/
|
||||
venv.bak/
|
||||
|
||||
# Spyder project settings
|
||||
.spyderproject
|
||||
.spyproject
|
||||
|
||||
# Rope project settings
|
||||
.ropeproject
|
||||
|
||||
# mkdocs documentation
|
||||
/site
|
||||
|
||||
# mypy
|
||||
.mypy_cache/
|
||||
.dmypy.json
|
||||
dmypy.json
|
||||
|
||||
# Pyre type checker
|
||||
.pyre/
|
||||
|
||||
# pytype static type analyzer
|
||||
.pytype/
|
||||
|
||||
# Cython debug symbols
|
||||
cython_debug/
|
||||
|
||||
# PyCharm
|
||||
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
|
||||
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
|
||||
# and can be added to the global gitignore or merged into this file. For a more nuclear
|
||||
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
|
||||
#.idea/
|
||||
|
||||
# Ruff stuff:
|
||||
.ruff_cache/
|
||||
|
||||
# PyPI configuration file
|
||||
.pypirc
|
||||
|
||||
# unsloth compiled cache
|
||||
unsloth_compiled_cache
|
||||
0
unsloth/registry/__init__.py
Normal file
0
unsloth/registry/__init__.py
Normal file
77
unsloth/registry/_llama.py
Normal file
77
unsloth/registry/_llama.py
Normal file
|
|
@ -0,0 +1,77 @@
|
|||
from unsloth.registry.registry import ModelInfo, ModelMeta, _register_models
|
||||
|
||||
_IS_LLAMA_REGISTERED = False
|
||||
|
||||
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
|
||||
|
||||
|
||||
# Llama 3.1
|
||||
LlamaMeta3_1 = ModelMeta(
|
||||
org="meta-llama",
|
||||
base_name="Llama",
|
||||
instruct_tags=[None, "Instruct"],
|
||||
model_version="3.1",
|
||||
model_sizes=[8],
|
||||
model_info_cls=LlamaModelInfo,
|
||||
is_multimodal=False,
|
||||
quant_types=[None, "bnb", "unsloth"],
|
||||
)
|
||||
|
||||
# Llama 3.2
|
||||
LlamaMeta3_2 = ModelMeta(
|
||||
org="meta-llama",
|
||||
base_name="Llama",
|
||||
instruct_tags=[None, "Instruct"],
|
||||
model_version="3.2",
|
||||
model_sizes=[1, 3],
|
||||
model_info_cls=LlamaModelInfo,
|
||||
is_multimodal=False,
|
||||
quant_types=[None, "bnb", "unsloth"],
|
||||
)
|
||||
|
||||
# Llama 3.2 Vision
|
||||
LlamaMeta3_2_Vision = ModelMeta(
|
||||
org="meta-llama",
|
||||
base_name="Llama",
|
||||
instruct_tags=[None, "Instruct"],
|
||||
model_version="3.2",
|
||||
model_sizes=[11, 90],
|
||||
model_info_cls=LlamaVisionModelInfo,
|
||||
is_multimodal=True,
|
||||
quant_types=[None, "bnb", "unsloth"],
|
||||
)
|
||||
|
||||
|
||||
def register_llama_models():
|
||||
global _IS_LLAMA_REGISTERED
|
||||
if _IS_LLAMA_REGISTERED:
|
||||
return
|
||||
_register_models(LlamaMeta3_1)
|
||||
_register_models(LlamaMeta3_2)
|
||||
_IS_LLAMA_REGISTERED = True
|
||||
|
||||
register_llama_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}")
|
||||
|
|
@ -1,11 +1,8 @@
|
|||
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"]
|
||||
from unsloth.registry._llama import LlamaMeta3_1, LlamaMeta3_2
|
||||
from unsloth.registry.common import ModelInfo, ModelMeta
|
||||
|
||||
_IS_LLAMA_REGISTERED = False
|
||||
_IS_LLAMA_VISION_REGISTERED = False
|
||||
|
|
@ -19,222 +16,97 @@ _IS_PHI_REGISTERED = False
|
|||
_IS_PHI_INSTRUCT_REGISTERED = False
|
||||
|
||||
|
||||
@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 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 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 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 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 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
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelMeta:
|
||||
org: str
|
||||
base_name: str
|
||||
model_version: str
|
||||
model_info_cls: type[ModelInfo]
|
||||
model_sizes: list[str] = field(default_factory=list)
|
||||
instruct_tags: list[str] = field(default_factory=list)
|
||||
quant_types: list[Literal[None, "bnb", "unsloth"]] = field(
|
||||
default_factory=list
|
||||
)
|
||||
is_multimodal: bool = False
|
||||
|
||||
|
||||
LlamaMeta3_1 = ModelMeta(
|
||||
org="meta-llama",
|
||||
base_name="Llama",
|
||||
instruct_tags=[None, "Instruct"],
|
||||
model_version="3.1",
|
||||
model_sizes=[8],
|
||||
model_info_cls=LlamaModelInfo,
|
||||
is_multimodal=False,
|
||||
quant_types=[None, "bnb", "unsloth"],
|
||||
)
|
||||
|
||||
LlamaMeta3_2 = ModelMeta(
|
||||
org="meta-llama",
|
||||
base_name="Llama",
|
||||
instruct_tags=[None, "Instruct"],
|
||||
model_version="3.2",
|
||||
model_sizes=[1, 3],
|
||||
model_info_cls=LlamaModelInfo,
|
||||
is_multimodal=False,
|
||||
quant_types=[None, "bnb", "unsloth"],
|
||||
)
|
||||
|
||||
|
||||
# # Llama text only models
|
||||
# _LLAMA_INFO = {
|
||||
# "org": "meta-llama",
|
||||
# "base_name": "Llama",
|
||||
# # 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": ["3.2", "3.1"],
|
||||
# "model_sizes": {"3.2": [1, 3], "3.1": [8]},
|
||||
# "model_versions": ["2.5"],
|
||||
# "model_sizes": {"2.5": [3, 7]},
|
||||
# "is_multimodal": False,
|
||||
# "model_info_cls": LlamaModelInfo,
|
||||
# "model_info_cls": QwenModelInfo,
|
||||
# }
|
||||
|
||||
_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,
|
||||
# }
|
||||
|
||||
_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,
|
||||
# }
|
||||
|
||||
_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_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,
|
||||
}
|
||||
# _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 = {}
|
||||
MODEL_REGISTRY: dict[str, ModelInfo] = {}
|
||||
|
||||
|
||||
def register_model(
|
||||
|
|
@ -243,9 +115,9 @@ def register_model(
|
|||
base_name: str,
|
||||
version: str,
|
||||
size: int,
|
||||
instruct_tag: str = None,
|
||||
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(
|
||||
|
|
@ -323,7 +195,6 @@ def _register_models(model_meta: ModelMeta):
|
|||
is_multimodal=is_multimodal,
|
||||
)
|
||||
|
||||
|
||||
def register_llama_models():
|
||||
global _IS_LLAMA_REGISTERED
|
||||
if _IS_LLAMA_REGISTERED:
|
||||
149
unsloth/registry/registry.py
Normal file
149
unsloth/registry/registry.py
Normal file
|
|
@ -0,0 +1,149 @@
|
|||
from dataclasses import dataclass, field
|
||||
from typing import Literal
|
||||
|
||||
BNB_QUANTIZED_TAG = "bnb-4bit"
|
||||
UNSLOTH_DYNAMIC_QUANT_TAG = "unsloth" + "-" + BNB_QUANTIZED_TAG
|
||||
QUANT_TYPE_MAP = {
|
||||
"bnb": BNB_QUANTIZED_TAG,
|
||||
"unsloth": UNSLOTH_DYNAMIC_QUANT_TAG,
|
||||
"GGUF": "GGUF",
|
||||
}
|
||||
QUANT_TYPES = list(QUANT_TYPE_MAP.keys())
|
||||
|
||||
|
||||
@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", "GGUF"] = None
|
||||
):
|
||||
if quant_type:
|
||||
if quant_type == "bnb":
|
||||
key = "-".join([key, QUANT_TYPE_MAP["bnb"]])
|
||||
elif quant_type == "unsloth":
|
||||
key = "-".join([key, QUANT_TYPE_MAP["unsloth"]])
|
||||
elif quant_type == "GGUF":
|
||||
key = "-".join([key, QUANT_TYPE_MAP["GGUF"]])
|
||||
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}"
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelMeta:
|
||||
org: str
|
||||
base_name: str
|
||||
model_version: str
|
||||
model_info_cls: type[ModelInfo]
|
||||
model_sizes: list[str] = field(default_factory=list)
|
||||
instruct_tags: list[str] = field(default_factory=list)
|
||||
quant_types: list[Literal[None, "bnb", "unsloth"]] = field(default_factory=list)
|
||||
is_multimodal: bool = False
|
||||
|
||||
|
||||
MODEL_REGISTRY: dict[str, ModelInfo] = {}
|
||||
|
||||
|
||||
def register_model(
|
||||
model_info_cls: ModelInfo,
|
||||
org: str,
|
||||
base_name: str,
|
||||
version: str,
|
||||
size: int,
|
||||
instruct_tag: str = None,
|
||||
quant_type: Literal["bnb", "unsloth"] = None,
|
||||
is_multimodal: bool = False,
|
||||
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 _check_model_info(model_id: str, properties: list[str] = ["lastModified"]):
|
||||
from huggingface_hub import HfApi
|
||||
from huggingface_hub import ModelInfo as HfModelInfo
|
||||
api = HfApi()
|
||||
|
||||
try:
|
||||
model_info: HfModelInfo = 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 _register_models(model_meta: ModelMeta):
|
||||
org = model_meta.org
|
||||
base_name = model_meta.base_name
|
||||
instruct_tags = model_meta.instruct_tags
|
||||
model_version = model_meta.model_version
|
||||
model_sizes = model_meta.model_sizes
|
||||
is_multimodal = model_meta.is_multimodal
|
||||
quant_types = model_meta.quant_types
|
||||
model_info_cls = model_meta.model_info_cls
|
||||
|
||||
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
|
||||
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=quant_type,
|
||||
is_multimodal=is_multimodal,
|
||||
)
|
||||
Loading…
Add table
Add a link
Reference in a new issue